From 07fed14f21b43c91c8dfe837044333e411b10077 Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Wed, 30 Sep 2026 00:31:20 -0700 Subject: [PATCH 1/3] perf(cake_sparse_mla): regenerate sm_103a programs from Cake 56f72c78456 (round-7 export) Co-Authored-By: Claude Fable 5.1 --- .../cake_dsv4_064850af5858dd03c104_binding.cu | 606 ++++ .../cake_dsv4_064850af5858dd03c104_kernel.cu | 2363 +++++++++++++ .../cake_dsv4_0a925b9bcb3044777c76_binding.cu | 441 +++ .../cake_dsv4_0a925b9bcb3044777c76_kernel.cu | 2967 +++++++++++++++++ .../cake_dsv4_1bad879f597b2e2adf3e_binding.cu | 389 +++ .../cake_dsv4_1bad879f597b2e2adf3e_kernel.cu | 1742 ++++++++++ .../cake_dsv4_1e3cf4361de31c58b69b_binding.cu | 436 +++ .../cake_dsv4_1e3cf4361de31c58b69b_kernel.cu | 2556 ++++++++++++++ .../cake_dsv4_1ff52fece183b851b441_binding.cu | 534 +++ .../cake_dsv4_1ff52fece183b851b441_kernel.cu | 1686 ++++++++++ .../cake_dsv4_4ab0487d5fde70d4569e_binding.cu | 606 ++++ .../cake_dsv4_4ab0487d5fde70d4569e_kernel.cu | 2418 ++++++++++++++ .../cake_dsv4_868bb2291a1464f14e05_binding.cu | 525 +++ .../cake_dsv4_868bb2291a1464f14e05_kernel.cu | 2920 ++++++++++++++++ .../cake_dsv4_983db4c490ceed3c455a_binding.cu | 461 +++ .../cake_dsv4_983db4c490ceed3c455a_kernel.cu | 2802 ++++++++++++++++ .../cake_dsv4_98db1f55c8c129e16005_binding.cu | 461 +++ .../cake_dsv4_98db1f55c8c129e16005_kernel.cu | 2835 ++++++++++++++++ .../cake_dsv4_c1d444bba4cdb8273c7c_binding.cu | 526 +++ .../cake_dsv4_c1d444bba4cdb8273c7c_kernel.cu | 2655 +++++++++++++++ .../cake_dsv4_d9754db4fe394a927e67_binding.cu | 385 +++ .../cake_dsv4_d9754db4fe394a927e67_kernel.cu | 1734 ++++++++++ .../cake_dsv4_e48c5e0a89a41eb1e57a_binding.cu | 526 +++ .../cake_dsv4_e48c5e0a89a41eb1e57a_kernel.cu | 2734 +++++++++++++++ ...sv4_family_0d09e8b242706c4e05a9_binding.cu | 703 ++++ ...sv4_family_ebb85adb872529ca9e83_binding.cu | 703 ++++ flashinfer/jit/cake_dsv4.py | 181 +- 27 files changed, 36804 insertions(+), 91 deletions(-) create mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_064850af5858dd03c104_binding.cu create mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_064850af5858dd03c104_kernel.cu create mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_0a925b9bcb3044777c76_binding.cu create mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_0a925b9bcb3044777c76_kernel.cu create mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_1bad879f597b2e2adf3e_binding.cu create mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_1bad879f597b2e2adf3e_kernel.cu create mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_1e3cf4361de31c58b69b_binding.cu create mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_1e3cf4361de31c58b69b_kernel.cu create mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_1ff52fece183b851b441_binding.cu create mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_1ff52fece183b851b441_kernel.cu create mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_4ab0487d5fde70d4569e_binding.cu create mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_4ab0487d5fde70d4569e_kernel.cu create mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_868bb2291a1464f14e05_binding.cu create mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_868bb2291a1464f14e05_kernel.cu create mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_983db4c490ceed3c455a_binding.cu create mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_983db4c490ceed3c455a_kernel.cu create mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_98db1f55c8c129e16005_binding.cu create mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_98db1f55c8c129e16005_kernel.cu create mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_c1d444bba4cdb8273c7c_binding.cu create mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_c1d444bba4cdb8273c7c_kernel.cu create mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_d9754db4fe394a927e67_binding.cu create mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_d9754db4fe394a927e67_kernel.cu create mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_e48c5e0a89a41eb1e57a_binding.cu create mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_e48c5e0a89a41eb1e57a_kernel.cu create mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_family_0d09e8b242706c4e05a9_binding.cu create mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_family_ebb85adb872529ca9e83_binding.cu diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_064850af5858dd03c104_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_064850af5858dd03c104_binding.cu new file mode 100644 index 00000000000..0392d451107 --- /dev/null +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_064850af5858dd03c104_binding.cu @@ -0,0 +1,606 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +// Generated by the Cake source exporter — do not edit. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_064850af5858dd03c104'. +#include +#include +#include +#include + +#include "tvm_ffi_utils.h" + +#include +#include +#include +#include +#include +#include +#include + +extern "C" __global__ void kernel_cake_dsv4_064850af5858dd03c104(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, const __grid_constant__ CUtensorMap O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); + + +namespace cake_host_shim_66599af918c96e07 { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t CakeDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +// 4D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + int64_t d3 = t.size(t.ndim() - 3); + TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(d3)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 64u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_swa_k' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_swa_k(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_swa_k' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_swa_k' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_swa_k' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_swa_k"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_swa_k' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_swa_k' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_k'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_swa_k') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_compressed_k' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_compressed_k(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_compressed_k' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_compressed_k' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_compressed_k' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_compressed_k"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_compressed_k' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_compressed_k' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_k'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_k') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_swa_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_swa_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_swa_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_swa_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_swa_v' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_swa_v"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_swa_v' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_swa_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_v'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_swa_v') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_compressed_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_compressed_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_compressed_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_compressed_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_compressed_v' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_compressed_v"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_compressed_v' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_compressed_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_v'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_v') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'O' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_O(const TensorView& t, int64_t tp_num_heads, int64_t tp_num_query_tokens) { + TVM_FFI_CHECK(t.ndim() >= 1, ValueError) + << "TMA source 'O' must have at least 1 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'O' must have unit innermost stride, got " << t.stride(-1); + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(8), (uint64_t)(4), (uint64_t)((tp_num_query_tokens * tp_num_heads))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'O' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1] && 1u <= global_dim[2] && 32u <= global_dim[3], ValueError) + << "TMA box (64, 1, 1, 32) exceeds resolved global dims for 'O'"; + int64_t carrier_stride_0 = 64; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 512; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = 2048; + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 1u, 1u, 32u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'O') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_compressed_k, TensorView arg_tmap_swa_v, TensorView arg_tmap_compressed_v, TensorView arg_O, TensorView arg_partial_lse, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_num_query_tokens, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckContiguous(arg_tmap_q, "tmap_q"); + CheckCudaTensor(arg_tmap_swa_k, "tmap_swa_k"); + CheckDtype(arg_tmap_swa_k, "tmap_swa_k", 4, 16, 1); + CheckCudaTensor(arg_tmap_compressed_k, "tmap_compressed_k"); + CheckDtype(arg_tmap_compressed_k, "tmap_compressed_k", 4, 16, 1); + CheckCudaTensor(arg_tmap_swa_v, "tmap_swa_v"); + CheckDtype(arg_tmap_swa_v, "tmap_swa_v", 4, 16, 1); + CheckCudaTensor(arg_tmap_compressed_v, "tmap_compressed_v"); + CheckDtype(arg_tmap_compressed_v, "tmap_compressed_v", 4, 16, 1); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_partial_lse, "partial_lse"); + CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); + CheckContiguous(arg_partial_lse, "partial_lse"); + CheckCudaTensor(arg_swa_indices, "swa_indices"); + CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); + CheckContiguous(arg_swa_indices, "swa_indices"); + CheckCudaTensor(arg_compressed_indices, "compressed_indices"); + CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); + CheckContiguous(arg_compressed_indices, "compressed_indices"); + CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); + CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); + CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); + CheckContiguous(arg_bmm1_scale, "bmm1_scale"); + CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); + CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); + CheckContiguous(arg_bmm2_scale, "bmm2_scale"); + TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) + << "scalar 'num_heads' value " << arg_num_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_query_tokens >= -2147483648LL && arg_num_query_tokens <= 2147483647LL, ValueError) + << "scalar 'num_query_tokens' value " << arg_num_query_tokens + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) + << "scalar 'swa_index_stride' value " << arg_swa_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) + << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) + << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) + << "scalar 'sparse_topk' value " << arg_sparse_topk + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) + << "scalar 'has_sinks' value " << arg_has_sinks + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_swa_k, arg_tmap_q, "tmap_swa_k", "tmap_q"); + CheckSameCudaDevice(arg_tmap_compressed_k, arg_tmap_q, "tmap_compressed_k", "tmap_q"); + CheckSameCudaDevice(arg_tmap_swa_v, arg_tmap_q, "tmap_swa_v", "tmap_q"); + CheckSameCudaDevice(arg_tmap_compressed_v, arg_tmap_q, "tmap_compressed_v", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_partial_lse, arg_tmap_q, "partial_lse", "tmap_q"); + CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); + CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); + CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); + CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") must be divisible by cluster dims (2, 1, 1)"; + + cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_swa_k = EncodeTma_tmap_swa_k(arg_tmap_swa_k); + CUtensorMap p_tmap_compressed_k = EncodeTma_tmap_compressed_k(arg_tmap_compressed_k); + CUtensorMap p_tmap_swa_v = EncodeTma_tmap_swa_v(arg_tmap_swa_v); + CUtensorMap p_tmap_compressed_v = EncodeTma_tmap_compressed_v(arg_tmap_compressed_v); + CUtensorMap p_O = EncodeTma_O(arg_O, (int64_t)arg_num_heads, (int64_t)arg_num_query_tokens); + float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); + int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); + int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); + int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); + float* p_sinks = static_cast(arg_sinks.data_ptr()); + float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); + float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); + int32_t v_num_heads = (int32_t)arg_num_heads; + int32_t v_num_query_tokens = (int32_t)arg_num_query_tokens; + int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; + int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; + int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; + int32_t v_sparse_topk = (int32_t)arg_sparse_topk; + int32_t v_has_sinks = (int32_t)arg_has_sinks; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + void* kargs[] = {&p_tmap_q, &p_tmap_swa_k, &p_tmap_compressed_k, &p_tmap_swa_v, &p_tmap_compressed_v, &p_O, &p_partial_lse, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_num_query_tokens, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_has_sinks, &v_total_work_items}; + + static std::mutex smem_mu; + static auto* smem_status_by_device = + new std::unordered_map(); + cudaError_t smem_status = cudaSuccess; + { + std::lock_guard lock(smem_mu); + auto it = smem_status_by_device->find(dev.device_id); + if (it == smem_status_by_device->end()) { + int current_device = -1; + smem_status = cudaGetDevice(¤t_device); + if (smem_status == cudaSuccess && current_device != dev.device_id) { + smem_status = cudaErrorInvalidDevice; + } + if (smem_status == cudaSuccess) { + smem_status = cudaFuncSetAttribute( + reinterpret_cast(kernel_cake_dsv4_064850af5858dd03c104), + cudaFuncAttributeMaxDynamicSharedMemorySize, + 220032); + } + smem_status_by_device->emplace(dev.device_id, smem_status); + } else { + smem_status = it->second; + } + } + TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) + << "cudaFuncSetAttribute for kernel_cake_dsv4_064850af5858dd03c104 failed: " + << cudaGetErrorString(smem_status); + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(512u, 1u, 1u); + + // Extended direct-source launch for a linked CUDA kernel symbol. + cudaLaunchAttribute attrs[2]{}; + int n = 0; + attrs[n].id = cudaLaunchAttributeClusterDimension; + attrs[n].val.clusterDim.x = 2u; + attrs[n].val.clusterDim.y = 1u; + attrs[n].val.clusterDim.z = 1u; + ++n; + attrs[n].id = cudaLaunchAttributeClusterSchedulingPolicyPreference; + attrs[n].val.clusterSchedulingPolicyPreference = cudaClusterSchedulingPolicySpread; + ++n; + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes = 220032u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; + cudaError_t launch_status = cudaLaunchKernelExC( + &config, reinterpret_cast(kernel_cake_dsv4_064850af5858dd03c104), kargs); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernelExC for kernel_cake_dsv4_064850af5858dd03c104 failed: " + << cudaGetErrorString(launch_status); +} + +} // namespace cake_host_shim_66599af918c96e07 + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_66599af918c96e07::Run); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_064850af5858dd03c104_kernel.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_064850af5858dd03c104_kernel.cu new file mode 100644 index 00000000000..1016950b3e6 --- /dev/null +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_064850af5858dd03c104_kernel.cu @@ -0,0 +1,2363 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define CAKE_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_SCRATCH_OFFSET 0 +#define NUM_K_PIPE_STAGES 3 +#define NUM_V_PIPE_STAGES 2 +#define SMEM_SMEM_Q_OFF 1024 +#define SMEM_SMEM_Q_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_STRIDE 8192 +#define SMEM_SMEM_K_OFF 66560 +#define SMEM_SMEM_K_STAGE_BYTES 16384 +#define SMEM_SMEM_K_STRIDE 16384 +#define SMEM_SMEM_V_LO_OFF 115712 +#define SMEM_SMEM_V_LO_STAGE_BYTES 16384 +#define SMEM_SMEM_V_LO_STRIDE 32768 +#define SMEM_SMEM_V_HI_OFF 132096 +#define SMEM_SMEM_V_HI_STAGE_BYTES 16384 +#define SMEM_SMEM_V_HI_STRIDE 32768 +#define SMEM_SMEM_P_OFF 181248 +#define SMEM_SMEM_P_STAGE_BYTES 8192 +#define SMEM_SMEM_P_STRIDE 8192 +#define SMEM_SMEM_INDICES_OFF 214016 +#define SMEM_SMEM_INDICES_STAGE_BYTES 2048 +#define SMEM_SMEM_INDICES_STRIDE 2048 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_OFF 218624 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STRIDE 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_OFF 219136 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STRIDE 512 +#define SMEM_SMEM_INDEX_VALID_OFF 219648 +#define SMEM_SMEM_INDEX_VALID_STAGE_BYTES 288 +#define SMEM_SMEM_INDEX_VALID_STRIDE 288 +#define SMEM_TOTAL 220032 +#define THREADS 512 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count) : "memory"); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16_cta2( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" + "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" + "}\n" + :: "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x4(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x4.b32" + " {%0, %1, %2, %3}, [%4];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); +} + + +__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ float2 ex2_emulation_f32x2_value(float2 value) { + const float c0 = 1.0f, c1 = 0.695146143436431884765625f; + const float c2 = 0.227564394474029541015625f, c3 = 0.077119089663028717041015625f; + const float magic = 12582912.0f; + float x0 = max_noftz(value.x, -127.0f), x1 = max_noftz(value.y, -127.0f); + float2 xc2 = make_float2(x0, x1), magic2 = make_float2(magic, magic); + float2 xr2; + asm("add.rm.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xr2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&magic2)); + float2 c3_2 = make_float2(c3, c3), c2_2 = make_float2(c2, c2); + float2 c1_2 = make_float2(c1, c1), c0_2 = make_float2(c0, c0); + float2 xrb2, xfrac2; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xrb2) + : "l"(*(unsigned long long*)&xr2), "l"(*(unsigned long long*)&magic2)); + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xfrac2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&xrb2)); + float2 poly2; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&c3_2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c2_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c1_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c0_2)); + int x0r_i, x1r_i, p0_i, p1_i; + asm("mov.b64 {%0, %1}, %2;" : "=r"(x0r_i), "=r"(x1r_i) : "l"(*(unsigned long long*)&xr2)); + asm("mov.b64 {%0, %1}, %2;" : "=r"(p0_i), "=r"(p1_i) : "l"(*(unsigned long long*)&poly2)); + float r0, r1; + asm("mov.b32 %0, %1;" : "=f"(r0) : "r"((x0r_i << 23) + p0_i)); + asm("mov.b32 %0, %1;" : "=f"(r1) : "r"((x1r_i << 23) + p1_i)); + return make_float2(r0, r1); +} + +__device__ __forceinline__ void ex2_emulation_f32x2(float* x0_ptr, float* x1_ptr) { + float2 result = ex2_emulation_f32x2_value(make_float2(*x0_ptr, *x1_ptr)); + *x0_ptr = result.x; *x1_ptr = result.y; +} + +__device__ __forceinline__ void softmax_frag_exp2_cast( + float* sv, uint32_t* pv, int use_emu) +{ + #pragma unroll + for (int j = 0; j < 16; j++) { + if (use_emu && j >= 12) + ex2_emulation_f32x2(&sv[j*2], &sv[j*2+1]); + else { + sv[j*2] = approx_exp2(sv[j*2]); + sv[j*2+1] = approx_exp2(sv[j*2+1]); + } + } + #pragma unroll + for (int j = 0; j < 16; j++) { + __nv_bfloat162 bf = __float22bfloat162_rn({sv[j*2], sv[j*2+1]}); + pv[j] = reinterpret_cast(bf); + } +} + + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_2d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3}], [%4];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr) { + // Canonical .shared::cta form for non-multicast gather4, matching + // trtllm-gen / cuda_ptx and the PTX ISA qualifier order + // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, + // .shared::cluster is reserved for the multicast variant (ctaMask). + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem_cta2( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr) { + // Canonical .shared::cta form; see tma_gather4_gmem2smem above. + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem_mc( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr, unsigned short cta_mask) { + // Multicast variant: the PTX grammar ties the .shared::cluster + // destination to .multicast::cluster + ctaMask (cf. cuda_ptx / + // SM100_TMA_LOAD_MULTICAST_2D_GATHER4 in CUTLASS). + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" + ".mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem_mc_cta2( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr, unsigned short cta_mask) { + // Multicast + cta_group::2 variant; see tma_gather4_gmem2smem_mc. + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" + ".mbarrier::complete_tx::bytes.multicast::cluster.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void tma_store_4d( + const void *tmap, int x, int y, int z, int w, unsigned smem_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.global.shared::cta.tile.bulk_group" + " [%0, {%1, %2, %3, %4}], [%5];" + :: "l"(tmap), "r"(x), "r"(y), "r"(z), "r"(w), "r"(smem_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %1;\n\t" + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" + "}\n" + :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); +} + +extern "C" { + +__global__ __launch_bounds__(512, 1) __cluster_dims__(2,1,1) void +kernel_cake_dsv4_064850af5858dd03c104(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, const __grid_constant__ CUtensorMap O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int mbar_base = smem; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define k_full_addr (mbar_base + 16) + #define k_empty_addr (mbar_base + 40) + #define v_full_addr (mbar_base + 64) + #define v_empty_addr (mbar_base + 80) + #define s_full_addr (mbar_base + 96) + #define s_empty_addr (mbar_base + 112) + #define p_full_addr (mbar_base + 128) + #define p_empty_addr (mbar_base + 144) + #define o_empty_addr (mbar_base + 160) + #define stats_addr (mbar_base + 168) + #define stats_empty_addr (mbar_base + 184) + #define o_full_addr (mbar_base + 200) + #define epi_q_free_addr (mbar_base + 208) + #define o_first_slice_seeded_addr (mbar_base + 216) + #define tmem_dealloc_peer_addr (mbar_base + 224) + #define index_valid_full_addr (mbar_base + 232) + #define index_valid_empty_addr (mbar_base + 248) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 2; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + + // Kernel setup ops + __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_q_addr = smem + 1024; + __nv_bfloat16* smem_k = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_k_addr = smem + 66560; + __nv_bfloat16* smem_v_lo = reinterpret_cast<__nv_bfloat16*>(smem_raw + 115712); + const int smem_v_lo_addr = smem + 115712; + __nv_bfloat16* smem_v_hi = reinterpret_cast<__nv_bfloat16*>(smem_raw + 132096); + const int smem_v_hi_addr = smem + 132096; + __nv_bfloat16* smem_p = reinterpret_cast<__nv_bfloat16*>(smem_raw + 181248); + const int smem_p_addr = smem + 181248; + int* smem_indices = reinterpret_cast(smem_raw + 214016); + const int smem_indices_addr = smem + 214016; + float* smem_softmax_exchange = reinterpret_cast(smem_raw + 218624); + const int smem_softmax_exchange_addr = smem + 218624; + float* smem_epilogue_exchange = reinterpret_cast(smem_raw + 219136); + const int smem_epilogue_exchange_addr = smem + 219136; + unsigned int* smem_index_valid = reinterpret_cast(smem_raw + 219648); + const int smem_index_valid_addr = smem + 219648; + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // Mbarrier init (19 pipeline groups, 0 ordered-sequence groups, 33 barriers) + // Mbarriers at smem_raw[0..264) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // q_full: 1 barriers, init_count=1 + mbarrier_init(smem + 0, 1); + // q_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 8, 1); + // --- pipeline 'k_pipe' --- + // k_full: 3 barriers, init_count=1 + mbarrier_init(smem + 16, 1); + mbarrier_init(smem + 24, 1); + mbarrier_init(smem + 32, 1); + // k_empty: 3 barriers, init_count=1 + mbarrier_init(smem + 40, 1); + mbarrier_init(smem + 48, 1); + mbarrier_init(smem + 56, 1); + // --- pipeline 'v_pipe' --- + // v_full: 2 barriers, init_count=1 + mbarrier_init(smem + 64, 1); + mbarrier_init(smem + 72, 1); + // v_empty: 2 barriers, init_count=1 + mbarrier_init(smem + 80, 1); + mbarrier_init(smem + 88, 1); + // s_full: 2 barriers, init_count=1 + mbarrier_init(smem + 96, 1); + mbarrier_init(smem + 104, 1); + // s_empty: 2 barriers, init_count=256 + mbarrier_init(smem + 112, 256); + mbarrier_init(smem + 120, 256); + // p_full: 2 barriers, init_count=256 + mbarrier_init(smem + 128, 256); + mbarrier_init(smem + 136, 256); + // p_empty: 2 barriers, init_count=1 + mbarrier_init(smem + 144, 1); + mbarrier_init(smem + 152, 1); + // o_empty: 1 barriers, init_count=256 + mbarrier_init(smem + 160, 256); + // stats: 2 barriers, init_count=128 + mbarrier_init(smem + 168, 128); + mbarrier_init(smem + 176, 128); + // stats_empty: 2 barriers, init_count=128 + mbarrier_init(smem + 184, 128); + mbarrier_init(smem + 192, 128); + // o_full: 1 barriers, init_count=1 + mbarrier_init(smem + 200, 1); + // epi_q_free: 1 barriers, init_count=4 + mbarrier_init(smem + 208, 4); + // o_first_slice_seeded: 1 barriers, init_count=256 + mbarrier_init(smem + 216, 256); + // tmem_dealloc_peer: 1 barriers, init_count=32 + mbarrier_init(smem + 224, 32); + // index_valid_full: 2 barriers, init_count=32 + mbarrier_init(smem + 232, 32); + mbarrier_init(smem + 240, 32); + // index_valid_empty: 2 barriers, init_count=128 + mbarrier_init(smem + 248, 128); + mbarrier_init(smem + 256, 128); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 264); + if (warp == 0) { + int _tmem_hold = smem + 264; + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); + } + + __syncthreads(); + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_tmem_scratch = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 15) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 56;"); + } + + // ---- Role: softmax_wg ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); + { // softmax_wg_main + const int softmax_dummy = 0; + float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; + const int head_row_base = warp % 2 * 32; + const int score_row_base = warp % 4 * 32; + const int n_half = warp % 4 / 2; + const int my_row = head_row_base + lane; + const int exchange_idx = n_half * 64 + my_row; + int tile_cursor = 0; + int valid_iter = 0; + float seed_zero[4]; + #pragma unroll + for (int seed_i = 0; seed_i < 4; seed_i++) { + seed_zero[seed_i] = 0.0f; + } + #pragma unroll + for (int seed_stage = 0; seed_stage < 2; seed_stage++) { + #pragma unroll + for (int seed_col = 128 + seed_stage * 128; seed_col < 128 + (seed_stage + 1) * 128; seed_col += 4) { + int seed_addr = taddr + (unsigned int)seed_col + (unsigned int)(score_row_base << 16); + tmem_st_x4_f32(seed_addr, seed_zero); + } + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(o_first_slice_seeded_addr); + int peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_first_slice_seeded_addr), "r"(peer_rank) : "memory"); + #pragma unroll 1 + for (unsigned int work_idx = cluster_id; work_idx < total_work_items; work_idx += num_clusters) { + int split_idx = work_idx % (unsigned int)(((1) ? 4 : 1)); + int query_idx = work_idx / (unsigned int)(((1) ? 4 : 1)); + int _max_0 = ((sparse_topk_lens[query_idx] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx] + sparse_topk_lens_offset) : (0)); + int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); + int active_topk = _min_0; + int num_kv_tiles = ((sparse_topk + 128 - 1) / 128 + ((1) ? 4 : 1) - 1) / ((1) ? 4 : 1); + int first_kv_tile = split_idx * num_kv_tiles; + int head_idx = cta_rank * 64 + my_row; + float row_max_val = -CAKE_INF; + float row_sum_val = 0.0f; + if (has_sinks != 0 && head_idx < num_heads && split_idx == 0) { + row_max_val = sinks[head_idx] * 1.4426950408889634f / softmax_scale_log2; + if (n_half == 0) { + row_sum_val = 1.0f; + } + } + int valid_slot = valid_iter & 1; + int valid_full_phase = valid_iter >> 1 & 1; + mbarrier_wait(index_valid_full_addr + (valid_slot) * 8, valid_full_phase); + int valid_word_base = valid_slot * 36 + n_half * 2; + #pragma unroll 1 + for (int tile = 0; tile < num_kv_tiles; tile++) { + int pipeline_tile = tile_cursor + tile; + int score_phase = pipeline_tile & 1; + int s_full_phase = pipeline_tile >> 1 & 1; + mbarrier_wait(s_full_addr + (score_phase) * 8, s_full_phase); + int score_offset = ((score_phase != 0) ? 64 : 0); + int score_base = taddr + (unsigned int)score_offset + (unsigned int)(score_row_base << 16); + float _tmem_load_0[64]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]) + : "r"(score_base)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]) + : "r"(score_base + 32)); + int valid_cols = active_topk - (first_kv_tile + tile) * 128 - n_half * 64; + if (valid_cols < 0) { + valid_cols = 0; + } + if (valid_cols > 64) { + valid_cols = 64; + } + unsigned int keep_lo = smem_index_valid[valid_word_base + tile * 4]; + unsigned int keep_hi = smem_index_valid[valid_word_base + tile * 4 + 1]; + uint32_t _thresh_mask_0; + { + int _lim_0 = valid_cols; + if (_lim_0 <= 0) { _thresh_mask_0 = 0u; } + else if (_lim_0 >= 32) { _thresh_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_thresh_mask_0) : "r"(_lim_0)); + } + } + uint32_t _thresh_mask_1; + { + int _lim_1 = valid_cols - 32; + if (_lim_1 <= 0) { _thresh_mask_1 = 0u; } + else if (_lim_1 >= 32) { _thresh_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_thresh_mask_1) : "r"(_lim_1)); + } + } + if (!(_thresh_mask_0 & keep_lo & (1u << 0))) _tmem_load_0[0] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 1))) _tmem_load_0[1] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 2))) _tmem_load_0[2] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 3))) _tmem_load_0[3] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 4))) _tmem_load_0[4] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 5))) _tmem_load_0[5] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 6))) _tmem_load_0[6] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 7))) _tmem_load_0[7] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 8))) _tmem_load_0[8] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 9))) _tmem_load_0[9] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 10))) _tmem_load_0[10] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 11))) _tmem_load_0[11] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 12))) _tmem_load_0[12] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 13))) _tmem_load_0[13] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 14))) _tmem_load_0[14] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 15))) _tmem_load_0[15] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 16))) _tmem_load_0[16] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 17))) _tmem_load_0[17] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 18))) _tmem_load_0[18] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 19))) _tmem_load_0[19] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 20))) _tmem_load_0[20] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 21))) _tmem_load_0[21] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 22))) _tmem_load_0[22] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 23))) _tmem_load_0[23] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 24))) _tmem_load_0[24] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 25))) _tmem_load_0[25] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 26))) _tmem_load_0[26] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 27))) _tmem_load_0[27] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 28))) _tmem_load_0[28] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 29))) _tmem_load_0[29] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 30))) _tmem_load_0[30] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 31))) _tmem_load_0[31] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 0))) _tmem_load_0[32] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 1))) _tmem_load_0[33] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 2))) _tmem_load_0[34] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 3))) _tmem_load_0[35] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 4))) _tmem_load_0[36] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 5))) _tmem_load_0[37] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 6))) _tmem_load_0[38] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 7))) _tmem_load_0[39] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 8))) _tmem_load_0[40] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 9))) _tmem_load_0[41] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 10))) _tmem_load_0[42] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 11))) _tmem_load_0[43] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 12))) _tmem_load_0[44] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 13))) _tmem_load_0[45] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 14))) _tmem_load_0[46] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 15))) _tmem_load_0[47] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 16))) _tmem_load_0[48] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 17))) _tmem_load_0[49] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 18))) _tmem_load_0[50] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 19))) _tmem_load_0[51] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 20))) _tmem_load_0[52] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 21))) _tmem_load_0[53] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 22))) _tmem_load_0[54] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 23))) _tmem_load_0[55] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 24))) _tmem_load_0[56] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 25))) _tmem_load_0[57] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 26))) _tmem_load_0[58] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 27))) _tmem_load_0[59] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 28))) _tmem_load_0[60] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 29))) _tmem_load_0[61] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 30))) _tmem_load_0[62] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 31))) _tmem_load_0[63] = -CAKE_INF; + float2 _reg_reduce_max2_2 = {-CAKE_INF, -CAKE_INF}; + row_max_x32_accum(&_tmem_load_0[0], _reg_reduce_max2_2); + row_max_x32_accum(&_tmem_load_0[32], _reg_reduce_max2_2); + float _tmem_load_0_max = row_max_reduce(_reg_reduce_max2_2); + float _max_1 = max_noftz(_tmem_load_0_max, row_max_val); + float new_max = _max_1; + smem_softmax_exchange[exchange_idx] = new_max; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + float _max_2 = max_noftz(new_max, smem_softmax_exchange[exchange_idx ^ 64]); + new_max = _max_2; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + float no_correction = (((new_max - row_max_val) * softmax_scale_log2 <= 0.0f) ? 1.0f : 0.0f); + float delta = softmax_scale_log2 * (row_max_val - new_max); + float _exp2_0 = approx_exp2(delta); + float exp_delta = _exp2_0; + float acc_scale = ((row_max_val > -CAKE_INF) ? exp_delta : 1.0f); + row_max_val = new_max; + float safe_max = ((new_max == -CAKE_INF) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + const float2 _fma_b2_3 = {softmax_scale_log2, softmax_scale_log2}; + const float2 _fma_c2_4 = {-max_scaled, -max_scaled}; + #pragma unroll + for (int _lf = 0; _lf < 32; _lf++) + fma_f32x2_inplace(&reinterpret_cast(_tmem_load_0)[_lf], _fma_b2_3, _fma_c2_4); + #pragma unroll + for (int _le = 0; _le < 64; _le++) { + _tmem_load_0[_le] = approx_exp2(_tmem_load_0[_le]); + } + float2 _reg_reduce_sum2_5 = make_float2(0.0f, 0.0f); + softmax_block_sum(&_tmem_load_0[0], &_reg_reduce_sum2_5); + softmax_block_sum(&_tmem_load_0[32], &_reg_reduce_sum2_5); + float _tmem_load_0_sum = _reg_reduce_sum2_5.x + _reg_reduce_sum2_5.y; + row_sum_val = row_sum_val * acc_scale + _tmem_load_0_sum; + int stats_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + mbarrier_wait(stats_empty_addr + (score_phase) * 8, stats_empty_phase); + float meta[4]; + meta[0] = row_sum_val; + meta[1] = row_max_val; + meta[2] = acc_scale; + meta[3] = no_correction; + int meta_addr = taddr + 384 + (unsigned int)(score_phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(head_row_base << 16); + tmem_st_x4_f32(meta_addr, meta); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (score_phase) * 8); + int p_stage = pipeline_tile % 2; + int p_empty_phase = pipeline_tile / 2 & 1 ^ 1; + mbarrier_wait(p_empty_addr + (p_stage) * 8, p_empty_phase); + uint32_t _tmem_load_0_bf16[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_0[_lp*2 + 0], _tmem_load_0[_lp*2+1 + 0])); + _tmem_load_0_bf16[_lp] = *(uint32_t*)&_bf2; + } + int smem_p_stage = p_stage * 2 + n_half; + int p_base_smem = smem_p_addr + (unsigned int)(smem_p_stage * 8192); + #pragma unroll + for (int vec = 0; vec < 8; vec++) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_base_smem + (my_row * 128 + vec * 16 ^ (my_row * 128 + vec * 16 >> 7 & 7) << 4))), "r"(_tmem_load_0_bf16[vec * 4]), "r"(_tmem_load_0_bf16[vec * 4 + 1]), "r"(_tmem_load_0_bf16[vec * 4 + 2]), "r"(_tmem_load_0_bf16[vec * 4 + 3]) : "memory"); + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(p_full_addr + (unsigned int)(p_stage * 8)), "r"(0) : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(s_empty_addr + (unsigned int)(score_phase * 8)), "r"(0) : "memory"); + } + mbarrier_arrive(index_valid_empty_addr + (valid_slot) * 8); + valid_iter = valid_iter + 1; + tile_cursor = tile_cursor + num_kv_tiles; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + } + } + // ---- Role: correction_wg ---- + if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); + { // correction_wg_main + const int correction_dummy = 0; + float softmax_scale_log2_1 = bmm1_scale[0] * 1.4426950408889634f; + float output_scale = bmm2_scale[0]; + const int head_row_base_1 = warp % 2 * 32; + const int output_row_base = warp % 4 * 32; + const int n_half_1 = warp % 4 / 2; + const int my_row_1 = head_row_base_1 + lane; + const int corr_row = head_row_base_1 << 16; + int tile_cursor_1 = 0; + unsigned int _phase_o_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = cluster_id; work_idx_1 < total_work_items; work_idx_1 += num_clusters) { + int split_idx_1 = work_idx_1 % (unsigned int)(((1) ? 4 : 1)); + int query_idx_1 = work_idx_1 / (unsigned int)(((1) ? 4 : 1)); + int num_kv_tiles_1 = ((sparse_topk + 128 - 1) / 128 + ((1) ? 4 : 1) - 1) / ((1) ? 4 : 1); + float final_sum = 0.0f; + float final_max = -CAKE_INF; + #pragma unroll 1 + for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { + int pipeline_tile_1 = tile_cursor_1 + tile_1; + int score_phase_1 = pipeline_tile_1 & 1; + int stats_phase = pipeline_tile_1 >> 1 & 1; + mbarrier_wait(stats_addr + (score_phase_1) * 8, stats_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int meta_addr_1 = taddr + 384 + (unsigned int)(score_phase_1 * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row; + float _tmem_load_1[4]; + tmem_ld_x4(&_tmem_load_1[0], meta_addr_1); + final_sum = _tmem_load_1[0]; + final_max = _tmem_load_1[1]; + float acc_scale_1 = _tmem_load_1[2]; + float no_correction_1 = _tmem_load_1[3]; + if (tile_1 > 0) { + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _vote_0 = __all_sync(0xFFFFFFFF, no_correction_1 == 1.0f); + int skip_correction = _vote_0; + if (skip_correction == 0) { + #pragma unroll + for (int local_slice = 0; local_slice < 2; local_slice++) { + #pragma unroll + for (int acc_stage = 0; acc_stage < 2; acc_stage++) { + int o_base = taddr + 128 + (unsigned int)(acc_stage * 128) + (unsigned int)(local_slice * 64) + (unsigned int)(output_row_base << 16); + #pragma unroll + for (int c = 0; c < 64; c += 16) { + float _tmem_load_2[16]; + tmem_ld_x16(&_tmem_load_2[0], o_base + c); + const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + mul_f32x2_inplace(&reinterpret_cast(_tmem_load_2)[_ls], _scale2_0); + tmem_st_x16_f32(o_base + c, _tmem_load_2); + } + } + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + } + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + } + mbarrier_arrive(stats_empty_addr + (score_phase_1) * 8); + } + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + const int exchange_idx_1 = n_half_1 * 64 + my_row_1; + smem_epilogue_exchange[exchange_idx_1] = final_sum; + asm volatile("barrier.sync 4, 128;" ::: "memory"); + final_sum = final_sum + smem_epilogue_exchange[exchange_idx_1 ^ 64]; + asm volatile("barrier.sync 4, 128;" ::: "memory"); + float _rcp_0 = approx_rcp(final_sum); + float inv_sum = ((final_sum > 0.0f) ? _rcp_0 : 0.0f); + int head_idx_1 = cta_rank * 64 + my_row_1; + int output_offset = ((1) ? ((query_idx_1 * num_heads + head_idx_1) * ((1) ? 4 : 1) + split_idx_1) * 512 : (query_idx_1 * num_heads + head_idx_1) * 512); + if (n_half_1 == 0 && head_idx_1 < num_heads) { + int stat_offset = (query_idx_1 * num_heads + head_idx_1) * ((1) ? 4 : 1) + split_idx_1; + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(final_sum)); + partial_lse[stat_offset] = ((final_sum > 0.0f) ? final_max * softmax_scale_log2_1 + _log2_0 : -CAKE_INF); + } + float epi_scale = inv_sum * output_scale; + int epi_box_base = smem_q_addr + (unsigned int)(warp % 4 * 16384); + int epi_row0 = query_idx_1 * num_heads + cta_rank * 64 + head_row_base_1; + #pragma unroll + for (int acc_stage_1 = 0; acc_stage_1 < 2; acc_stage_1++) { + #pragma unroll + for (int local_slice_1 = 0; local_slice_1 < 2; local_slice_1++) { + int logical_slice = acc_stage_1 * 2 * 2 + n_half_1 * 2 + local_slice_1; + int epi_box = epi_box_base + (acc_stage_1 * 2 + local_slice_1) * 4096; + int o_base_1 = taddr + 128 + (unsigned int)(acc_stage_1 * 128) + (unsigned int)(local_slice_1 * 64) + (unsigned int)(output_row_base << 16); + #pragma unroll + for (int c_1 = 0; c_1 < 64; c_1 += 32) { + float _tmem_load_3[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]) + : "r"(o_base_1 + c_1)); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + float epi_scaled[32]; + #pragma unroll + for (int e = 0; e < 32; e++) { + epi_scaled[e] = _tmem_load_3[e] * epi_scale; + } + unsigned int epi_packed[16]; + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(epi_scaled[_lp*2 + 0], epi_scaled[_lp*2+1 + 0])); + epi_packed[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((epi_box + (lane * 128 + c_1 * 2 ^ (lane * 128 + c_1 * 2 >> 7 & 7) << 4))), "r"(epi_packed[0]), "r"(epi_packed[1]), "r"(epi_packed[2]), "r"(epi_packed[3]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((epi_box + (lane * 128 + (c_1 * 2 + 16) ^ (lane * 128 + (c_1 * 2 + 16) >> 7 & 7) << 4))), "r"(epi_packed[4]), "r"(epi_packed[5]), "r"(epi_packed[6]), "r"(epi_packed[7]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((epi_box + (lane * 128 + (c_1 * 2 + 32) ^ (lane * 128 + (c_1 * 2 + 32) >> 7 & 7) << 4))), "r"(epi_packed[8]), "r"(epi_packed[9]), "r"(epi_packed[10]), "r"(epi_packed[11]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((epi_box + (lane * 128 + (c_1 * 2 + 48) ^ (lane * 128 + (c_1 * 2 + 48) >> 7 & 7) << 4))), "r"(epi_packed[12]), "r"(epi_packed[13]), "r"(epi_packed[14]), "r"(epi_packed[15]) : "memory"); + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + __syncwarp(); + if (elect_sync()) { + tma_store_4d((&O), 0, logical_slice, split_idx_1, epi_row0, epi_box); + asm volatile("cp.async.bulk.commit_group;"); + } + } + } + if (elect_sync()) { + asm volatile("cp.async.bulk.wait_group.read 0;"); + mbarrier_arrive(epi_q_free_addr); + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + tile_cursor_1 = tile_cursor_1 + num_kv_tiles_1; + } + if (elect_sync()) { + asm volatile("cp.async.bulk.wait_group 0;"); + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + } + } + // ---- Role: mma_warp ---- + if (warp == 8) { + { // mma_warp_main + const int mma_dummy = 0; + unsigned int k_stage = 0; + unsigned int v_stage = 0; + int tile_cursor_2 = 0; + unsigned int _phase_o_first_slice_seeded_0 = 0; + mbarrier_wait(o_first_slice_seeded_addr, _phase_o_first_slice_seeded_0); + _phase_o_first_slice_seeded_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + unsigned int _phase_q_full_0 = 0; + unsigned int _phase_k_full = 0; + unsigned int _phase_o_empty_0 = 1; + unsigned int _phase_v_full = 0; + #pragma unroll 1 + for (unsigned int work_idx_2 = cluster_id; work_idx_2 < total_work_items; work_idx_2 += num_clusters) { + int num_kv_tiles_2 = ((sparse_topk + 128 - 1) / 128 + ((1) ? 4 : 1) - 1) / ((1) ? 4 : 1); + if (cta_rank == 0) { + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + int first_pv = 1; + #pragma unroll 1 + for (int tile_2 = 0; tile_2 < num_kv_tiles_2; tile_2++) { + int pipeline_tile_2 = tile_cursor_2 + tile_2; + int score_phase_2 = pipeline_tile_2 & 1; + int s_empty_phase = pipeline_tile_2 >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (score_phase_2) * 8, s_empty_phase); + int score_col = ((score_phase_2 != 0) ? 64 : 0); + #pragma unroll + for (int qk_stage = 0; qk_stage < 4; qk_stage++) { + mbarrier_wait(k_full_addr + (k_stage) * 8, _phase_k_full); + int _mma_a_lo_0 = (((smem_q_addr) >> 4) & 0x3FFF) + (qk_stage * 2) * 512; + int _mma_b_lo_0 = (((smem_k_addr) >> 4) & 0x3FFF) + (k_stage) * 1024; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (score_col))), "r"(((qk_stage == 0) ? 0 : 1))); + int _mma_a_lo_1 = (((smem_q_addr) >> 4) & 0x3FFF) + (qk_stage * 2 + 1) * 512; + int _mma_b_lo_1 = (((smem_k_addr + 8192) >> 4) & 0x3FFF) + (k_stage) * 1024; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + elect_commit_cg2_multicast(k_empty_addr + (k_stage) * 8, (uint16_t)(3)); + k_stage += 1; + if (k_stage == 3) { k_stage = 0; _phase_k_full ^= 1; } + } + elect_commit_cg2_multicast(s_full_addr + (score_phase_2) * 8, (uint16_t)(3)); + if (tile_2 > 0) { + int prev_pipeline_tile = pipeline_tile_2 - 1; + int p_stage_1 = prev_pipeline_tile % 2; + int p_full_phase = prev_pipeline_tile / 2 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (p_stage_1) * 8, p_full_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps = 0; ps < 2; ps++) { + int smem_p_stage_1 = p_stage_1 * 2 + ps; + mbarrier_wait(v_full_addr + (v_stage) * 8, _phase_v_full); + #pragma unroll + for (int acc_stage_2 = 0; acc_stage_2 < 2; acc_stage_2++) { + int output_col = 128 + acc_stage_2 * 128; + if (acc_stage_2 == 0) { + int _mma_a_lo_2 = (((smem_p_addr) >> 4) & 0x3FFF) + (smem_p_stage_1) * 512; + int _mma_b_lo_2 = ((((smem_v_lo_addr) >> 4) & 0x3FFF) | 0x2000000) + (v_stage) * 2048; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (output_col))), "r"(((((ps == 0) ? first_pv : 0)) ? 0 : 1))); + } else { + int _mma_a_lo_3 = (((smem_p_addr) >> 4) & 0x3FFF) + (smem_p_stage_1) * 512; + int _mma_b_lo_3 = ((((smem_v_hi_addr) >> 4) & 0x3FFF) | 0x2000000) + (v_stage) * 2048; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"((tmem_tmem_scratch + (output_col))), "r"(((((ps == 0) ? first_pv : 0)) ? 0 : 1))); + } + } + elect_commit_cg2_multicast(v_empty_addr + (v_stage) * 8, (uint16_t)(3)); + v_stage += 1; + if (v_stage == 2) { v_stage = 0; _phase_v_full ^= 1; } + } + first_pv = 0; + elect_commit_cg2_multicast(p_empty_addr + (p_stage_1) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + } + } + int last_pipeline_tile = tile_cursor_2 + num_kv_tiles_2 - 1; + int last_p_stage = last_pipeline_tile % 2; + int last_p_phase = last_pipeline_tile / 2 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (last_p_stage) * 8, last_p_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_1 = 0; ps_1 < 2; ps_1++) { + int smem_p_stage_2 = last_p_stage * 2 + ps_1; + mbarrier_wait(v_full_addr + (v_stage) * 8, _phase_v_full); + #pragma unroll + for (int acc_stage_3 = 0; acc_stage_3 < 2; acc_stage_3++) { + int output_col_1 = 128 + acc_stage_3 * 128; + if (acc_stage_3 == 0) { + int _mma_a_lo_4 = (((smem_p_addr) >> 4) & 0x3FFF) + (smem_p_stage_2) * 512; + int _mma_b_lo_4 = ((((smem_v_lo_addr) >> 4) & 0x3FFF) | 0x2000000) + (v_stage) * 2048; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"((tmem_tmem_scratch + (output_col_1))), "r"(((((ps_1 == 0) ? first_pv : 0)) ? 0 : 1))); + } else { + int _mma_a_lo_5 = (((smem_p_addr) >> 4) & 0x3FFF) + (smem_p_stage_2) * 512; + int _mma_b_lo_5 = ((((smem_v_hi_addr) >> 4) & 0x3FFF) | 0x2000000) + (v_stage) * 2048; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_5), "r"(_mma_b_lo_5), "r"((tmem_tmem_scratch + (output_col_1))), "r"(((((ps_1 == 0) ? first_pv : 0)) ? 0 : 1))); + } + } + elect_commit_cg2_multicast(v_empty_addr + (v_stage) * 8, (uint16_t)(3)); + v_stage += 1; + if (v_stage == 2) { v_stage = 0; _phase_v_full ^= 1; } + } + elect_commit_cg2_multicast(p_empty_addr + (last_p_stage) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + } + tile_cursor_2 = tile_cursor_2 + num_kv_tiles_2; + } + if (cta_rank == 0) { + #pragma unroll + for (int tail_offset = 0; tail_offset < 2; tail_offset++) { + int tail_tile = tile_cursor_2 + tail_offset; + int tail_stage = tail_tile & 1; + int tail_phase = tail_tile >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (tail_stage) * 8, tail_phase); + } + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + } + asm volatile("tcgen05.fence::before_thread_sync;"); + int peer_rank_1 = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_dealloc_peer_addr), "r"(peer_rank_1) : "memory"); + mbarrier_wait(tmem_dealloc_peer_addr, 0); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + } + // ---- Role: load_wg ---- + if (warp >= 9 && warp <= 12) { + { // load_wg_main + const int load_dummy = 0; + const int load_warp_rank = warp - 9; + unsigned int k_stage_1 = 0; + unsigned int v_stage_1 = 0; + unsigned int _phase_q_empty_0 = 1; + unsigned int _phase_epi_q_free_0 = 1; + unsigned int _phase_k_empty = 1; + unsigned int _phase_v_empty = 1; + #pragma unroll 1 + for (unsigned int work_idx_3 = cluster_id; work_idx_3 < total_work_items; work_idx_3 += num_clusters) { + int split_idx_2 = work_idx_3 % (unsigned int)(((1) ? 4 : 1)); + int query_idx_2 = work_idx_3 / (unsigned int)(((1) ? 4 : 1)); + int num_kv_tiles_3 = ((sparse_topk + 128 - 1) / 128 + ((1) ? 4 : 1) - 1) / ((1) ? 4 : 1); + int first_kv_tile_1 = split_idx_2 * num_kv_tiles_3; + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + mbarrier_wait(epi_q_free_addr, _phase_epi_q_free_0); + _phase_epi_q_free_0 ^= 1; + if (load_warp_rank == 0) { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(131072)) : "memory"); + } + } + if (elect_sync()) { + #pragma unroll + for (int q_stage = 0; q_stage < 8; q_stage++) { + tma_4d_gmem2smem_cta2(smem_q_addr + (unsigned int)(q_stage * 8192), (&tmap_q), 0, cta_rank * 64, q_stage, query_idx_2, ((q_full_addr) & 0xFEFFFFFF)); + } + } + } + asm volatile("barrier.sync 2, 160;" ::: "memory"); + #pragma unroll + for (int qk_stage_1 = 0; qk_stage_1 < 4; qk_stage_1++) { + mbarrier_wait(k_empty_addr + (k_stage_1) * 8, _phase_k_empty); + if (load_warp_rank == 0) { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + } + } + } + int group = load_warp_rank * 4 + lane; + if (lane < 4) { + int index_offset = cta_rank * 64 + group * 4; + int raw_rows[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows[0])), "=r"(*reinterpret_cast(&raw_rows[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows[(0) + 3])) + : "r"(smem_indices_addr + (unsigned int)(index_offset * 4))); + if (first_kv_tile_1 == 0) { + tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group * 512), (&tmap_swa_k), qk_stage_1 * 128, ((raw_rows[0] >= 0) ? raw_rows[0] : ((1) ? group * 4 : 0)), ((raw_rows[1] >= 0) ? raw_rows[1] : ((1) ? group * 4 + 1 : 0)), ((raw_rows[2] >= 0) ? raw_rows[2] : ((1) ? group * 4 + 2 : 0)), ((raw_rows[3] >= 0) ? raw_rows[3] : ((1) ? group * 4 + 3 : 0)), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group * 512) + 8192, (&tmap_swa_k), qk_stage_1 * 128 + 64, ((raw_rows[0] >= 0) ? raw_rows[0] : ((1) ? group * 4 : 0)), ((raw_rows[1] >= 0) ? raw_rows[1] : ((1) ? group * 4 + 1 : 0)), ((raw_rows[2] >= 0) ? raw_rows[2] : ((1) ? group * 4 + 2 : 0)), ((raw_rows[3] >= 0) ? raw_rows[3] : ((1) ? group * 4 + 3 : 0)), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + } else { + tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group * 512), (&tmap_compressed_k), qk_stage_1 * 128, ((raw_rows[0] >= 0) ? raw_rows[0] : ((1) ? group * 4 : 0)), ((raw_rows[1] >= 0) ? raw_rows[1] : ((1) ? group * 4 + 1 : 0)), ((raw_rows[2] >= 0) ? raw_rows[2] : ((1) ? group * 4 + 2 : 0)), ((raw_rows[3] >= 0) ? raw_rows[3] : ((1) ? group * 4 + 3 : 0)), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group * 512) + 8192, (&tmap_compressed_k), qk_stage_1 * 128 + 64, ((raw_rows[0] >= 0) ? raw_rows[0] : ((1) ? group * 4 : 0)), ((raw_rows[1] >= 0) ? raw_rows[1] : ((1) ? group * 4 + 1 : 0)), ((raw_rows[2] >= 0) ? raw_rows[2] : ((1) ? group * 4 + 2 : 0)), ((raw_rows[3] >= 0) ? raw_rows[3] : ((1) ? group * 4 + 3 : 0)), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + } + } + k_stage_1 += 1; + if (k_stage_1 == 3) { k_stage_1 = 0; _phase_k_empty ^= 1; } + } + #pragma unroll 1 + for (int tile_3 = 1; tile_3 < num_kv_tiles_3; tile_3++) { + #pragma unroll + for (int qk_stage_2 = 0; qk_stage_2 < 4; qk_stage_2++) { + mbarrier_wait(k_empty_addr + (k_stage_1) * 8, _phase_k_empty); + if (load_warp_rank == 0) { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + } + } + } + int group_1 = load_warp_rank * 4 + lane; + if (lane < 4) { + int index_offset_1 = tile_3 * 128 + cta_rank * 64 + group_1 * 4; + int raw_rows_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_1[0])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 3])) + : "r"(smem_indices_addr + (unsigned int)(index_offset_1 * 4))); + tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group_1 * 512), (&tmap_compressed_k), qk_stage_2 * 128, ((raw_rows_1[0] >= 0) ? raw_rows_1[0] : ((1) ? group_1 * 4 : 0)), ((raw_rows_1[1] >= 0) ? raw_rows_1[1] : ((1) ? group_1 * 4 + 1 : 0)), ((raw_rows_1[2] >= 0) ? raw_rows_1[2] : ((1) ? group_1 * 4 + 2 : 0)), ((raw_rows_1[3] >= 0) ? raw_rows_1[3] : ((1) ? group_1 * 4 + 3 : 0)), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group_1 * 512) + 8192, (&tmap_compressed_k), qk_stage_2 * 128 + 64, ((raw_rows_1[0] >= 0) ? raw_rows_1[0] : ((1) ? group_1 * 4 : 0)), ((raw_rows_1[1] >= 0) ? raw_rows_1[1] : ((1) ? group_1 * 4 + 1 : 0)), ((raw_rows_1[2] >= 0) ? raw_rows_1[2] : ((1) ? group_1 * 4 + 2 : 0)), ((raw_rows_1[3] >= 0) ? raw_rows_1[3] : ((1) ? group_1 * 4 + 3 : 0)), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + } + k_stage_1 += 1; + if (k_stage_1 == 3) { k_stage_1 = 0; _phase_k_empty ^= 1; } + } + int prev_tile = tile_3 - 1; + #pragma unroll + for (int ps_2 = 0; ps_2 < 2; ps_2++) { + mbarrier_wait(v_empty_addr + (v_stage_1) * 8, _phase_v_empty); + if (load_warp_rank == 0) { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(65536)) : "memory"); + } + } + } + int group_2 = load_warp_rank * 4 + lane; + if (lane < 4) { + int index_offset_2 = prev_tile * 128 + ps_2 * 64 + group_2 * 4; + int raw_rows_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_2[0])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 3])) + : "r"(smem_indices_addr + (unsigned int)(index_offset_2 * 4))); + #pragma unroll + for (int acc_stage_4 = 0; acc_stage_4 < 2; acc_stage_4++) { + int dst_v = ((acc_stage_4 == 0) ? smem_v_lo_addr + v_stage_1 * 32768 : smem_v_hi_addr + v_stage_1 * 32768); + if (first_kv_tile_1 + prev_tile == 0) { + tma_gather4_gmem2smem_mc_cta2(dst_v + group_2 * 512, (&tmap_swa_v), acc_stage_4 * 128 * 2 + cta_rank * 128, ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : ((1) ? group_2 * 4 : 0)), ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : ((1) ? group_2 * 4 + 1 : 0)), ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : ((1) ? group_2 * 4 + 2 : 0)), ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : ((1) ? group_2 * 4 + 3 : 0)), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + tma_gather4_gmem2smem_mc_cta2(dst_v + group_2 * 512 + 8192, (&tmap_swa_v), acc_stage_4 * 128 * 2 + cta_rank * 128 + 64, ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : ((1) ? group_2 * 4 : 0)), ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : ((1) ? group_2 * 4 + 1 : 0)), ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : ((1) ? group_2 * 4 + 2 : 0)), ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : ((1) ? group_2 * 4 + 3 : 0)), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + } else { + tma_gather4_gmem2smem_mc_cta2(dst_v + group_2 * 512, (&tmap_compressed_v), acc_stage_4 * 128 * 2 + cta_rank * 128, ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : ((1) ? group_2 * 4 : 0)), ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : ((1) ? group_2 * 4 + 1 : 0)), ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : ((1) ? group_2 * 4 + 2 : 0)), ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : ((1) ? group_2 * 4 + 3 : 0)), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + tma_gather4_gmem2smem_mc_cta2(dst_v + group_2 * 512 + 8192, (&tmap_compressed_v), acc_stage_4 * 128 * 2 + cta_rank * 128 + 64, ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : ((1) ? group_2 * 4 : 0)), ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : ((1) ? group_2 * 4 + 1 : 0)), ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : ((1) ? group_2 * 4 + 2 : 0)), ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : ((1) ? group_2 * 4 + 3 : 0)), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + } + } + } + v_stage_1 += 1; + if (v_stage_1 == 2) { v_stage_1 = 0; _phase_v_empty ^= 1; } + } + } + int last_tile = num_kv_tiles_3 - 1; + #pragma unroll + for (int ps_3 = 0; ps_3 < 2; ps_3++) { + mbarrier_wait(v_empty_addr + (v_stage_1) * 8, _phase_v_empty); + if (load_warp_rank == 0) { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(65536)) : "memory"); + } + } + } + int group_3 = load_warp_rank * 4 + lane; + if (lane < 4) { + int index_offset_3 = last_tile * 128 + ps_3 * 64 + group_3 * 4; + int raw_rows_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_3[0])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 3])) + : "r"(smem_indices_addr + (unsigned int)(index_offset_3 * 4))); + #pragma unroll + for (int acc_stage_5 = 0; acc_stage_5 < 2; acc_stage_5++) { + int dst_v_1 = ((acc_stage_5 == 0) ? smem_v_lo_addr + v_stage_1 * 32768 : smem_v_hi_addr + v_stage_1 * 32768); + if (first_kv_tile_1 + last_tile == 0) { + tma_gather4_gmem2smem_mc_cta2(dst_v_1 + group_3 * 512, (&tmap_swa_v), acc_stage_5 * 128 * 2 + cta_rank * 128, ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : ((1) ? group_3 * 4 : 0)), ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : ((1) ? group_3 * 4 + 1 : 0)), ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : ((1) ? group_3 * 4 + 2 : 0)), ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : ((1) ? group_3 * 4 + 3 : 0)), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + tma_gather4_gmem2smem_mc_cta2(dst_v_1 + group_3 * 512 + 8192, (&tmap_swa_v), acc_stage_5 * 128 * 2 + cta_rank * 128 + 64, ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : ((1) ? group_3 * 4 : 0)), ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : ((1) ? group_3 * 4 + 1 : 0)), ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : ((1) ? group_3 * 4 + 2 : 0)), ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : ((1) ? group_3 * 4 + 3 : 0)), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + } else { + tma_gather4_gmem2smem_mc_cta2(dst_v_1 + group_3 * 512, (&tmap_compressed_v), acc_stage_5 * 128 * 2 + cta_rank * 128, ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : ((1) ? group_3 * 4 : 0)), ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : ((1) ? group_3 * 4 + 1 : 0)), ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : ((1) ? group_3 * 4 + 2 : 0)), ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : ((1) ? group_3 * 4 + 3 : 0)), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + tma_gather4_gmem2smem_mc_cta2(dst_v_1 + group_3 * 512 + 8192, (&tmap_compressed_v), acc_stage_5 * 128 * 2 + cta_rank * 128 + 64, ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : ((1) ? group_3 * 4 : 0)), ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : ((1) ? group_3 * 4 + 1 : 0)), ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : ((1) ? group_3 * 4 + 2 : 0)), ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : ((1) ? group_3 * 4 + 3 : 0)), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + } + } + } + v_stage_1 += 1; + if (v_stage_1 == 2) { v_stage_1 = 0; _phase_v_empty ^= 1; } + } + asm volatile("barrier.sync 2, 160;" ::: "memory"); + } + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + } + } + // ---- Role: index_warp ---- + if (warp == 13) { + { // index_warp_main + const int index_dummy = 0; + int valid_iter_1 = 0; + #pragma unroll 1 + for (unsigned int work_idx_4 = cluster_id; work_idx_4 < total_work_items; work_idx_4 += num_clusters) { + int split_idx_3 = work_idx_4 % (unsigned int)(((1) ? 4 : 1)); + int query_idx_3 = work_idx_4 / (unsigned int)(((1) ? 4 : 1)); + int owner_col_base = split_idx_3 * (4 / ((1) ? 4 : 1)) * 128; + int* swa_row = swa_indices + (query_idx_3 * swa_index_stride); + int* compressed_row = compressed_indices + (query_idx_3 * compressed_index_stride); + int valid_slot_1 = valid_iter_1 & 1; + int valid_empty_phase = valid_iter_1 >> 1 & 1 ^ 1; + mbarrier_wait(index_valid_empty_addr + (valid_slot_1) * 8, valid_empty_phase); + #pragma unroll + for (int index_pass = 0; index_pass < 4 / ((1) ? 4 : 1); index_pass++) { + int index_offset_4 = index_pass * 128 + lane * 4; + int combined_col = owner_col_base + index_offset_4; + int* index_src = ((combined_col < 128) ? swa_row + combined_col : compressed_row + (combined_col - 128)); + int values[4]; + #pragma unroll + for (int i = 0; i < 4; i++) { + values[i] = -1; + } + if (combined_col + 4 <= sparse_topk) { + int _vec_load_0[4]; + { + const int4* _ivptr_0 = reinterpret_cast(index_src + 0); + int4 _ivld_0; + _ivld_0 = *_ivptr_0; + _vec_load_0[0 + 0] = _ivld_0.x; + _vec_load_0[0 + 1] = _ivld_0.y; + _vec_load_0[0 + 2] = _ivld_0.z; + _vec_load_0[0 + 3] = _ivld_0.w; + } + #pragma unroll + for (int i_1 = 0; i_1 < 4; i_1++) { + values[i_1] = _vec_load_0[i_1]; + } + } + if (combined_col < sparse_topk && combined_col + 4 > sparse_topk) { + #pragma unroll + for (int i_2 = 0; i_2 < 4; i_2++) { + if (combined_col + i_2 < sparse_topk) { + values[i_2] = index_src[i_2]; + } + } + } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(smem_indices_addr + (unsigned int)(index_offset_4 * 4)), "r"(values[0]), "r"(values[1]), "r"(values[2]), "r"(values[3]) : "memory"); + unsigned int valid_word = (((values[0] >= 0) ? 1 : 0) | ((values[1] >= 0) ? 1 : 0) << 1 | ((values[2] >= 0) ? 1 : 0) << 2 | ((values[3] >= 0) ? 1 : 0) << 3) << (lane & 7) * 4; + unsigned int _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, valid_word, 1); + valid_word = valid_word | _shfl_xor_0; + unsigned int _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, valid_word, 2); + valid_word = valid_word | _shfl_xor_1; + unsigned int _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, valid_word, 4); + valid_word = valid_word | _shfl_xor_2; + if ((lane & 7) == 0) { + smem_index_valid[valid_slot_1 * 36 + index_pass * 4 + (lane >> 3)] = valid_word; + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(index_valid_full_addr + (valid_slot_1) * 8); + asm volatile("barrier.sync 2, 160;" ::: "memory"); + asm volatile("barrier.sync 2, 160;" ::: "memory"); + valid_iter_1 = valid_iter_1 + 1; + } + } + } + // ---- Role: empty_warps ---- + if (warp >= 14 && warp <= 15) { + { // empty_warps_main + const int empty_dummy = 0; + #pragma unroll 1 + for (unsigned int work_idx_5 = cluster_id; work_idx_5 < total_work_items; work_idx_5 += num_clusters) { + asm volatile("tcgen05.fence::after_thread_sync;"); + } + } + } + + // Cleanup +} + +} // extern "C" diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_0a925b9bcb3044777c76_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_0a925b9bcb3044777c76_binding.cu new file mode 100644 index 00000000000..f5e9697f131 --- /dev/null +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_0a925b9bcb3044777c76_binding.cu @@ -0,0 +1,441 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +// Generated by the Cake source exporter — do not edit. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_0a925b9bcb3044777c76'. +#include +#include +#include +#include + +#include "tvm_ffi_utils.h" + +#include +#include +#include +#include +#include +#include +#include + +extern "C" __global__ void kernel_cake_dsv4_0a925b9bcb3044777c76(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ cum_seq_lens_q, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); + + +namespace cake_host_shim_52ed16aa7da8966c { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t CakeDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +// 4D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + int64_t d3 = t.size(t.ndim() - 3); + TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + TVM_FFI_CHECK(d2 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d2 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(d1), (uint64_t)(64), (uint64_t)((d2 / 64)), (uint64_t)(d3)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (128, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = (64 * d1); + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + (uint64_t)((carrier_stride_1 * 8) / 8), + (uint64_t)((carrier_stride_2 * 8) / 8), + }; + uint32_t box_dim[4] = {128u, 64u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_L2_128B, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_swa_kv' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_swa_kv' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; + TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (128, 1) exceeds resolved global dims for 'tmap_swa_kv'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + }; + uint32_t box_dim[2] = {128u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_compressed_kv' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_compressed_kv(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_compressed_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_compressed_kv' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_compressed_kv' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_compressed_kv"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_compressed_kv' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved a non-positive global dim"; + TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (128, 1) exceeds resolved global dims for 'tmap_compressed_kv'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + }; + uint32_t box_dim[2] = {128u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_kv') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_compressed_kv, TensorView arg_O, TensorView arg_cum_seq_lens_q, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 1, 8, 1); + CheckContiguous(arg_tmap_q, "tmap_q"); + CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); + CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 1, 8, 1); + CheckCudaTensor(arg_tmap_compressed_kv, "tmap_compressed_kv"); + CheckDtype(arg_tmap_compressed_kv, "tmap_compressed_kv", 1, 8, 1); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_cum_seq_lens_q, "cum_seq_lens_q"); + CheckDtype(arg_cum_seq_lens_q, "cum_seq_lens_q", 0, 32, 1); + CheckContiguous(arg_cum_seq_lens_q, "cum_seq_lens_q"); + CheckCudaTensor(arg_swa_indices, "swa_indices"); + CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); + CheckContiguous(arg_swa_indices, "swa_indices"); + CheckCudaTensor(arg_compressed_indices, "compressed_indices"); + CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); + CheckContiguous(arg_compressed_indices, "compressed_indices"); + CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); + CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); + CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); + CheckContiguous(arg_bmm1_scale, "bmm1_scale"); + CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); + CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); + CheckContiguous(arg_bmm2_scale, "bmm2_scale"); + TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) + << "scalar 'num_heads' value " << arg_num_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) + << "scalar 'swa_index_stride' value " << arg_swa_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) + << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) + << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) + << "scalar 'sparse_topk' value " << arg_sparse_topk + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) + << "scalar 'has_sinks' value " << arg_has_sinks + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); + CheckSameCudaDevice(arg_tmap_compressed_kv, arg_tmap_q, "tmap_compressed_kv", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_cum_seq_lens_q, arg_tmap_q, "cum_seq_lens_q", "tmap_q"); + CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); + CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); + CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); + CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + + cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); + CUtensorMap p_tmap_compressed_kv = EncodeTma_tmap_compressed_kv(arg_tmap_compressed_kv); + __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); + int* p_cum_seq_lens_q = static_cast(arg_cum_seq_lens_q.data_ptr()); + int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); + int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); + int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); + float* p_sinks = static_cast(arg_sinks.data_ptr()); + float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); + float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); + int32_t v_num_heads = (int32_t)arg_num_heads; + int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; + int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; + int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; + int32_t v_sparse_topk = (int32_t)arg_sparse_topk; + int32_t v_has_sinks = (int32_t)arg_has_sinks; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_tmap_compressed_kv, &p_O, &p_cum_seq_lens_q, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_has_sinks, &v_total_work_items}; + + static std::mutex smem_mu; + static auto* smem_status_by_device = + new std::unordered_map(); + cudaError_t smem_status = cudaSuccess; + { + std::lock_guard lock(smem_mu); + auto it = smem_status_by_device->find(dev.device_id); + if (it == smem_status_by_device->end()) { + int current_device = -1; + smem_status = cudaGetDevice(¤t_device); + if (smem_status == cudaSuccess && current_device != dev.device_id) { + smem_status = cudaErrorInvalidDevice; + } + if (smem_status == cudaSuccess) { + smem_status = cudaFuncSetAttribute( + reinterpret_cast(kernel_cake_dsv4_0a925b9bcb3044777c76), + cudaFuncAttributeMaxDynamicSharedMemorySize, + 203136); + } + smem_status_by_device->emplace(dev.device_id, smem_status); + } else { + smem_status = it->second; + } + } + TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) + << "cudaFuncSetAttribute for kernel_cake_dsv4_0a925b9bcb3044777c76 failed: " + << cudaGetErrorString(smem_status); + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(512u, 1u, 1u); + + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes = 203136u; + config.stream = stream; + cudaError_t launch_status = cudaLaunchKernelExC( + &config, reinterpret_cast(kernel_cake_dsv4_0a925b9bcb3044777c76), kargs); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernelExC for kernel_cake_dsv4_0a925b9bcb3044777c76 failed: " + << cudaGetErrorString(launch_status); +} + +} // namespace cake_host_shim_52ed16aa7da8966c + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_52ed16aa7da8966c::Run); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_0a925b9bcb3044777c76_kernel.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_0a925b9bcb3044777c76_kernel.cu new file mode 100644 index 00000000000..6b6f53ff26e --- /dev/null +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_0a925b9bcb3044777c76_kernel.cu @@ -0,0 +1,2967 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +#define CAKE_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_OFFSET 0 +#define NUM_Q_PIPE_STAGES 2 +#define NUM_KV_PIPE_STAGES 4 +#define NUM_PAGE_PIPE_STAGES 6 +#define NUM_SCORE_PIPE_STAGES 2 +#define NUM_STATS_PIPE_STAGES 2 +#define NUM_WORK_PIPE_STAGES 2 +#define NUM_THROTTLE_PIPE_STAGES 2 +#define SMEM_SMEM_Q_OFF 0 +#define SMEM_SMEM_Q_STAGE_BYTES 32768 +#define SMEM_SMEM_Q_STRIDE 32768 +#define SMEM_SMEM_K_OFF 65536 +#define SMEM_SMEM_K_STAGE_BYTES 32768 +#define SMEM_SMEM_K_STRIDE 32768 +#define SMEM_SMEM_V_OFF 65536 +#define SMEM_SMEM_V_STAGE_BYTES 32768 +#define SMEM_SMEM_V_STRIDE 32768 +#define SMEM_SMEM_PAGE_OFFSETS_OFF 196608 +#define SMEM_SMEM_PAGE_OFFSETS_STAGE_BYTES 1024 +#define SMEM_SMEM_PAGE_OFFSETS_STRIDE 1024 +#define SMEM_WORK_RESPONSE_OFF 202752 +#define SMEM_WORK_RESPONSE_STAGE_BYTES 16 +#define SMEM_WORK_RESPONSE_STRIDE 16 +#define SMEM_TOTAL 203136 +#define THREADS 512 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count) : "memory"); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f8f6f4( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ float2 ex2_emulation_f32x2_value(float2 value) { + const float c0 = 1.0f, c1 = 0.695146143436431884765625f; + const float c2 = 0.227564394474029541015625f, c3 = 0.077119089663028717041015625f; + const float magic = 12582912.0f; + float x0 = max_noftz(value.x, -127.0f), x1 = max_noftz(value.y, -127.0f); + float2 xc2 = make_float2(x0, x1), magic2 = make_float2(magic, magic); + float2 xr2; + asm("add.rm.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xr2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&magic2)); + float2 c3_2 = make_float2(c3, c3), c2_2 = make_float2(c2, c2); + float2 c1_2 = make_float2(c1, c1), c0_2 = make_float2(c0, c0); + float2 xrb2, xfrac2; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xrb2) + : "l"(*(unsigned long long*)&xr2), "l"(*(unsigned long long*)&magic2)); + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xfrac2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&xrb2)); + float2 poly2; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&c3_2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c2_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c1_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c0_2)); + int x0r_i, x1r_i, p0_i, p1_i; + asm("mov.b64 {%0, %1}, %2;" : "=r"(x0r_i), "=r"(x1r_i) : "l"(*(unsigned long long*)&xr2)); + asm("mov.b64 {%0, %1}, %2;" : "=r"(p0_i), "=r"(p1_i) : "l"(*(unsigned long long*)&poly2)); + float r0, r1; + asm("mov.b32 %0, %1;" : "=f"(r0) : "r"((x0r_i << 23) + p0_i)); + asm("mov.b32 %0, %1;" : "=f"(r1) : "r"((x1r_i << 23) + p1_i)); + return make_float2(r0, r1); +} + +__device__ __forceinline__ void ex2_emulation_f32x2(float* x0_ptr, float* x1_ptr) { + float2 result = ex2_emulation_f32x2_value(make_float2(*x0_ptr, *x1_ptr)); + *x0_ptr = result.x; *x1_ptr = result.y; +} + +__device__ __forceinline__ void softmax_frag_exp2_cast( + float* sv, uint32_t* pv, int use_emu) +{ + #pragma unroll + for (int j = 0; j < 16; j++) { + if (use_emu && j >= 12) + ex2_emulation_f32x2(&sv[j*2], &sv[j*2+1]); + else { + sv[j*2] = approx_exp2(sv[j*2]); + sv[j*2+1] = approx_exp2(sv[j*2+1]); + } + } + #pragma unroll + for (int j = 0; j < 16; j++) { + __nv_bfloat162 bf = __float22bfloat162_rn({sv[j*2], sv[j*2+1]}); + pv[j] = reinterpret_cast(bf); + } +} + + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_2d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3}], [%4];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr) { + // Canonical .shared::cta form for non-multicast gather4, matching + // trtllm-gen / cuda_ptx and the PTX ISA qualifier order + // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, + // .shared::cluster is reserved for the multicast variant (ctaMask). + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + + +template +struct CakePipelineState { + uint32_t index_; + uint32_t phase_; + uint32_t count_; + + __device__ __forceinline__ uint32_t index() const { return index_; } + __device__ __forceinline__ uint32_t phase() const { return phase_; } + __device__ __forceinline__ uint32_t count() const { return count_; } + __device__ __forceinline__ CakePipelineState& operator++() { + ++count_; + ++index_; + if (index_ == Stages) { + index_ = 0; + phase_ ^= 1; + } + return *this; + } +}; + +template +__device__ __forceinline__ CakePipelineState +cake_pipeline_state_at(CakePipelineState state) { + #pragma unroll + for (int step = 0; step < Offset; ++step) { + ++state; + } + return state; +} + +extern "C" { + +__global__ __launch_bounds__(512, 1) void +kernel_cake_dsv4_0a925b9bcb3044777c76(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ cum_seq_lens_q, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int mbar_base = smem + 202784; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 16) + #define kv_full_addr (mbar_base + 32) + #define kv_empty_addr (mbar_base + 64) + #define page_full_addr (mbar_base + 96) + #define page_empty_addr (mbar_base + 144) + #define score_full_addr (mbar_base + 192) + #define score_empty_addr (mbar_base + 208) + #define stats_full_addr (mbar_base + 224) + #define stats_empty_addr (mbar_base + 240) + #define o_full_addr (mbar_base + 256) + #define o_empty_addr (mbar_base + 264) + #define work_full_addr (mbar_base + 272) + #define work_empty_addr (mbar_base + 288) + #define throttle_full_addr (mbar_base + 304) + #define throttle_empty_addr (mbar_base + 320) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + + // Kernel setup ops + uint8_t* smem_q = reinterpret_cast(smem_raw + 0); + const int smem_q_addr = smem + 0; + uint8_t* smem_k = reinterpret_cast(smem_raw + 65536); + const int smem_k_addr = smem + 65536; + uint8_t* smem_v = reinterpret_cast(smem_raw + 65536); + const int smem_v_addr = smem + 65536; + int* smem_page_offsets = reinterpret_cast(smem_raw + 196608); + const int smem_page_offsets_addr = smem + 196608; + unsigned int* work_response = reinterpret_cast(smem_raw + 202752); + const int work_response_addr = smem + 202752; + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_q))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_swa_kv))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_compressed_kv))) : "memory"); + + // Mbarrier init (16 pipeline groups, 0 ordered-sequence groups, 42 barriers) + // Mbarriers at smem_raw[202784..203120) + + if (warp == 12) { + uint32_t leader = elect_sync(); + if (leader) { + // --- pipeline 'q_pipe' --- + // q_full: 2 barriers, init_count=1 + mbarrier_init(smem + 202784, 1); + mbarrier_init(smem + 202792, 1); + // q_empty: 2 barriers, init_count=1 + mbarrier_init(smem + 202800, 1); + mbarrier_init(smem + 202808, 1); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // --- pipeline 'kv_pipe' --- + // kv_full: 4 barriers, init_count=4 + mbarrier_init(smem + 202816, 4); + mbarrier_init(smem + 202824, 4); + mbarrier_init(smem + 202832, 4); + mbarrier_init(smem + 202840, 4); + // kv_empty: 4 barriers, init_count=1 + mbarrier_init(smem + 202848, 1); + mbarrier_init(smem + 202856, 1); + mbarrier_init(smem + 202864, 1); + mbarrier_init(smem + 202872, 1); + // --- pipeline 'page_pipe' --- + // page_full: 6 barriers, init_count=32 + mbarrier_init(smem + 202880, 32); + mbarrier_init(smem + 202888, 32); + mbarrier_init(smem + 202896, 32); + mbarrier_init(smem + 202904, 32); + mbarrier_init(smem + 202912, 32); + mbarrier_init(smem + 202920, 32); + // page_empty: 6 barriers, init_count=128 + mbarrier_init(smem + 202928, 128); + mbarrier_init(smem + 202936, 128); + mbarrier_init(smem + 202944, 128); + mbarrier_init(smem + 202952, 128); + mbarrier_init(smem + 202960, 128); + mbarrier_init(smem + 202968, 128); + // --- pipeline 'score_pipe' --- + // score_full: 2 barriers, init_count=1 + mbarrier_init(smem + 202976, 1); + mbarrier_init(smem + 202984, 1); + // score_empty: 2 barriers, init_count=128 + mbarrier_init(smem + 202992, 128); + mbarrier_init(smem + 203000, 128); + // --- pipeline 'stats_pipe' --- + // stats_full: 2 barriers, init_count=128 + mbarrier_init(smem + 203008, 128); + mbarrier_init(smem + 203016, 128); + // stats_empty: 2 barriers, init_count=128 + mbarrier_init(smem + 203024, 128); + mbarrier_init(smem + 203032, 128); + // o_full: 1 barriers, init_count=1 + mbarrier_init(smem + 203040, 1); + // o_empty: 1 barriers, init_count=128 + mbarrier_init(smem + 203048, 128); + // --- pipeline 'work_pipe' --- + // work_full: 2 barriers, init_count=1 + mbarrier_init(smem + 203056, 1); + mbarrier_init(smem + 203064, 1); + // work_empty: 2 barriers, init_count=512 + mbarrier_init(smem + 203072, 512); + mbarrier_init(smem + 203080, 512); + // --- pipeline 'throttle_pipe' --- + // throttle_full: 2 barriers, init_count=128 + mbarrier_init(smem + 203088, 128); + mbarrier_init(smem + 203096, 128); + // throttle_empty: 2 barriers, init_count=32 + mbarrier_init(smem + 203104, 32); + mbarrier_init(smem + 203112, 32); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncwarp(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 203120); + if (warp == 0) { + int _tmem_hold = smem + 203120; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_tmem = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Inc phase consumes the registers released above. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 136;"); + } + + // ---- Role: softmax ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 152;"); + { // softmax_main + const int warp_rank = warp; + const int logical_row_origin = warp_rank * 16; + const int my_row = logical_row_origin + lane % 16; + const int col_half = lane / 16; + int score_cursor = 0; + int stats_cursor = 0; + unsigned int softmax_work_stage = 0; + unsigned int softmax_q = blockIdx.x; + unsigned int softmax_head = blockIdx.y; + unsigned int softmax_batch = blockIdx.z; + unsigned int _phase_work_full = 0; + #pragma unroll 1 + for (unsigned int _work = 0; _work < total_work_items; _work++) { + int q_begin = cum_seq_lens_q[(int)softmax_batch]; + int q_end = cum_seq_lens_q[(int)(softmax_batch + 1)]; + int q_local = (int)softmax_q; + int work_valid = q_local < q_end - q_begin; + asm volatile("griddepcontrol.wait;" ::: "memory"); + float softmax_rescale_log2 = bmm1_scale[0] * 1.4426950408889634f; + asm volatile("griddepcontrol.wait;" ::: "memory"); + float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; + if (work_valid != 0) { + int query_idx = q_begin + q_local; + int _max_4 = ((sparse_topk_lens[query_idx] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx] + sparse_topk_lens_offset) : (0)); + int _min_3 = ((_max_4) < (sparse_topk) ? (_max_4) : (sparse_topk)); + int active_topk = _min_3; + int num_tiles = (active_topk + 128 - 1) / 128; + float row_max = -CAKE_INF; + float row_sum = 0.0f; + #pragma unroll 1 + for (int tile = 0; tile < num_tiles; tile++) { + int pipeline_tile = score_cursor + tile; + int score_stage = pipeline_tile & 1; + int score_phase = pipeline_tile >> 1 & 1; + mbarrier_wait(score_full_addr + (score_stage) * 8, score_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int score_col = ((score_stage != 0) ? 128 : 0); + float _tmem_load_0[64]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_0[0])), "=r"(*reinterpret_cast(&_tmem_load_0[1])), "=r"(*reinterpret_cast(&_tmem_load_0[2])), "=r"(*reinterpret_cast(&_tmem_load_0[3])), "=r"(*reinterpret_cast(&_tmem_load_0[4])), "=r"(*reinterpret_cast(&_tmem_load_0[5])), "=r"(*reinterpret_cast(&_tmem_load_0[6])), "=r"(*reinterpret_cast(&_tmem_load_0[7])), "=r"(*reinterpret_cast(&_tmem_load_0[8])), "=r"(*reinterpret_cast(&_tmem_load_0[9])), "=r"(*reinterpret_cast(&_tmem_load_0[10])), "=r"(*reinterpret_cast(&_tmem_load_0[11])), "=r"(*reinterpret_cast(&_tmem_load_0[12])), "=r"(*reinterpret_cast(&_tmem_load_0[13])), "=r"(*reinterpret_cast(&_tmem_load_0[14])), "=r"(*reinterpret_cast(&_tmem_load_0[15])), "=r"(*reinterpret_cast(&_tmem_load_0[16])), "=r"(*reinterpret_cast(&_tmem_load_0[17])), "=r"(*reinterpret_cast(&_tmem_load_0[18])), "=r"(*reinterpret_cast(&_tmem_load_0[19])), "=r"(*reinterpret_cast(&_tmem_load_0[20])), "=r"(*reinterpret_cast(&_tmem_load_0[21])), "=r"(*reinterpret_cast(&_tmem_load_0[22])), "=r"(*reinterpret_cast(&_tmem_load_0[23])), "=r"(*reinterpret_cast(&_tmem_load_0[24])), "=r"(*reinterpret_cast(&_tmem_load_0[25])), "=r"(*reinterpret_cast(&_tmem_load_0[26])), "=r"(*reinterpret_cast(&_tmem_load_0[27])), "=r"(*reinterpret_cast(&_tmem_load_0[28])), "=r"(*reinterpret_cast(&_tmem_load_0[29])), "=r"(*reinterpret_cast(&_tmem_load_0[30])), "=r"(*reinterpret_cast(&_tmem_load_0[31])) + : "r"(taddr + (unsigned int)score_col)); + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_0[32])), "=r"(*reinterpret_cast(&_tmem_load_0[33])), "=r"(*reinterpret_cast(&_tmem_load_0[34])), "=r"(*reinterpret_cast(&_tmem_load_0[35])), "=r"(*reinterpret_cast(&_tmem_load_0[36])), "=r"(*reinterpret_cast(&_tmem_load_0[37])), "=r"(*reinterpret_cast(&_tmem_load_0[38])), "=r"(*reinterpret_cast(&_tmem_load_0[39])), "=r"(*reinterpret_cast(&_tmem_load_0[40])), "=r"(*reinterpret_cast(&_tmem_load_0[41])), "=r"(*reinterpret_cast(&_tmem_load_0[42])), "=r"(*reinterpret_cast(&_tmem_load_0[43])), "=r"(*reinterpret_cast(&_tmem_load_0[44])), "=r"(*reinterpret_cast(&_tmem_load_0[45])), "=r"(*reinterpret_cast(&_tmem_load_0[46])), "=r"(*reinterpret_cast(&_tmem_load_0[47])), "=r"(*reinterpret_cast(&_tmem_load_0[48])), "=r"(*reinterpret_cast(&_tmem_load_0[49])), "=r"(*reinterpret_cast(&_tmem_load_0[50])), "=r"(*reinterpret_cast(&_tmem_load_0[51])), "=r"(*reinterpret_cast(&_tmem_load_0[52])), "=r"(*reinterpret_cast(&_tmem_load_0[53])), "=r"(*reinterpret_cast(&_tmem_load_0[54])), "=r"(*reinterpret_cast(&_tmem_load_0[55])), "=r"(*reinterpret_cast(&_tmem_load_0[56])), "=r"(*reinterpret_cast(&_tmem_load_0[57])), "=r"(*reinterpret_cast(&_tmem_load_0[58])), "=r"(*reinterpret_cast(&_tmem_load_0[59])), "=r"(*reinterpret_cast(&_tmem_load_0[60])), "=r"(*reinterpret_cast(&_tmem_load_0[61])), "=r"(*reinterpret_cast(&_tmem_load_0[62])), "=r"(*reinterpret_cast(&_tmem_load_0[63])) + : "r"(taddr + (unsigned int)score_col + 32)); + int valid_cols = active_topk - tile * 128 - col_half * 64; + if (valid_cols < 0) { + valid_cols = 0; + } + if (valid_cols > 64) { + valid_cols = 64; + } + uint32_t _slice_lo_mask_0; + { + int _lim_0 = valid_cols; + if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_0 & (1u << 0))) _tmem_load_0[0] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 1))) _tmem_load_0[1] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 2))) _tmem_load_0[2] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 3))) _tmem_load_0[3] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 4))) _tmem_load_0[4] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 5))) _tmem_load_0[5] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 6))) _tmem_load_0[6] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 7))) _tmem_load_0[7] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 8))) _tmem_load_0[8] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 9))) _tmem_load_0[9] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 10))) _tmem_load_0[10] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 11))) _tmem_load_0[11] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 12))) _tmem_load_0[12] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 13))) _tmem_load_0[13] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 14))) _tmem_load_0[14] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 15))) _tmem_load_0[15] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 16))) _tmem_load_0[16] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 17))) _tmem_load_0[17] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 18))) _tmem_load_0[18] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 19))) _tmem_load_0[19] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 20))) _tmem_load_0[20] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 21))) _tmem_load_0[21] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 22))) _tmem_load_0[22] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 23))) _tmem_load_0[23] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 24))) _tmem_load_0[24] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 25))) _tmem_load_0[25] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 26))) _tmem_load_0[26] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 27))) _tmem_load_0[27] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 28))) _tmem_load_0[28] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 29))) _tmem_load_0[29] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 30))) _tmem_load_0[30] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 31))) _tmem_load_0[31] = -CAKE_INF; + uint32_t _slice_lo_mask_1; + { + int _lim_1 = valid_cols - 32; + if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) _tmem_load_0[32] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) _tmem_load_0[33] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) _tmem_load_0[34] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) _tmem_load_0[35] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) _tmem_load_0[36] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) _tmem_load_0[37] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) _tmem_load_0[38] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) _tmem_load_0[39] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) _tmem_load_0[40] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) _tmem_load_0[41] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) _tmem_load_0[42] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) _tmem_load_0[43] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) _tmem_load_0[44] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) _tmem_load_0[45] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) _tmem_load_0[46] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) _tmem_load_0[47] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) _tmem_load_0[48] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) _tmem_load_0[49] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) _tmem_load_0[50] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) _tmem_load_0[51] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) _tmem_load_0[52] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) _tmem_load_0[53] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) _tmem_load_0[54] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) _tmem_load_0[55] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) _tmem_load_0[56] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) _tmem_load_0[57] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) _tmem_load_0[58] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) _tmem_load_0[59] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) _tmem_load_0[60] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) _tmem_load_0[61] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) _tmem_load_0[62] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) _tmem_load_0[63] = -CAKE_INF; + float2 _reg_reduce_max2_2 = {-CAKE_INF, -CAKE_INF}; + row_max_x32_accum(&_tmem_load_0[0], _reg_reduce_max2_2); + row_max_x32_accum(&_tmem_load_0[32], _reg_reduce_max2_2); + float _tmem_load_0_max = row_max_reduce(_reg_reduce_max2_2); + float _max_5 = max_noftz(row_max, _tmem_load_0_max); + float new_max = _max_5; + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, new_max, 16); + float _max_6 = max_noftz(new_max, _shfl_xor_0); + new_max = _max_6; + float _fma_0 = __fmaf_rn(row_max, softmax_rescale_log2, (-new_max) * softmax_rescale_log2); + float delta = _fma_0; + float _exp2_0 = approx_exp2(delta); + float acc_scale = ((row_max > -CAKE_INF) ? _exp2_0 : 1.0f); + int stats_tile = stats_cursor + tile; + int stats_stage = stats_tile & 1; + int stats_phase = stats_tile >> 1 & 1; + mbarrier_wait(stats_empty_addr + (stats_stage) * 8, 1 - stats_phase); + float stats_pair[2]; + stats_pair[0] = row_max; + stats_pair[1] = new_max; + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(taddr + (unsigned int)(stats_stage * 128)), "f"(stats_pair[0]), "f"(stats_pair[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(stats_full_addr + (stats_stage) * 8); + float safe_max = ((new_max == -CAKE_INF) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + const float2 _fma_b2_3 = {softmax_scale_log2, softmax_scale_log2}; + const float2 _fma_c2_4 = {-max_scaled + 8.8073549f, -max_scaled + 8.8073549f}; + #pragma unroll + for (int _lf = 0; _lf < 32; _lf++) + fma_f32x2_inplace(&reinterpret_cast(_tmem_load_0)[_lf], _fma_b2_3, _fma_c2_4); + #pragma unroll + for (int _le = 0; _le < 64; _le++) { + _tmem_load_0[_le] = approx_exp2(_tmem_load_0[_le]); + } + float2 _reg_reduce_sum2_5 = make_float2(0.0f, 0.0f); + softmax_block_sum(&_tmem_load_0[0], &_reg_reduce_sum2_5); + softmax_block_sum(&_tmem_load_0[32], &_reg_reduce_sum2_5); + float _tmem_load_0_sum = _reg_reduce_sum2_5.x + _reg_reduce_sum2_5.y; + float block_sum = _tmem_load_0_sum; + float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, block_sum, 16); + block_sum = block_sum + _shfl_xor_1; + float _fma_1 = __fmaf_rn(row_sum, acc_scale, block_sum); + row_sum = _fma_1; + uint32_t _fp8_0[16]; + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_0[0]), "f"(_tmem_load_0[1]), + "f"(_tmem_load_0[2]), "f"(_tmem_load_0[3])); + _fp8_0[0] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_0[4]), "f"(_tmem_load_0[5]), + "f"(_tmem_load_0[6]), "f"(_tmem_load_0[7])); + _fp8_0[1] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_0[8]), "f"(_tmem_load_0[9]), + "f"(_tmem_load_0[10]), "f"(_tmem_load_0[11])); + _fp8_0[2] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_0[12]), "f"(_tmem_load_0[13]), + "f"(_tmem_load_0[14]), "f"(_tmem_load_0[15])); + _fp8_0[3] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_0[16]), "f"(_tmem_load_0[17]), + "f"(_tmem_load_0[18]), "f"(_tmem_load_0[19])); + _fp8_0[4] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_0[20]), "f"(_tmem_load_0[21]), + "f"(_tmem_load_0[22]), "f"(_tmem_load_0[23])); + _fp8_0[5] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_0[24]), "f"(_tmem_load_0[25]), + "f"(_tmem_load_0[26]), "f"(_tmem_load_0[27])); + _fp8_0[6] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_0[28]), "f"(_tmem_load_0[29]), + "f"(_tmem_load_0[30]), "f"(_tmem_load_0[31])); + _fp8_0[7] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_0[32]), "f"(_tmem_load_0[33]), + "f"(_tmem_load_0[34]), "f"(_tmem_load_0[35])); + _fp8_0[8] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_0[36]), "f"(_tmem_load_0[37]), + "f"(_tmem_load_0[38]), "f"(_tmem_load_0[39])); + _fp8_0[9] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_0[40]), "f"(_tmem_load_0[41]), + "f"(_tmem_load_0[42]), "f"(_tmem_load_0[43])); + _fp8_0[10] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_0[44]), "f"(_tmem_load_0[45]), + "f"(_tmem_load_0[46]), "f"(_tmem_load_0[47])); + _fp8_0[11] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_0[48]), "f"(_tmem_load_0[49]), + "f"(_tmem_load_0[50]), "f"(_tmem_load_0[51])); + _fp8_0[12] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_0[52]), "f"(_tmem_load_0[53]), + "f"(_tmem_load_0[54]), "f"(_tmem_load_0[55])); + _fp8_0[13] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_0[56]), "f"(_tmem_load_0[57]), + "f"(_tmem_load_0[58]), "f"(_tmem_load_0[59])); + _fp8_0[14] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_0[60]), "f"(_tmem_load_0[61]), + "f"(_tmem_load_0[62]), "f"(_tmem_load_0[63])); + _fp8_0[15] = _packed; + } + int p_col = ((score_stage != 0) ? 160 : 32); + asm volatile( + "tcgen05.st.sync.aligned.16x32bx2.x16.b32" + " [%0], 16, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(taddr + (unsigned int)p_col), "r"(*reinterpret_cast(&_fp8_0[0])), "r"(*reinterpret_cast(&_fp8_0[1])), "r"(*reinterpret_cast(&_fp8_0[2])), "r"(*reinterpret_cast(&_fp8_0[3])), "r"(*reinterpret_cast(&_fp8_0[4])), "r"(*reinterpret_cast(&_fp8_0[5])), "r"(*reinterpret_cast(&_fp8_0[6])), "r"(*reinterpret_cast(&_fp8_0[7])), "r"(*reinterpret_cast(&_fp8_0[8])), "r"(*reinterpret_cast(&_fp8_0[9])), "r"(*reinterpret_cast(&_fp8_0[10])), "r"(*reinterpret_cast(&_fp8_0[11])), "r"(*reinterpret_cast(&_fp8_0[12])), "r"(*reinterpret_cast(&_fp8_0[13])), "r"(*reinterpret_cast(&_fp8_0[14])), "r"(*reinterpret_cast(&_fp8_0[15]))); + asm volatile( + "tcgen05.st.sync.aligned.16x32bx2.x16.b32" + " [%0], 16, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(taddr + 1048576 + (unsigned int)p_col), "r"(*reinterpret_cast(&_fp8_0[0])), "r"(*reinterpret_cast(&_fp8_0[1])), "r"(*reinterpret_cast(&_fp8_0[2])), "r"(*reinterpret_cast(&_fp8_0[3])), "r"(*reinterpret_cast(&_fp8_0[4])), "r"(*reinterpret_cast(&_fp8_0[5])), "r"(*reinterpret_cast(&_fp8_0[6])), "r"(*reinterpret_cast(&_fp8_0[7])), "r"(*reinterpret_cast(&_fp8_0[8])), "r"(*reinterpret_cast(&_fp8_0[9])), "r"(*reinterpret_cast(&_fp8_0[10])), "r"(*reinterpret_cast(&_fp8_0[11])), "r"(*reinterpret_cast(&_fp8_0[12])), "r"(*reinterpret_cast(&_fp8_0[13])), "r"(*reinterpret_cast(&_fp8_0[14])), "r"(*reinterpret_cast(&_fp8_0[15]))); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + if (num_tiles > tile + 1) { + mbarrier_arrive(score_empty_addr + (score_stage) * 8); + } + row_max = new_max; + } + int final_stats_tile = stats_cursor + num_tiles; + int final_stats_stage = final_stats_tile & 1; + int final_stats_phase = final_stats_tile >> 1 & 1; + mbarrier_wait(stats_empty_addr + (final_stats_stage) * 8, 1 - final_stats_phase); + float final_stats[2]; + final_stats[0] = row_sum; + final_stats[1] = row_max; + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(taddr + (unsigned int)(final_stats_stage * 128)), "f"(final_stats[0]), "f"(final_stats[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(stats_full_addr + (final_stats_stage) * 8); + int tail_stats_tile = final_stats_tile + 1; + int tail_stats_stage = tail_stats_tile & 1; + int tail_stats_phase = tail_stats_tile >> 1 & 1; + mbarrier_wait(stats_empty_addr + (tail_stats_stage) * 8, 1 - tail_stats_phase); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(stats_full_addr + (tail_stats_stage) * 8); + int last_score_tile = score_cursor + num_tiles - 1; + mbarrier_arrive(score_empty_addr + (last_score_tile & 1) * 8); + stats_cursor = stats_cursor + num_tiles + 2; + score_cursor = score_cursor + num_tiles; + } + mbarrier_wait(work_full_addr + (softmax_work_stage) * 8, _phase_work_full); + uint32_t _clc_valid_3 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_3) + : "r"(work_response_addr + softmax_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_9 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_9) + : "r"(work_response_addr + softmax_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_10 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_10) + : "r"(work_response_addr + softmax_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_11 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_11) + : "r"(work_response_addr + softmax_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(work_empty_addr + (softmax_work_stage) * 8); + softmax_work_stage += 1; + if (softmax_work_stage == 2) { softmax_work_stage = 0; _phase_work_full ^= 1; } + if (_clc_valid_3 == 0) { + break; + } + softmax_q = _clc_ctaid_9; + softmax_head = _clc_ctaid_10; + softmax_batch = _clc_ctaid_11; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + } + } + // ---- Role: correction ---- + if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 144;"); + { // correction_main + const int warp_rank_1 = warp - 4; + const int logical_row_origin_1 = warp_rank_1 * 16; + const int my_row_1 = logical_row_origin_1 + lane % 16; + const int col_half_1 = lane / 16; + int score_cursor_1 = 0; + int stats_cursor_1 = 0; + unsigned int correction_work_stage = 0; + unsigned int correction_q = blockIdx.x; + unsigned int correction_head = blockIdx.y; + unsigned int correction_batch = blockIdx.z; + unsigned int _phase_work_full_1 = 0; + #pragma unroll 1 + for (unsigned int _work_1 = 0; _work_1 < total_work_items; _work_1++) { + int q_begin_1 = cum_seq_lens_q[(int)correction_batch]; + int q_end_1 = cum_seq_lens_q[(int)(correction_batch + 1)]; + int q_local_1 = (int)correction_q; + int work_valid_1 = q_local_1 < q_end_1 - q_begin_1; + asm volatile("griddepcontrol.wait;" ::: "memory"); + float softmax_scale_log2_1 = bmm1_scale[0] * 1.4426950408889634f; + float output_scale = bmm2_scale[0]; + if (work_valid_1 != 0) { + int query_idx_1 = q_begin_1 + q_local_1; + int head_base = (int)correction_head * 64; + int head_idx = head_base + my_row_1; + int _max_7 = ((sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) : (0)); + int _min_4 = ((_max_7) < (sparse_topk) ? (_max_7) : (sparse_topk)); + int active_topk_1 = _min_4; + int num_tiles_1 = (active_topk_1 + 128 - 1) / 128; + int first_stats_tile = stats_cursor_1; + int first_stats_stage = first_stats_tile & 1; + int first_stats_phase = first_stats_tile >> 1 & 1; + mbarrier_wait(stats_full_addr + (first_stats_stage) * 8, first_stats_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + mbarrier_arrive(stats_empty_addr + (first_stats_stage) * 8); + #pragma unroll 1 + for (int tile_1 = 1; tile_1 < num_tiles_1; tile_1++) { + int stats_tile_1 = stats_cursor_1 + tile_1; + int stats_stage_1 = stats_tile_1 & 1; + int stats_phase_1 = stats_tile_1 >> 1 & 1; + mbarrier_wait(stats_full_addr + (stats_stage_1) * 8, stats_phase_1); + asm volatile("tcgen05.fence::after_thread_sync;"); + float _tmem_load_1[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]) + : "r"(taddr + (unsigned int)(stats_stage_1 * 128))); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + mbarrier_arrive(stats_empty_addr + (stats_stage_1) * 8); + int output_tile = score_cursor_1 + tile_1 - 1; + mbarrier_wait(o_full_addr, output_tile & 1); + asm volatile("tcgen05.fence::after_thread_sync;"); + float max_delta = _tmem_load_1[0] - _tmem_load_1[1]; + float _exp2_1 = approx_exp2(max_delta * softmax_scale_log2_1); + float acc_scale_1 = ((max_delta != 0.0f) ? _exp2_1 : 1.0f); + int _vote_0 = __any_sync(0xFFFFFFFF, acc_scale_1 < 1.0f); + int any_rescale = _vote_0; + if (any_rescale != 0) { + #pragma unroll + for (int output_half = 0; output_half < 2; output_half++) { + int row_bank = ((output_half != 0) ? 1048576 : 0); + #pragma unroll + for (int col = 0; col < 128; col += 64) { + float _tmem_load_2[64]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 128;" + : "=r"(*reinterpret_cast(&_tmem_load_2[0])), "=r"(*reinterpret_cast(&_tmem_load_2[1])), "=r"(*reinterpret_cast(&_tmem_load_2[2])), "=r"(*reinterpret_cast(&_tmem_load_2[3])), "=r"(*reinterpret_cast(&_tmem_load_2[4])), "=r"(*reinterpret_cast(&_tmem_load_2[5])), "=r"(*reinterpret_cast(&_tmem_load_2[6])), "=r"(*reinterpret_cast(&_tmem_load_2[7])), "=r"(*reinterpret_cast(&_tmem_load_2[8])), "=r"(*reinterpret_cast(&_tmem_load_2[9])), "=r"(*reinterpret_cast(&_tmem_load_2[10])), "=r"(*reinterpret_cast(&_tmem_load_2[11])), "=r"(*reinterpret_cast(&_tmem_load_2[12])), "=r"(*reinterpret_cast(&_tmem_load_2[13])), "=r"(*reinterpret_cast(&_tmem_load_2[14])), "=r"(*reinterpret_cast(&_tmem_load_2[15])), "=r"(*reinterpret_cast(&_tmem_load_2[16])), "=r"(*reinterpret_cast(&_tmem_load_2[17])), "=r"(*reinterpret_cast(&_tmem_load_2[18])), "=r"(*reinterpret_cast(&_tmem_load_2[19])), "=r"(*reinterpret_cast(&_tmem_load_2[20])), "=r"(*reinterpret_cast(&_tmem_load_2[21])), "=r"(*reinterpret_cast(&_tmem_load_2[22])), "=r"(*reinterpret_cast(&_tmem_load_2[23])), "=r"(*reinterpret_cast(&_tmem_load_2[24])), "=r"(*reinterpret_cast(&_tmem_load_2[25])), "=r"(*reinterpret_cast(&_tmem_load_2[26])), "=r"(*reinterpret_cast(&_tmem_load_2[27])), "=r"(*reinterpret_cast(&_tmem_load_2[28])), "=r"(*reinterpret_cast(&_tmem_load_2[29])), "=r"(*reinterpret_cast(&_tmem_load_2[30])), "=r"(*reinterpret_cast(&_tmem_load_2[31])) + : "r"(taddr + 256 + (unsigned int)row_bank + (unsigned int)col)); + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 128;" + : "=r"(*reinterpret_cast(&_tmem_load_2[32])), "=r"(*reinterpret_cast(&_tmem_load_2[33])), "=r"(*reinterpret_cast(&_tmem_load_2[34])), "=r"(*reinterpret_cast(&_tmem_load_2[35])), "=r"(*reinterpret_cast(&_tmem_load_2[36])), "=r"(*reinterpret_cast(&_tmem_load_2[37])), "=r"(*reinterpret_cast(&_tmem_load_2[38])), "=r"(*reinterpret_cast(&_tmem_load_2[39])), "=r"(*reinterpret_cast(&_tmem_load_2[40])), "=r"(*reinterpret_cast(&_tmem_load_2[41])), "=r"(*reinterpret_cast(&_tmem_load_2[42])), "=r"(*reinterpret_cast(&_tmem_load_2[43])), "=r"(*reinterpret_cast(&_tmem_load_2[44])), "=r"(*reinterpret_cast(&_tmem_load_2[45])), "=r"(*reinterpret_cast(&_tmem_load_2[46])), "=r"(*reinterpret_cast(&_tmem_load_2[47])), "=r"(*reinterpret_cast(&_tmem_load_2[48])), "=r"(*reinterpret_cast(&_tmem_load_2[49])), "=r"(*reinterpret_cast(&_tmem_load_2[50])), "=r"(*reinterpret_cast(&_tmem_load_2[51])), "=r"(*reinterpret_cast(&_tmem_load_2[52])), "=r"(*reinterpret_cast(&_tmem_load_2[53])), "=r"(*reinterpret_cast(&_tmem_load_2[54])), "=r"(*reinterpret_cast(&_tmem_load_2[55])), "=r"(*reinterpret_cast(&_tmem_load_2[56])), "=r"(*reinterpret_cast(&_tmem_load_2[57])), "=r"(*reinterpret_cast(&_tmem_load_2[58])), "=r"(*reinterpret_cast(&_tmem_load_2[59])), "=r"(*reinterpret_cast(&_tmem_load_2[60])), "=r"(*reinterpret_cast(&_tmem_load_2[61])), "=r"(*reinterpret_cast(&_tmem_load_2[62])), "=r"(*reinterpret_cast(&_tmem_load_2[63])) + : "r"(taddr + 256 + (unsigned int)row_bank + (unsigned int)col + 32)); + const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 32; _ls++) + mul_f32x2_inplace(&reinterpret_cast(_tmem_load_2)[_ls], _scale2_0); + asm volatile( + "tcgen05.st.sync.aligned.16x32bx2.x64.b32" + " [%0], 128, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64};" + :: "r"(taddr + 256 + (unsigned int)row_bank + (unsigned int)col), "r"(*reinterpret_cast(&_tmem_load_2[0])), "r"(*reinterpret_cast(&_tmem_load_2[1])), "r"(*reinterpret_cast(&_tmem_load_2[2])), "r"(*reinterpret_cast(&_tmem_load_2[3])), "r"(*reinterpret_cast(&_tmem_load_2[4])), "r"(*reinterpret_cast(&_tmem_load_2[5])), "r"(*reinterpret_cast(&_tmem_load_2[6])), "r"(*reinterpret_cast(&_tmem_load_2[7])), "r"(*reinterpret_cast(&_tmem_load_2[8])), "r"(*reinterpret_cast(&_tmem_load_2[9])), "r"(*reinterpret_cast(&_tmem_load_2[10])), "r"(*reinterpret_cast(&_tmem_load_2[11])), "r"(*reinterpret_cast(&_tmem_load_2[12])), "r"(*reinterpret_cast(&_tmem_load_2[13])), "r"(*reinterpret_cast(&_tmem_load_2[14])), "r"(*reinterpret_cast(&_tmem_load_2[15])), "r"(*reinterpret_cast(&_tmem_load_2[16])), "r"(*reinterpret_cast(&_tmem_load_2[17])), "r"(*reinterpret_cast(&_tmem_load_2[18])), "r"(*reinterpret_cast(&_tmem_load_2[19])), "r"(*reinterpret_cast(&_tmem_load_2[20])), "r"(*reinterpret_cast(&_tmem_load_2[21])), "r"(*reinterpret_cast(&_tmem_load_2[22])), "r"(*reinterpret_cast(&_tmem_load_2[23])), "r"(*reinterpret_cast(&_tmem_load_2[24])), "r"(*reinterpret_cast(&_tmem_load_2[25])), "r"(*reinterpret_cast(&_tmem_load_2[26])), "r"(*reinterpret_cast(&_tmem_load_2[27])), "r"(*reinterpret_cast(&_tmem_load_2[28])), "r"(*reinterpret_cast(&_tmem_load_2[29])), "r"(*reinterpret_cast(&_tmem_load_2[30])), "r"(*reinterpret_cast(&_tmem_load_2[31])), "r"(*reinterpret_cast(&_tmem_load_2[32])), "r"(*reinterpret_cast(&_tmem_load_2[33])), "r"(*reinterpret_cast(&_tmem_load_2[34])), "r"(*reinterpret_cast(&_tmem_load_2[35])), "r"(*reinterpret_cast(&_tmem_load_2[36])), "r"(*reinterpret_cast(&_tmem_load_2[37])), "r"(*reinterpret_cast(&_tmem_load_2[38])), "r"(*reinterpret_cast(&_tmem_load_2[39])), "r"(*reinterpret_cast(&_tmem_load_2[40])), "r"(*reinterpret_cast(&_tmem_load_2[41])), "r"(*reinterpret_cast(&_tmem_load_2[42])), "r"(*reinterpret_cast(&_tmem_load_2[43])), "r"(*reinterpret_cast(&_tmem_load_2[44])), "r"(*reinterpret_cast(&_tmem_load_2[45])), "r"(*reinterpret_cast(&_tmem_load_2[46])), "r"(*reinterpret_cast(&_tmem_load_2[47])), "r"(*reinterpret_cast(&_tmem_load_2[48])), "r"(*reinterpret_cast(&_tmem_load_2[49])), "r"(*reinterpret_cast(&_tmem_load_2[50])), "r"(*reinterpret_cast(&_tmem_load_2[51])), "r"(*reinterpret_cast(&_tmem_load_2[52])), "r"(*reinterpret_cast(&_tmem_load_2[53])), "r"(*reinterpret_cast(&_tmem_load_2[54])), "r"(*reinterpret_cast(&_tmem_load_2[55])), "r"(*reinterpret_cast(&_tmem_load_2[56])), "r"(*reinterpret_cast(&_tmem_load_2[57])), "r"(*reinterpret_cast(&_tmem_load_2[58])), "r"(*reinterpret_cast(&_tmem_load_2[59])), "r"(*reinterpret_cast(&_tmem_load_2[60])), "r"(*reinterpret_cast(&_tmem_load_2[61])), "r"(*reinterpret_cast(&_tmem_load_2[62])), "r"(*reinterpret_cast(&_tmem_load_2[63]))); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + } + } + } + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(o_empty_addr); + } + int final_stats_tile_1 = stats_cursor_1 + num_tiles_1; + int final_stats_stage_1 = final_stats_tile_1 & 1; + int final_stats_phase_1 = final_stats_tile_1 >> 1 & 1; + mbarrier_wait(stats_full_addr + (final_stats_stage_1) * 8, final_stats_phase_1); + asm volatile("tcgen05.fence::after_thread_sync;"); + float _tmem_load_3[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]) + : "r"(taddr + (unsigned int)(final_stats_stage_1 * 128))); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + mbarrier_arrive(stats_empty_addr + (final_stats_stage_1) * 8); + int tail_stats_tile_1 = final_stats_tile_1 + 1; + int tail_stats_stage_1 = tail_stats_tile_1 & 1; + int tail_stats_phase_1 = tail_stats_tile_1 >> 1 & 1; + mbarrier_wait(stats_full_addr + (tail_stats_stage_1) * 8, tail_stats_phase_1); + asm volatile("tcgen05.fence::after_thread_sync;"); + mbarrier_arrive(stats_empty_addr + (tail_stats_stage_1) * 8); + int last_output_tile = score_cursor_1 + num_tiles_1 - 1; + mbarrier_wait(o_full_addr, last_output_tile & 1); + asm volatile("tcgen05.fence::after_thread_sync;"); + float total_sum = _tmem_load_3[0]; + float final_max = _tmem_load_3[1]; + int output_in_bounds = head_idx < num_heads; + if (((int)(has_sinks != 0) & output_in_bounds) != 0) { + float _exp2_2 = approx_exp2(sinks[head_idx] * 1.4426950408889634f - final_max * softmax_scale_log2_1); + total_sum = total_sum + _exp2_2 * 448.0f; + } + float _rcp_0 = approx_rcp(total_sum); + float inv_sum = ((total_sum > 0.0f) ? output_scale * _rcp_0 : 0.0f); + int output_base = (query_idx_1 * num_heads + head_idx) * 512; + #pragma unroll + for (int output_half_1 = 0; output_half_1 < 2; output_half_1++) { + int row_bank_1 = ((output_half_1 != 0) ? 1048576 : 0); + int col_base = col_half_1 * 128; + #pragma unroll 2 + for (int wcol = 0; wcol < 128; wcol += 16) { + float _tmem_load_4[16]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16], 128;" + : "=r"(*reinterpret_cast(&_tmem_load_4[0])), "=r"(*reinterpret_cast(&_tmem_load_4[1])), "=r"(*reinterpret_cast(&_tmem_load_4[2])), "=r"(*reinterpret_cast(&_tmem_load_4[3])), "=r"(*reinterpret_cast(&_tmem_load_4[4])), "=r"(*reinterpret_cast(&_tmem_load_4[5])), "=r"(*reinterpret_cast(&_tmem_load_4[6])), "=r"(*reinterpret_cast(&_tmem_load_4[7])), "=r"(*reinterpret_cast(&_tmem_load_4[8])), "=r"(*reinterpret_cast(&_tmem_load_4[9])), "=r"(*reinterpret_cast(&_tmem_load_4[10])), "=r"(*reinterpret_cast(&_tmem_load_4[11])), "=r"(*reinterpret_cast(&_tmem_load_4[12])), "=r"(*reinterpret_cast(&_tmem_load_4[13])), "=r"(*reinterpret_cast(&_tmem_load_4[14])), "=r"(*reinterpret_cast(&_tmem_load_4[15])) + : "r"(taddr + 256 + (unsigned int)row_bank_1 + (unsigned int)wcol)); + const float2 _scale2_1 = {inv_sum, inv_sum}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + mul_f32x2_inplace(&reinterpret_cast(_tmem_load_4)[_ls], _scale2_1); + uint32_t _tmem_load_4_bf16[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_4[_lp*2 + 0], _tmem_load_4[_lp*2+1 + 0])); + _tmem_load_4_bf16[_lp] = *(uint32_t*)&_bf2; + } + if (output_in_bounds != 0) { + { + const unsigned* _raw_stv8_2 = reinterpret_cast(_tmem_load_4_bf16); + asm volatile( + "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"((void*)(O + (output_base + output_half_1 * 256 + col_base + wcol))), "r"(_raw_stv8_2[0]), "r"(_raw_stv8_2[1]), "r"(_raw_stv8_2[2]), "r"(_raw_stv8_2[3]), "r"(_raw_stv8_2[4]), "r"(_raw_stv8_2[5]), "r"(_raw_stv8_2[6]), "r"(_raw_stv8_2[7]) : "memory"); + } + } + } + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(o_empty_addr); + stats_cursor_1 = stats_cursor_1 + num_tiles_1 + 2; + score_cursor_1 = score_cursor_1 + num_tiles_1; + } + mbarrier_wait(work_full_addr + (correction_work_stage) * 8, _phase_work_full_1); + uint32_t _clc_valid_4 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_4) + : "r"(work_response_addr + correction_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_12 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_12) + : "r"(work_response_addr + correction_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_13 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_13) + : "r"(work_response_addr + correction_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_14 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_14) + : "r"(work_response_addr + correction_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(work_empty_addr + (correction_work_stage) * 8); + correction_work_stage += 1; + if (correction_work_stage == 2) { correction_work_stage = 0; _phase_work_full_1 ^= 1; } + if (_clc_valid_4 == 0) { + break; + } + correction_q = _clc_ctaid_12; + correction_head = _clc_ctaid_13; + correction_batch = _clc_ctaid_14; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("barrier.sync 7, 128;" ::: "memory"); + if (warp == 4) { + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + } + } + // ---- Role: mma ---- + if (warp == 8) { + { // mma_main + CakePipelineState<4> _kv_cons_state_state{0u, 0u, 0u}; + unsigned int q_cons_stage = 0; + int score_prod_stage = 0; + int score_prod_phase = 0; + int p_cons_stage = 0; + int p_cons_phase = 0; + int o_prod_phase = 0; + unsigned int mma_work_stage = 0; + unsigned int mma_q = blockIdx.x; + unsigned int mma_head = blockIdx.y; + unsigned int mma_batch = blockIdx.z; + #pragma unroll + for (int score_lookahead = 0; score_lookahead < 2; score_lookahead++) { + int lookahead_stage = score_lookahead & 1; + int lookahead_phase = score_lookahead >> 1 & 1; + mbarrier_wait(score_empty_addr + (lookahead_stage) * 8, 1 - lookahead_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + } + unsigned int _phase_q_full = 0; + unsigned int _phase_work_full_2 = 0; + #pragma unroll 1 + for (unsigned int _work_2 = 0; _work_2 < total_work_items; _work_2++) { + int q_begin_2 = cum_seq_lens_q[(int)mma_batch]; + int q_end_2 = cum_seq_lens_q[(int)(mma_batch + 1)]; + int q_local_2 = (int)mma_q; + int work_valid_2 = q_local_2 < q_end_2 - q_begin_2; + if (work_valid_2 != 0) { + int query_idx_2 = q_begin_2 + q_local_2; + int _max_3 = ((sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) : (0)); + int _min_2 = ((_max_3) < (sparse_topk) ? (_max_3) : (sparse_topk)); + int active_topk_2 = _min_2; + int num_tiles_2 = (active_topk_2 + 128 - 1) / 128; + if (num_tiles_2 > 0) { + mbarrier_wait(q_full_addr + (q_cons_stage) * 8, _phase_q_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + int score_col_1 = ((score_prod_stage != 0) ? 128 : 0); + mbarrier_wait(kv_full_addr + (_kv_cons_state_state.index()) * 8, _kv_cons_state_state.phase()); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_0 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (q_cons_stage) * 2048); + int _mma_b_lo_0 = make_warp_uniform((((smem_k_addr) >> 4) & 0x3FFF) + (_kv_cons_state_state.index()) * 2048); + { + uint64_t _mma_ss_a_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_0); + uint64_t _mma_ss_b_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_0); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69206032, 0); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 506U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 1018U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69206032, 1); + } + } + elect_commit(kv_empty_addr + (_kv_cons_state_state.index()) * 8); + ++_kv_cons_state_state; + mbarrier_wait(kv_full_addr + (_kv_cons_state_state.index()) * 8, _kv_cons_state_state.phase()); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_1 = make_warp_uniform((((smem_q_addr + 16384) >> 4) & 0x3FFF) + (q_cons_stage) * 2048); + int _mma_b_lo_1 = make_warp_uniform((((smem_k_addr) >> 4) & 0x3FFF) + (_kv_cons_state_state.index()) * 2048); + { + uint64_t _mma_ss_a_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_1); + uint64_t _mma_ss_b_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_1); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 506U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 1018U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69206032, 1); + } + } + elect_commit(score_full_addr + (score_prod_stage) * 8); + elect_commit(kv_empty_addr + (_kv_cons_state_state.index()) * 8); + ++_kv_cons_state_state; + score_prod_stage ^= 1; + if (score_prod_stage == 0) { + score_prod_phase ^= 1; + } + int first_pv = 1; + #pragma unroll 1 + for (int _steady = 0; _steady < num_tiles_2 - 1; _steady++) { + score_col_1 = ((score_prod_stage != 0) ? 128 : 0); + mbarrier_wait(kv_full_addr + (_kv_cons_state_state.index()) * 8, _kv_cons_state_state.phase()); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_2 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (q_cons_stage) * 2048); + int _mma_b_lo_2 = make_warp_uniform((((smem_k_addr) >> 4) & 0x3FFF) + (_kv_cons_state_state.index()) * 2048); + { + uint64_t _mma_ss_a_desc_2 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_2); + uint64_t _mma_ss_b_desc_2 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_2); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 69206032, 0); + } + incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_2, 506U); + incr_smem_desc_lo(_mma_ss_b_desc_2, 1018U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 69206032, 1); + } + } + elect_commit(kv_empty_addr + (_kv_cons_state_state.index()) * 8); + ++_kv_cons_state_state; + mbarrier_wait(kv_full_addr + (_kv_cons_state_state.index()) * 8, _kv_cons_state_state.phase()); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_3 = make_warp_uniform((((smem_q_addr + 16384) >> 4) & 0x3FFF) + (q_cons_stage) * 2048); + int _mma_b_lo_3 = make_warp_uniform((((smem_k_addr) >> 4) & 0x3FFF) + (_kv_cons_state_state.index()) * 2048); + { + uint64_t _mma_ss_a_desc_3 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_3); + uint64_t _mma_ss_b_desc_3 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_3); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_3, 506U); + incr_smem_desc_lo(_mma_ss_b_desc_3, 1018U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 69206032, 1); + } + } + elect_commit(score_full_addr + (score_prod_stage) * 8); + elect_commit(kv_empty_addr + (_kv_cons_state_state.index()) * 8); + ++_kv_cons_state_state; + score_prod_stage ^= 1; + if (score_prod_stage == 0) { + score_prod_phase ^= 1; + } + mbarrier_wait(score_empty_addr + (score_prod_stage) * 8, 1 - score_prod_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + mbarrier_wait(o_empty_addr, 1 - o_prod_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int v_half = 0; v_half < 2; v_half++) { + mbarrier_wait(kv_full_addr + (_kv_cons_state_state.index()) * 8, _kv_cons_state_state.phase()); + asm volatile("tcgen05.fence::after_thread_sync;"); + int row_bank_2 = ((v_half != 0) ? 1048576 : 0); + int p_col_1 = ((p_cons_stage != 0) ? 160 : 32) + row_bank_2; + int _mma_b_lo_4 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (_kv_cons_state_state.index()) * 2048); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 dhi, blo, ta, id;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 71368720;\n\t" + "mov.b32 ta, %2;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p0;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" + "}\n" + :: "r"((tmem_tmem + (256 + row_bank_2))), "r"(_mma_b_lo_4), "r"(tmem_tmem + p_col_1), "r"(((first_pv) ? 0 : 1))); + elect_commit(kv_empty_addr + (_kv_cons_state_state.index()) * 8); + ++_kv_cons_state_state; + } + first_pv = 0; + elect_commit(o_full_addr); + p_cons_stage ^= 1; + if (p_cons_stage == 0) { + p_cons_phase ^= 1; + } + o_prod_phase ^= 1; + } + elect_commit(q_empty_addr + (q_cons_stage) * 8); + q_cons_stage += 1; + if (q_cons_stage == 2) { q_cons_stage = 0; _phase_q_full ^= 1; } + int tail_score_stage = score_prod_stage ^ 1; + int tail_score_phase = score_prod_phase ^ score_prod_stage; + mbarrier_wait(score_empty_addr + (tail_score_stage) * 8, 1 - tail_score_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + mbarrier_wait(o_empty_addr, 1 - o_prod_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int v_half_1 = 0; v_half_1 < 2; v_half_1++) { + mbarrier_wait(kv_full_addr + (_kv_cons_state_state.index()) * 8, _kv_cons_state_state.phase()); + asm volatile("tcgen05.fence::after_thread_sync;"); + int row_bank_3 = ((v_half_1 != 0) ? 1048576 : 0); + int p_col_2 = ((p_cons_stage != 0) ? 160 : 32) + row_bank_3; + int _mma_b_lo_5 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (_kv_cons_state_state.index()) * 2048); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 dhi, blo, ta, id;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 71368720;\n\t" + "mov.b32 ta, %2;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p0;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" + "}\n" + :: "r"((tmem_tmem + (256 + row_bank_3))), "r"(_mma_b_lo_5), "r"(tmem_tmem + p_col_2), "r"(((first_pv) ? 0 : 1))); + elect_commit(kv_empty_addr + (_kv_cons_state_state.index()) * 8); + ++_kv_cons_state_state; + } + elect_commit(o_full_addr); + p_cons_stage ^= 1; + if (p_cons_stage == 0) { + p_cons_phase ^= 1; + } + o_prod_phase ^= 1; + } + } + mbarrier_wait(work_full_addr + (mma_work_stage) * 8, _phase_work_full_2); + uint32_t _clc_valid_2 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_2) + : "r"(work_response_addr + mma_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_6 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_6) + : "r"(work_response_addr + mma_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_7 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_7) + : "r"(work_response_addr + mma_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_8 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_8) + : "r"(work_response_addr + mma_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(work_empty_addr + (mma_work_stage) * 8); + mma_work_stage += 1; + if (mma_work_stage == 2) { mma_work_stage = 0; _phase_work_full_2 ^= 1; } + if (_clc_valid_2 == 0) { + break; + } + mma_q = _clc_ctaid_6; + mma_head = _clc_ctaid_7; + mma_batch = _clc_ctaid_8; + } + } + } + // ---- Role: page_loader ---- + if (warp == 9) { + { // page_loader_main + unsigned int page_prod_stage = 0; + unsigned int page_prod_phase = 1; + unsigned int page_work_stage = 0; + unsigned int page_q = blockIdx.x; + unsigned int page_head = blockIdx.y; + unsigned int page_batch = blockIdx.z; + unsigned int _phase_work_full_3 = 0; + #pragma unroll 1 + for (unsigned int _work_3 = 0; _work_3 < total_work_items; _work_3++) { + int q_begin_3 = cum_seq_lens_q[(int)page_batch]; + int q_end_3 = cum_seq_lens_q[(int)(page_batch + 1)]; + int q_local_3 = (int)page_q; + int work_valid_3 = q_local_3 < q_end_3 - q_begin_3; + if (work_valid_3 != 0) { + int query_idx_3 = q_begin_3 + q_local_3; + int _max_0 = ((sparse_topk_lens[query_idx_3] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_3] + sparse_topk_lens_offset) : (0)); + int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); + int active_topk_3 = _min_0; + int num_tiles_3 = (active_topk_3 + 128 - 1) / 128; + int num_page_passes = (num_tiles_3 + 1) / 2; + int* swa_row = swa_indices + (query_idx_3 * swa_index_stride); + int* compressed_row = compressed_indices + (query_idx_3 * compressed_index_stride); + int _max_1 = (((active_topk_3 - 1) / 4 * 4) > (0) ? ((active_topk_3 - 1) / 4 * 4) : (0)); + int last_vec_base = _max_1; + #pragma unroll 1 + for (int page_pass = 0; page_pass < num_page_passes; page_pass++) { + #pragma unroll + for (int _duplicate = 0; _duplicate < 2; _duplicate++) { + mbarrier_wait(page_empty_addr + (page_prod_stage) * 8, page_prod_phase); + int page_dst = smem_page_offsets_addr + page_prod_stage * 1024; + #pragma unroll + for (int page_half = 0; page_half < 2; page_half++) { + int page_offset = page_pass * 256 + page_half * 128 + lane * 4; + int clamped_offset = ((page_offset < active_topk_3) ? page_offset : last_vec_base); + int* page_src = ((clamped_offset < 128) ? swa_row + clamped_offset : compressed_row + (clamped_offset - 128)); + asm volatile("cp.async.cg.shared::cta.global.L2::128B [%0], [%1], 16;" + :: "r"(page_dst + (page_half * 128 + lane * 4) * 4), "l"(page_src)); + } + asm volatile("cp.async.commit_group;"); + asm volatile("cp.async.wait_group 0;"); + mbarrier_arrive(page_full_addr + (page_prod_stage) * 8); + page_prod_stage += 1; + if (page_prod_stage == 6) { page_prod_stage = 0; page_prod_phase ^= 1; } + } + } + } + mbarrier_wait(work_full_addr + (page_work_stage) * 8, _phase_work_full_3); + uint32_t _clc_valid_0 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_0) + : "r"(work_response_addr + page_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_0 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_0) + : "r"(work_response_addr + page_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_1 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_1) + : "r"(work_response_addr + page_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_2 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_2) + : "r"(work_response_addr + page_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(work_empty_addr + (page_work_stage) * 8); + page_work_stage += 1; + if (page_work_stage == 2) { page_work_stage = 0; _phase_work_full_3 ^= 1; } + if (_clc_valid_0 == 0) { + break; + } + page_q = _clc_ctaid_0; + page_head = _clc_ctaid_1; + page_batch = _clc_ctaid_2; + } + } + } + // ---- Role: scheduler ---- + if (warp == 10) { + { // scheduler_main + unsigned int scheduler_work_stage = 0; + unsigned int scheduler_throttle_stage = 0; + unsigned int _phase_throttle_full = 0; + unsigned int _phase_work_empty = 1; + unsigned int _phase_work_full_4 = 0; + #pragma unroll 1 + for (unsigned int _work_4 = 0; _work_4 < total_work_items; _work_4++) { + mbarrier_wait(throttle_full_addr + (scheduler_throttle_stage) * 8, _phase_throttle_full); + mbarrier_arrive(throttle_empty_addr + (scheduler_throttle_stage) * 8); + scheduler_throttle_stage += 1; + if (scheduler_throttle_stage == 2) { scheduler_throttle_stage = 0; _phase_throttle_full ^= 1; } + if (elect_sync()) { + mbarrier_wait(work_empty_addr + (scheduler_work_stage) * 8, _phase_work_empty); + mbarrier_arrive_expect_tx(work_full_addr + (scheduler_work_stage) * 8, 16); + asm volatile( + "fence.proxy.async.shared::cta;\n\t" + "clusterlaunchcontrol.try_cancel.async.shared::cta" + ".mbarrier::complete_tx::bytes.b128" + " [%0], [%1];" + :: "r"(work_response_addr + scheduler_work_stage * 16 + 0 * 16), "r"(work_full_addr + scheduler_work_stage * 8) + : "memory"); + } + mbarrier_wait(work_full_addr + (scheduler_work_stage) * 8, _phase_work_full_4); + uint32_t _clc_valid_5 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_5) + : "r"(work_response_addr + scheduler_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_15 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_15) + : "r"(work_response_addr + scheduler_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_16 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_16) + : "r"(work_response_addr + scheduler_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_17 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_17) + : "r"(work_response_addr + scheduler_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(work_empty_addr + (scheduler_work_stage) * 8); + scheduler_work_stage += 1; + if (scheduler_work_stage == 2) { scheduler_work_stage = 0; _phase_work_empty ^= 1; _phase_work_full_4 ^= 1; } + if (_clc_valid_5 == 0) { + break; + } + } + } + } + // ---- Role: padding ---- + if (warp == 11) { + { // padding_main + unsigned int padding_work_stage = 0; + unsigned int _phase_work_full_5 = 0; + #pragma unroll 1 + for (unsigned int _work_5 = 0; _work_5 < total_work_items; _work_5++) { + mbarrier_wait(work_full_addr + (padding_work_stage) * 8, _phase_work_full_5); + uint32_t _clc_valid_6 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_6) + : "r"(work_response_addr + padding_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_18 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_18) + : "r"(work_response_addr + padding_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_19 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_19) + : "r"(work_response_addr + padding_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_20 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_20) + : "r"(work_response_addr + padding_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(work_empty_addr + (padding_work_stage) * 8); + padding_work_stage += 1; + if (padding_work_stage == 2) { padding_work_stage = 0; _phase_work_full_5 ^= 1; } + if (_clc_valid_6 == 0) { + break; + } + } + } + } + // ---- Role: load ---- + if (warp >= 12 && warp <= 15) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 72;"); + { // load_main + const int load_rank = warp - 12; + unsigned int q_prod_stage = 0; + unsigned int kv_prod_stage = 0; + unsigned int page_k_stage = 0; + unsigned int page_k_phase = 0; + unsigned int page_v_stage = 1; + unsigned int page_v_phase = 0; + unsigned int load_work_stage = 0; + unsigned int throttle_stage = 0; + unsigned int load_q = blockIdx.x; + unsigned int load_head = blockIdx.y; + unsigned int load_batch = blockIdx.z; + unsigned int _phase_throttle_empty = 1; + unsigned int _phase_q_empty = 1; + unsigned int _phase_kv_empty = 1; + unsigned int _phase_work_full_6 = 0; + #pragma unroll 1 + for (unsigned int _work_6 = 0; _work_6 < total_work_items; _work_6++) { + mbarrier_wait(throttle_empty_addr + (throttle_stage) * 8, _phase_throttle_empty); + mbarrier_arrive(throttle_full_addr + (throttle_stage) * 8); + throttle_stage += 1; + if (throttle_stage == 2) { throttle_stage = 0; _phase_throttle_empty ^= 1; } + int q_begin_4 = cum_seq_lens_q[(int)load_batch]; + int q_end_4 = cum_seq_lens_q[(int)(load_batch + 1)]; + int q_local_4 = (int)load_q; + int work_valid_4 = q_local_4 < q_end_4 - q_begin_4; + if (work_valid_4 != 0) { + int query_idx_4 = q_begin_4 + q_local_4; + int _max_2 = ((sparse_topk_lens[query_idx_4] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_4] + sparse_topk_lens_offset) : (0)); + int _min_1 = ((_max_2) < (sparse_topk) ? (_max_2) : (sparse_topk)); + int active_topk_4 = _min_1; + int num_tiles_4 = (active_topk_4 + 128 - 1) / 128; + mbarrier_wait(q_empty_addr + (q_prod_stage) * 8, _phase_q_empty); + if (load_rank == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(q_full_addr + (q_prod_stage) * 8, 32768); + } + } + asm volatile("griddepcontrol.wait;" ::: "memory"); + if (load_rank == 0) { + if (elect_sync()) { + #pragma unroll + for (int q_chunk = 0; q_chunk < 4; q_chunk++) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global.tile.mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(smem_q_addr + q_prod_stage * 32768 + (unsigned int)(q_chunk * 8192)), "l"((&tmap_q)), "r"(q_chunk * 128), "r"(0), "r"((int)load_head), "r"(query_idx_4), "r"(q_full_addr + (q_prod_stage) * 8) : "memory"); + } + } + } + q_prod_stage += 1; + if (q_prod_stage == 2) { q_prod_stage = 0; _phase_q_empty ^= 1; } + #pragma unroll 1 + for (int tile_2 = 0; tile_2 < num_tiles_4; tile_2++) { + if ((tile_2 & 1) == 0) { + mbarrier_wait(page_full_addr + (page_k_stage) * 8, page_k_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + } + int page_k_base = smem_page_offsets_addr + page_k_stage * 1024; + #pragma unroll + for (int k_half = 0; k_half < 2; k_half++) { + mbarrier_wait(kv_empty_addr + (kv_prod_stage) * 8, _phase_kv_empty); + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (kv_prod_stage) * 8, 8192); + if (tile_2 == 0) { + #pragma unroll + for (int local_group = 0; local_group < 8; local_group++) { + int group = load_rank * 8 + local_group; + int group_offset = (tile_2 & 1) * 128 + group * 4; + int rows[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&rows[0])), "=r"(*reinterpret_cast(&rows[(0) + 1])), "=r"(*reinterpret_cast(&rows[(0) + 2])), "=r"(*reinterpret_cast(&rows[(0) + 3])) + : "r"(page_k_base + group_offset * 4)); + int k_dst = smem_k_addr + kv_prod_stage * 32768; + tma_gather4_gmem2smem(k_dst + group * 512, (&tmap_swa_kv), k_half * 256, rows[0], rows[1], rows[2], rows[3], kv_full_addr + (kv_prod_stage) * 8); + tma_gather4_gmem2smem(k_dst + 16384 + group * 512, (&tmap_swa_kv), k_half * 256 + 128, rows[0], rows[1], rows[2], rows[3], kv_full_addr + (kv_prod_stage) * 8); + } + } else { + #pragma unroll + for (int local_group_1 = 0; local_group_1 < 8; local_group_1++) { + int group_1 = load_rank * 8 + local_group_1; + int group_offset_1 = (tile_2 & 1) * 128 + group_1 * 4; + int rows_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&rows_1[0])), "=r"(*reinterpret_cast(&rows_1[(0) + 1])), "=r"(*reinterpret_cast(&rows_1[(0) + 2])), "=r"(*reinterpret_cast(&rows_1[(0) + 3])) + : "r"(page_k_base + group_offset_1 * 4)); + int k_dst_1 = smem_k_addr + kv_prod_stage * 32768; + tma_gather4_gmem2smem(k_dst_1 + group_1 * 512, (&tmap_compressed_kv), k_half * 256, rows_1[0], rows_1[1], rows_1[2], rows_1[3], kv_full_addr + (kv_prod_stage) * 8); + tma_gather4_gmem2smem(k_dst_1 + 16384 + group_1 * 512, (&tmap_compressed_kv), k_half * 256 + 128, rows_1[0], rows_1[1], rows_1[2], rows_1[3], kv_full_addr + (kv_prod_stage) * 8); + } + } + } + kv_prod_stage += 1; + if (kv_prod_stage == 4) { kv_prod_stage = 0; _phase_kv_empty ^= 1; } + } + if ((tile_2 & 1) != 0 || tile_2 == num_tiles_4 - 1) { + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(page_empty_addr + (page_k_stage) * 8); + page_k_stage += 1; + if (page_k_stage == 6) { page_k_stage = 0; page_k_phase ^= 1; } + page_k_stage += 1; + if (page_k_stage == 6) { page_k_stage = 0; page_k_phase ^= 1; } + } + if (tile_2 > 0) { + int prev_tile = tile_2 - 1; + if ((prev_tile & 1) == 0) { + mbarrier_wait(page_full_addr + (page_v_stage) * 8, page_v_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + } + int page_v_base = smem_page_offsets_addr + page_v_stage * 1024; + #pragma unroll + for (int v_half_2 = 0; v_half_2 < 2; v_half_2++) { + mbarrier_wait(kv_empty_addr + (kv_prod_stage) * 8, _phase_kv_empty); + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (kv_prod_stage) * 8, 8192); + if (prev_tile == 0) { + #pragma unroll + for (int local_group_2 = 0; local_group_2 < 8; local_group_2++) { + int group_2 = load_rank * 8 + local_group_2; + int group_offset_2 = (prev_tile & 1) * 128 + group_2 * 4; + int rows_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&rows_2[0])), "=r"(*reinterpret_cast(&rows_2[(0) + 1])), "=r"(*reinterpret_cast(&rows_2[(0) + 2])), "=r"(*reinterpret_cast(&rows_2[(0) + 3])) + : "r"(page_v_base + group_offset_2 * 4)); + int v_dst = smem_v_addr + kv_prod_stage * 32768; + tma_gather4_gmem2smem(v_dst + group_2 * 512, (&tmap_swa_kv), v_half_2 * 256, rows_2[0], rows_2[1], rows_2[2], rows_2[3], kv_full_addr + (kv_prod_stage) * 8); + tma_gather4_gmem2smem(v_dst + 16384 + group_2 * 512, (&tmap_swa_kv), v_half_2 * 256 + 128, rows_2[0], rows_2[1], rows_2[2], rows_2[3], kv_full_addr + (kv_prod_stage) * 8); + } + } else { + #pragma unroll + for (int local_group_3 = 0; local_group_3 < 8; local_group_3++) { + int group_3 = load_rank * 8 + local_group_3; + int group_offset_3 = (prev_tile & 1) * 128 + group_3 * 4; + int rows_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&rows_3[0])), "=r"(*reinterpret_cast(&rows_3[(0) + 1])), "=r"(*reinterpret_cast(&rows_3[(0) + 2])), "=r"(*reinterpret_cast(&rows_3[(0) + 3])) + : "r"(page_v_base + group_offset_3 * 4)); + int v_dst_1 = smem_v_addr + kv_prod_stage * 32768; + tma_gather4_gmem2smem(v_dst_1 + group_3 * 512, (&tmap_compressed_kv), v_half_2 * 256, rows_3[0], rows_3[1], rows_3[2], rows_3[3], kv_full_addr + (kv_prod_stage) * 8); + tma_gather4_gmem2smem(v_dst_1 + 16384 + group_3 * 512, (&tmap_compressed_kv), v_half_2 * 256 + 128, rows_3[0], rows_3[1], rows_3[2], rows_3[3], kv_full_addr + (kv_prod_stage) * 8); + } + } + } + kv_prod_stage += 1; + if (kv_prod_stage == 4) { kv_prod_stage = 0; _phase_kv_empty ^= 1; } + } + if ((prev_tile & 1) != 0) { + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(page_empty_addr + (page_v_stage) * 8); + page_v_stage += 1; + if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } + page_v_stage += 1; + if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } + } + } + } + int last_tile = num_tiles_4 - 1; + mbarrier_wait(page_full_addr + (page_v_stage) * 8, page_v_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int page_v_base_1 = smem_page_offsets_addr + page_v_stage * 1024; + #pragma unroll + for (int v_half_3 = 0; v_half_3 < 2; v_half_3++) { + mbarrier_wait(kv_empty_addr + (kv_prod_stage) * 8, _phase_kv_empty); + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (kv_prod_stage) * 8, 8192); + if (last_tile == 0) { + #pragma unroll + for (int local_group_4 = 0; local_group_4 < 8; local_group_4++) { + int group_4 = load_rank * 8 + local_group_4; + int group_offset_4 = (last_tile & 1) * 128 + group_4 * 4; + int rows_4[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&rows_4[0])), "=r"(*reinterpret_cast(&rows_4[(0) + 1])), "=r"(*reinterpret_cast(&rows_4[(0) + 2])), "=r"(*reinterpret_cast(&rows_4[(0) + 3])) + : "r"(page_v_base_1 + group_offset_4 * 4)); + int v_dst_2 = smem_v_addr + kv_prod_stage * 32768; + tma_gather4_gmem2smem(v_dst_2 + group_4 * 512, (&tmap_swa_kv), v_half_3 * 256, rows_4[0], rows_4[1], rows_4[2], rows_4[3], kv_full_addr + (kv_prod_stage) * 8); + tma_gather4_gmem2smem(v_dst_2 + 16384 + group_4 * 512, (&tmap_swa_kv), v_half_3 * 256 + 128, rows_4[0], rows_4[1], rows_4[2], rows_4[3], kv_full_addr + (kv_prod_stage) * 8); + } + } else { + #pragma unroll + for (int local_group_5 = 0; local_group_5 < 8; local_group_5++) { + int group_5 = load_rank * 8 + local_group_5; + int group_offset_5 = (last_tile & 1) * 128 + group_5 * 4; + int rows_5[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&rows_5[0])), "=r"(*reinterpret_cast(&rows_5[(0) + 1])), "=r"(*reinterpret_cast(&rows_5[(0) + 2])), "=r"(*reinterpret_cast(&rows_5[(0) + 3])) + : "r"(page_v_base_1 + group_offset_5 * 4)); + int v_dst_3 = smem_v_addr + kv_prod_stage * 32768; + tma_gather4_gmem2smem(v_dst_3 + group_5 * 512, (&tmap_compressed_kv), v_half_3 * 256, rows_5[0], rows_5[1], rows_5[2], rows_5[3], kv_full_addr + (kv_prod_stage) * 8); + tma_gather4_gmem2smem(v_dst_3 + 16384 + group_5 * 512, (&tmap_compressed_kv), v_half_3 * 256 + 128, rows_5[0], rows_5[1], rows_5[2], rows_5[3], kv_full_addr + (kv_prod_stage) * 8); + } + } + } + kv_prod_stage += 1; + if (kv_prod_stage == 4) { kv_prod_stage = 0; _phase_kv_empty ^= 1; } + } + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(page_empty_addr + (page_v_stage) * 8); + page_v_stage += 1; + if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } + page_v_stage += 1; + if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } + } + mbarrier_wait(work_full_addr + (load_work_stage) * 8, _phase_work_full_6); + uint32_t _clc_valid_1 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_1) + : "r"(work_response_addr + load_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_3 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_3) + : "r"(work_response_addr + load_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_4 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_4) + : "r"(work_response_addr + load_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_5 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_5) + : "r"(work_response_addr + load_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(work_empty_addr + (load_work_stage) * 8); + load_work_stage += 1; + if (load_work_stage == 2) { load_work_stage = 0; _phase_work_full_6 ^= 1; } + if (_clc_valid_1 == 0) { + break; + } + load_q = _clc_ctaid_3; + load_head = _clc_ctaid_4; + load_batch = _clc_ctaid_5; + } + } + } + + // Cleanup +} + +} // extern "C" diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_1bad879f597b2e2adf3e_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_1bad879f597b2e2adf3e_binding.cu new file mode 100644 index 00000000000..8ef2566dad1 --- /dev/null +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_1bad879f597b2e2adf3e_binding.cu @@ -0,0 +1,389 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +// Generated by the Cake source exporter — do not edit. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_1bad879f597b2e2adf3e'. +#include +#include +#include +#include + +#include "tvm_ffi_utils.h" + +#include +#include +#include +#include +#include +#include +#include + +extern "C" __global__ void kernel_cake_dsv4_1bad879f597b2e2adf3e(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int num_head_tiles, int has_sinks); + + +namespace cake_host_shim_406af89a732a5e2c { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t CakeDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +// 4D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + int64_t d3 = t.size(t.ndim() - 3); + TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(d3)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 64u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_swa_kv' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_swa_kv' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_kv'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_O, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_num_head_tiles, int64_t arg_has_sinks, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckContiguous(arg_tmap_q, "tmap_q"); + CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); + CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 4, 16, 1); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_swa_indices, "swa_indices"); + CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); + CheckContiguous(arg_swa_indices, "swa_indices"); + CheckCudaTensor(arg_compressed_indices, "compressed_indices"); + CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); + CheckContiguous(arg_compressed_indices, "compressed_indices"); + CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); + CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); + CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); + CheckContiguous(arg_bmm1_scale, "bmm1_scale"); + CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); + CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); + CheckContiguous(arg_bmm2_scale, "bmm2_scale"); + TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) + << "scalar 'num_heads' value " << arg_num_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) + << "scalar 'swa_index_stride' value " << arg_swa_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) + << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) + << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) + << "scalar 'sparse_topk' value " << arg_sparse_topk + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_head_tiles >= -2147483648LL && arg_num_head_tiles <= 2147483647LL, ValueError) + << "scalar 'num_head_tiles' value " << arg_num_head_tiles + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) + << "scalar 'has_sinks' value " << arg_has_sinks + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); + CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); + CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); + CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + + cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); + __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); + int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); + int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); + int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); + float* p_sinks = static_cast(arg_sinks.data_ptr()); + float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); + float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); + int32_t v_num_heads = (int32_t)arg_num_heads; + int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; + int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; + int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; + int32_t v_sparse_topk = (int32_t)arg_sparse_topk; + int32_t v_num_head_tiles = (int32_t)arg_num_head_tiles; + int32_t v_has_sinks = (int32_t)arg_has_sinks; + void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_O, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_num_head_tiles, &v_has_sinks}; + + static std::mutex smem_mu; + static auto* smem_status_by_device = + new std::unordered_map(); + cudaError_t smem_status = cudaSuccess; + { + std::lock_guard lock(smem_mu); + auto it = smem_status_by_device->find(dev.device_id); + if (it == smem_status_by_device->end()) { + int current_device = -1; + smem_status = cudaGetDevice(¤t_device); + if (smem_status == cudaSuccess && current_device != dev.device_id) { + smem_status = cudaErrorInvalidDevice; + } + if (smem_status == cudaSuccess) { + smem_status = cudaFuncSetAttribute( + reinterpret_cast(kernel_cake_dsv4_1bad879f597b2e2adf3e), + cudaFuncAttributeMaxDynamicSharedMemorySize, + 198528); + } + smem_status_by_device->emplace(dev.device_id, smem_status); + } else { + smem_status = it->second; + } + } + TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) + << "cudaFuncSetAttribute for kernel_cake_dsv4_1bad879f597b2e2adf3e failed: " + << cudaGetErrorString(smem_status); + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(512u, 1u, 1u); + + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes = 198528u; + config.stream = stream; + cudaError_t launch_status = cudaLaunchKernelExC( + &config, reinterpret_cast(kernel_cake_dsv4_1bad879f597b2e2adf3e), kargs); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernelExC for kernel_cake_dsv4_1bad879f597b2e2adf3e failed: " + << cudaGetErrorString(launch_status); +} + +} // namespace cake_host_shim_406af89a732a5e2c + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_406af89a732a5e2c::Run); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_1bad879f597b2e2adf3e_kernel.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_1bad879f597b2e2adf3e_kernel.cu new file mode 100644 index 00000000000..0a4d274ffc7 --- /dev/null +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_1bad879f597b2e2adf3e_kernel.cu @@ -0,0 +1,1742 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +#define CAKE_INF CUDART_INF_F +#define TMEM_NCOLS 384 +#define TMEM_TMEM_OFFSET 0 +#define NUM_KV_PIPE_STAGES 4 +#define SMEM_SMEM_Q_OFF 1024 +#define SMEM_SMEM_Q_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_STRIDE 8192 +#define SMEM_SMEM_KV_OFF 66560 +#define SMEM_SMEM_KV_STAGE_BYTES 32768 +#define SMEM_SMEM_KV_STRIDE 32768 +#define SMEM_SMEM_V_OFF 66560 +#define SMEM_SMEM_V_STAGE_BYTES 32768 +#define SMEM_SMEM_V_STRIDE 32768 +#define SMEM_SMEM_SUM_OFF 197632 +#define SMEM_SMEM_SUM_STAGE_BYTES 256 +#define SMEM_SMEM_SUM_STRIDE 256 +#define SMEM_SMEM_INDEX_FLAGS_OFF 198400 +#define SMEM_SMEM_INDEX_FLAGS_STAGE_BYTES 16 +#define SMEM_SMEM_INDEX_FLAGS_STRIDE 16 +#define SMEM_SMEM_INDICES_OFF 197888 +#define SMEM_SMEM_INDICES_STAGE_BYTES 512 +#define SMEM_SMEM_INDICES_STRIDE 512 +#define SMEM_TOTAL 198528 +#define THREADS 512 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count) : "memory"); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_2d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3}], [%4];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr) { + // Canonical .shared::cta form for non-multicast gather4, matching + // trtllm-gen / cuda_ptx and the PTX ISA qualifier order + // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, + // .shared::cluster is reserved for the multicast variant (ctaMask). + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(512, 1) void +kernel_cake_dsv4_1bad879f597b2e2adf3e(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int num_head_tiles, int has_sinks) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int mbar_base = smem; + #define q_full_addr (mbar_base + 0) + #define kv_full_addr (mbar_base + 8) + #define kv_empty_addr (mbar_base + 40) + #define s_full_addr (mbar_base + 72) + #define p_full_addr (mbar_base + 80) + #define sum_ready_addr (mbar_base + 88) + #define o_done_addr (mbar_base + 96) + #define tmem_dealloc_addr (mbar_base + 104) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + + // Kernel setup ops + __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_q_addr = smem + 1024; + __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_kv_addr = smem + 66560; + __nv_bfloat16* smem_v = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_v_addr = smem + 66560; + float* smem_sum = reinterpret_cast(smem_raw + 197632); + const int smem_sum_addr = smem + 197632; + int* smem_index_flags = reinterpret_cast(smem_raw + 198400); + const int smem_index_flags_addr = smem + 198400; + int* smem_indices = reinterpret_cast(smem_raw + 197888); + const int smem_indices_addr = smem + 197888; + + // Mbarrier init (8 pipeline groups, 0 ordered-sequence groups, 14 barriers) + // Mbarriers at smem_raw[0..112) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // q_full: 1 barriers, init_count=1 + mbarrier_init(smem + 0, 1); + // --- pipeline 'kv_pipe' --- + // kv_full: 4 barriers, init_count=4 + mbarrier_init(smem + 8, 4); + mbarrier_init(smem + 16, 4); + mbarrier_init(smem + 24, 4); + mbarrier_init(smem + 32, 4); + // kv_empty: 4 barriers, init_count=1 + mbarrier_init(smem + 40, 1); + mbarrier_init(smem + 48, 1); + mbarrier_init(smem + 56, 1); + mbarrier_init(smem + 64, 1); + // s_full: 1 barriers, init_count=1 + mbarrier_init(smem + 72, 1); + // p_full: 1 barriers, init_count=128 + mbarrier_init(smem + 80, 128); + // sum_ready: 1 barriers, init_count=128 + mbarrier_init(smem + 88, 128); + // o_done: 1 barriers, init_count=1 + mbarrier_init(smem + 96, 1); + // tmem_dealloc: 1 barriers, init_count=256 + mbarrier_init(smem + 104, 256); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncwarp(); + + // TMEM alloc (512 columns, 384 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 112); + if (warp == 0) { + int _tmem_hold = smem + 112; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_tmem = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 64;"); + } + + // ---- Role: softmax ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); + { // softmax_main + float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; + int query_head_work = blockIdx.x >> 2; + int query_idx = query_head_work / num_head_tiles; + int head_tile = query_head_work % num_head_tiles; + int _max_0 = ((sparse_topk_lens[query_idx] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx] + sparse_topk_lens_offset) : (0)); + int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); + int active_topk = _min_0; + const int warp_in_compute = warp; + const int tmem_row_origin = warp_in_compute * 32; + const int logical_row_origin = warp_in_compute * 16; + const int my_row = logical_row_origin + lane % 16; + int head_idx = head_tile * 64 + my_row; + const int col_half = lane / 16; + unsigned int _phase_s_full_0 = 0; + mbarrier_wait(s_full_addr, _phase_s_full_0); + _phase_s_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int score_addr = taddr + (unsigned int)(tmem_row_origin << 16); + float _tmem_load_0[64]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_0[0])), "=r"(*reinterpret_cast(&_tmem_load_0[1])), "=r"(*reinterpret_cast(&_tmem_load_0[2])), "=r"(*reinterpret_cast(&_tmem_load_0[3])), "=r"(*reinterpret_cast(&_tmem_load_0[4])), "=r"(*reinterpret_cast(&_tmem_load_0[5])), "=r"(*reinterpret_cast(&_tmem_load_0[6])), "=r"(*reinterpret_cast(&_tmem_load_0[7])), "=r"(*reinterpret_cast(&_tmem_load_0[8])), "=r"(*reinterpret_cast(&_tmem_load_0[9])), "=r"(*reinterpret_cast(&_tmem_load_0[10])), "=r"(*reinterpret_cast(&_tmem_load_0[11])), "=r"(*reinterpret_cast(&_tmem_load_0[12])), "=r"(*reinterpret_cast(&_tmem_load_0[13])), "=r"(*reinterpret_cast(&_tmem_load_0[14])), "=r"(*reinterpret_cast(&_tmem_load_0[15])), "=r"(*reinterpret_cast(&_tmem_load_0[16])), "=r"(*reinterpret_cast(&_tmem_load_0[17])), "=r"(*reinterpret_cast(&_tmem_load_0[18])), "=r"(*reinterpret_cast(&_tmem_load_0[19])), "=r"(*reinterpret_cast(&_tmem_load_0[20])), "=r"(*reinterpret_cast(&_tmem_load_0[21])), "=r"(*reinterpret_cast(&_tmem_load_0[22])), "=r"(*reinterpret_cast(&_tmem_load_0[23])), "=r"(*reinterpret_cast(&_tmem_load_0[24])), "=r"(*reinterpret_cast(&_tmem_load_0[25])), "=r"(*reinterpret_cast(&_tmem_load_0[26])), "=r"(*reinterpret_cast(&_tmem_load_0[27])), "=r"(*reinterpret_cast(&_tmem_load_0[28])), "=r"(*reinterpret_cast(&_tmem_load_0[29])), "=r"(*reinterpret_cast(&_tmem_load_0[30])), "=r"(*reinterpret_cast(&_tmem_load_0[31])) + : "r"(score_addr)); + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_0[32])), "=r"(*reinterpret_cast(&_tmem_load_0[33])), "=r"(*reinterpret_cast(&_tmem_load_0[34])), "=r"(*reinterpret_cast(&_tmem_load_0[35])), "=r"(*reinterpret_cast(&_tmem_load_0[36])), "=r"(*reinterpret_cast(&_tmem_load_0[37])), "=r"(*reinterpret_cast(&_tmem_load_0[38])), "=r"(*reinterpret_cast(&_tmem_load_0[39])), "=r"(*reinterpret_cast(&_tmem_load_0[40])), "=r"(*reinterpret_cast(&_tmem_load_0[41])), "=r"(*reinterpret_cast(&_tmem_load_0[42])), "=r"(*reinterpret_cast(&_tmem_load_0[43])), "=r"(*reinterpret_cast(&_tmem_load_0[44])), "=r"(*reinterpret_cast(&_tmem_load_0[45])), "=r"(*reinterpret_cast(&_tmem_load_0[46])), "=r"(*reinterpret_cast(&_tmem_load_0[47])), "=r"(*reinterpret_cast(&_tmem_load_0[48])), "=r"(*reinterpret_cast(&_tmem_load_0[49])), "=r"(*reinterpret_cast(&_tmem_load_0[50])), "=r"(*reinterpret_cast(&_tmem_load_0[51])), "=r"(*reinterpret_cast(&_tmem_load_0[52])), "=r"(*reinterpret_cast(&_tmem_load_0[53])), "=r"(*reinterpret_cast(&_tmem_load_0[54])), "=r"(*reinterpret_cast(&_tmem_load_0[55])), "=r"(*reinterpret_cast(&_tmem_load_0[56])), "=r"(*reinterpret_cast(&_tmem_load_0[57])), "=r"(*reinterpret_cast(&_tmem_load_0[58])), "=r"(*reinterpret_cast(&_tmem_load_0[59])), "=r"(*reinterpret_cast(&_tmem_load_0[60])), "=r"(*reinterpret_cast(&_tmem_load_0[61])), "=r"(*reinterpret_cast(&_tmem_load_0[62])), "=r"(*reinterpret_cast(&_tmem_load_0[63])) + : "r"(score_addr + 32)); + int valid_cols = active_topk - col_half * 64; + int causal_cols = smem_index_flags[col_half * 2 + 1]; + valid_cols = ((valid_cols < causal_cols) ? valid_cols : causal_cols); + uint32_t _slice_lo_mask_0; + { + int _lim_0 = valid_cols; + if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_0 & (1u << 0))) _tmem_load_0[0] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 1))) _tmem_load_0[1] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 2))) _tmem_load_0[2] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 3))) _tmem_load_0[3] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 4))) _tmem_load_0[4] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 5))) _tmem_load_0[5] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 6))) _tmem_load_0[6] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 7))) _tmem_load_0[7] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 8))) _tmem_load_0[8] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 9))) _tmem_load_0[9] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 10))) _tmem_load_0[10] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 11))) _tmem_load_0[11] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 12))) _tmem_load_0[12] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 13))) _tmem_load_0[13] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 14))) _tmem_load_0[14] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 15))) _tmem_load_0[15] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 16))) _tmem_load_0[16] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 17))) _tmem_load_0[17] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 18))) _tmem_load_0[18] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 19))) _tmem_load_0[19] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 20))) _tmem_load_0[20] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 21))) _tmem_load_0[21] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 22))) _tmem_load_0[22] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 23))) _tmem_load_0[23] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 24))) _tmem_load_0[24] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 25))) _tmem_load_0[25] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 26))) _tmem_load_0[26] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 27))) _tmem_load_0[27] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 28))) _tmem_load_0[28] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 29))) _tmem_load_0[29] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 30))) _tmem_load_0[30] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 31))) _tmem_load_0[31] = -CAKE_INF; + uint32_t _slice_lo_mask_1; + { + int _lim_1 = valid_cols - 32; + if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) _tmem_load_0[32] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) _tmem_load_0[33] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) _tmem_load_0[34] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) _tmem_load_0[35] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) _tmem_load_0[36] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) _tmem_load_0[37] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) _tmem_load_0[38] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) _tmem_load_0[39] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) _tmem_load_0[40] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) _tmem_load_0[41] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) _tmem_load_0[42] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) _tmem_load_0[43] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) _tmem_load_0[44] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) _tmem_load_0[45] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) _tmem_load_0[46] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) _tmem_load_0[47] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) _tmem_load_0[48] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) _tmem_load_0[49] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) _tmem_load_0[50] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) _tmem_load_0[51] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) _tmem_load_0[52] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) _tmem_load_0[53] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) _tmem_load_0[54] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) _tmem_load_0[55] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) _tmem_load_0[56] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) _tmem_load_0[57] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) _tmem_load_0[58] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) _tmem_load_0[59] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) _tmem_load_0[60] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) _tmem_load_0[61] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) _tmem_load_0[62] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) _tmem_load_0[63] = -CAKE_INF; + float2 _reg_reduce_max2_2 = {-CAKE_INF, -CAKE_INF}; + row_max_x32_accum(&_tmem_load_0[0], _reg_reduce_max2_2); + row_max_x32_accum(&_tmem_load_0[32], _reg_reduce_max2_2); + float _tmem_load_0_max = row_max_reduce(_reg_reduce_max2_2); + float row_max = _tmem_load_0_max; + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, row_max, 16); + float _max_1 = max_noftz(row_max, _shfl_xor_0); + row_max = _max_1; + if (has_sinks != 0 && head_idx < num_heads) { + float sink_unscaled = sinks[head_idx] * 1.4426950408889634f / softmax_scale_log2; + float _max_2 = max_noftz(row_max, sink_unscaled); + row_max = _max_2; + } + bool empty_row = row_max == -CAKE_INF; + float safe_max = ((empty_row) ? 0.0f : row_max); + float max_scaled = safe_max * softmax_scale_log2; + float2 _f2_0 = make_float2(softmax_scale_log2, softmax_scale_log2); + float2 _f2_1 = make_float2(-max_scaled, -max_scaled); + float2 _f2_2 = make_float2(0.0f, 0.0f); + float2 sum0 = _f2_2; + float2 _f2_3 = make_float2(0.0f, 0.0f); + float2 sum1 = _f2_3; + float2 _f2_4 = make_float2(0.0f, 0.0f); + float2 sum2 = _f2_4; + float2 _f2_5 = make_float2(0.0f, 0.0f); + float2 sum3 = _f2_5; + #pragma unroll + for (int i = 0; i < 64; i += 8) { + float2 _f2_6 = make_float2(_tmem_load_0[i], _tmem_load_0[i + 1]); + float2 _f2_7 = make_float2(_tmem_load_0[i + 2], _tmem_load_0[i + 3]); + float2 _f2_8 = make_float2(_tmem_load_0[i + 4], _tmem_load_0[i + 5]); + float2 _f2_9 = make_float2(_tmem_load_0[i + 6], _tmem_load_0[i + 7]); + float2 affine0 = fma_f32x2_rn_ftz(_f2_6, _f2_0, _f2_1); + float2 affine1 = fma_f32x2_rn_ftz(_f2_7, _f2_0, _f2_1); + float2 affine2 = fma_f32x2_rn_ftz(_f2_8, _f2_0, _f2_1); + float2 affine3 = fma_f32x2_rn_ftz(_f2_9, _f2_0, _f2_1); + float _exp2_0 = approx_exp2(affine0.x); + float exp0 = _exp2_0; + float _exp2_1 = approx_exp2(affine0.y); + float exp1 = _exp2_1; + float _exp2_2 = approx_exp2(affine1.x); + float exp2 = _exp2_2; + float _exp2_3 = approx_exp2(affine1.y); + float exp3 = _exp2_3; + float _exp2_4 = approx_exp2(affine2.x); + float exp4 = _exp2_4; + float _exp2_5 = approx_exp2(affine2.y); + float exp5 = _exp2_5; + float _exp2_6 = approx_exp2(affine3.x); + float exp6 = _exp2_6; + float _exp2_7 = approx_exp2(affine3.y); + float exp7 = _exp2_7; + _tmem_load_0[i] = exp0; + _tmem_load_0[i + 1] = exp1; + _tmem_load_0[i + 2] = exp2; + _tmem_load_0[i + 3] = exp3; + _tmem_load_0[i + 4] = exp4; + _tmem_load_0[i + 5] = exp5; + _tmem_load_0[i + 6] = exp6; + _tmem_load_0[i + 7] = exp7; + float2 _f2_10 = make_float2(exp0, exp1); + sum0 = add_f32x2(sum0, _f2_10); + float2 _f2_11 = make_float2(exp2, exp3); + sum1 = add_f32x2(sum1, _f2_11); + float2 _f2_12 = make_float2(exp4, exp5); + sum2 = add_f32x2(sum2, _f2_12); + float2 _f2_13 = make_float2(exp6, exp7); + sum3 = add_f32x2(sum3, _f2_13); + } + float2 sum01 = add_f32x2(sum0, sum1); + float2 sum23 = add_f32x2(sum2, sum3); + float2 row_sum_pair = add_f32x2(sum01, sum23); + float row_sum = row_sum_pair.x + row_sum_pair.y; + float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, row_sum, 16); + row_sum = row_sum + _shfl_xor_1; + if (has_sinks != 0 && head_idx < num_heads) { + float _exp2_8 = approx_exp2(sinks[head_idx] * 1.4426950408889634f - max_scaled); + row_sum = row_sum + _exp2_8; + } + unsigned int packed_p[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_0[_lp*2 + 0], _tmem_load_0[_lp*2+1 + 0])); + packed_p[_lp] = *(uint32_t*)&_bf2; + } + int p_addr = taddr + 128 + (unsigned int)(tmem_row_origin << 16); + asm volatile( + "tcgen05.st.sync.aligned.16x32bx2.x16.b32" + " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(p_addr), "r"(*reinterpret_cast(&packed_p[0])), "r"(*reinterpret_cast(&packed_p[1])), "r"(*reinterpret_cast(&packed_p[2])), "r"(*reinterpret_cast(&packed_p[3])), "r"(*reinterpret_cast(&packed_p[4])), "r"(*reinterpret_cast(&packed_p[5])), "r"(*reinterpret_cast(&packed_p[6])), "r"(*reinterpret_cast(&packed_p[7])), "r"(*reinterpret_cast(&packed_p[8])), "r"(*reinterpret_cast(&packed_p[9])), "r"(*reinterpret_cast(&packed_p[10])), "r"(*reinterpret_cast(&packed_p[11])), "r"(*reinterpret_cast(&packed_p[12])), "r"(*reinterpret_cast(&packed_p[13])), "r"(*reinterpret_cast(&packed_p[14])), "r"(*reinterpret_cast(&packed_p[15]))); + asm volatile( + "tcgen05.st.sync.aligned.16x32bx2.x16.b32" + " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(p_addr + 16), "r"(*reinterpret_cast(&(packed_p + 16)[0])), "r"(*reinterpret_cast(&(packed_p + 16)[1])), "r"(*reinterpret_cast(&(packed_p + 16)[2])), "r"(*reinterpret_cast(&(packed_p + 16)[3])), "r"(*reinterpret_cast(&(packed_p + 16)[4])), "r"(*reinterpret_cast(&(packed_p + 16)[5])), "r"(*reinterpret_cast(&(packed_p + 16)[6])), "r"(*reinterpret_cast(&(packed_p + 16)[7])), "r"(*reinterpret_cast(&(packed_p + 16)[8])), "r"(*reinterpret_cast(&(packed_p + 16)[9])), "r"(*reinterpret_cast(&(packed_p + 16)[10])), "r"(*reinterpret_cast(&(packed_p + 16)[11])), "r"(*reinterpret_cast(&(packed_p + 16)[12])), "r"(*reinterpret_cast(&(packed_p + 16)[13])), "r"(*reinterpret_cast(&(packed_p + 16)[14])), "r"(*reinterpret_cast(&(packed_p + 16)[15]))); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + if (col_half == 0) { + float _max_3 = max_noftz(row_sum, 1.1754943508222875e-38f); + float published_sum = _max_3; + smem_sum[my_row] = published_sum; + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(p_full_addr); + mbarrier_arrive(sum_ready_addr); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_dealloc_addr); + } + } + // ---- Role: epilogue ---- + if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 128;"); + { // epilogue_main + float output_scale = bmm2_scale[0]; + int work_idx = blockIdx.x; + int query_head_work_1 = work_idx >> 2; + int query_idx_1 = query_head_work_1 / num_head_tiles; + int head_tile_1 = query_head_work_1 % num_head_tiles; + int v_chunk = work_idx & 3; + const int warp_in_role = warp - 4; + const int tmem_row_origin_1 = warp_in_role * 32; + const int logical_row_origin_1 = warp_in_role * 16; + const int my_row_1 = logical_row_origin_1 + lane % 16; + int head_idx_1 = head_tile_1 * 64 + my_row_1; + const int col_half_1 = lane / 16; + const int row_addr = tmem_row_origin_1 << 16; + unsigned int _phase_o_done_0 = 0; + mbarrier_wait(o_done_addr, _phase_o_done_0); + _phase_o_done_0 ^= 1; + unsigned int _phase_sum_ready_0 = 0; + mbarrier_wait(sum_ready_addr, _phase_sum_ready_0); + _phase_sum_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + float row_sum_1 = smem_sum[my_row_1]; + float _rcp_0 = approx_rcp(row_sum_1); + float inv_sum = _rcp_0; + int output_base = (query_idx_1 * num_heads + head_idx_1) * 512 + v_chunk * 128; + float _tmem_load_1[64]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_1[0])), "=r"(*reinterpret_cast(&_tmem_load_1[1])), "=r"(*reinterpret_cast(&_tmem_load_1[2])), "=r"(*reinterpret_cast(&_tmem_load_1[3])), "=r"(*reinterpret_cast(&_tmem_load_1[4])), "=r"(*reinterpret_cast(&_tmem_load_1[5])), "=r"(*reinterpret_cast(&_tmem_load_1[6])), "=r"(*reinterpret_cast(&_tmem_load_1[7])), "=r"(*reinterpret_cast(&_tmem_load_1[8])), "=r"(*reinterpret_cast(&_tmem_load_1[9])), "=r"(*reinterpret_cast(&_tmem_load_1[10])), "=r"(*reinterpret_cast(&_tmem_load_1[11])), "=r"(*reinterpret_cast(&_tmem_load_1[12])), "=r"(*reinterpret_cast(&_tmem_load_1[13])), "=r"(*reinterpret_cast(&_tmem_load_1[14])), "=r"(*reinterpret_cast(&_tmem_load_1[15])), "=r"(*reinterpret_cast(&_tmem_load_1[16])), "=r"(*reinterpret_cast(&_tmem_load_1[17])), "=r"(*reinterpret_cast(&_tmem_load_1[18])), "=r"(*reinterpret_cast(&_tmem_load_1[19])), "=r"(*reinterpret_cast(&_tmem_load_1[20])), "=r"(*reinterpret_cast(&_tmem_load_1[21])), "=r"(*reinterpret_cast(&_tmem_load_1[22])), "=r"(*reinterpret_cast(&_tmem_load_1[23])), "=r"(*reinterpret_cast(&_tmem_load_1[24])), "=r"(*reinterpret_cast(&_tmem_load_1[25])), "=r"(*reinterpret_cast(&_tmem_load_1[26])), "=r"(*reinterpret_cast(&_tmem_load_1[27])), "=r"(*reinterpret_cast(&_tmem_load_1[28])), "=r"(*reinterpret_cast(&_tmem_load_1[29])), "=r"(*reinterpret_cast(&_tmem_load_1[30])), "=r"(*reinterpret_cast(&_tmem_load_1[31])) + : "r"(taddr + 256 + (unsigned int)row_addr)); + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_1[32])), "=r"(*reinterpret_cast(&_tmem_load_1[33])), "=r"(*reinterpret_cast(&_tmem_load_1[34])), "=r"(*reinterpret_cast(&_tmem_load_1[35])), "=r"(*reinterpret_cast(&_tmem_load_1[36])), "=r"(*reinterpret_cast(&_tmem_load_1[37])), "=r"(*reinterpret_cast(&_tmem_load_1[38])), "=r"(*reinterpret_cast(&_tmem_load_1[39])), "=r"(*reinterpret_cast(&_tmem_load_1[40])), "=r"(*reinterpret_cast(&_tmem_load_1[41])), "=r"(*reinterpret_cast(&_tmem_load_1[42])), "=r"(*reinterpret_cast(&_tmem_load_1[43])), "=r"(*reinterpret_cast(&_tmem_load_1[44])), "=r"(*reinterpret_cast(&_tmem_load_1[45])), "=r"(*reinterpret_cast(&_tmem_load_1[46])), "=r"(*reinterpret_cast(&_tmem_load_1[47])), "=r"(*reinterpret_cast(&_tmem_load_1[48])), "=r"(*reinterpret_cast(&_tmem_load_1[49])), "=r"(*reinterpret_cast(&_tmem_load_1[50])), "=r"(*reinterpret_cast(&_tmem_load_1[51])), "=r"(*reinterpret_cast(&_tmem_load_1[52])), "=r"(*reinterpret_cast(&_tmem_load_1[53])), "=r"(*reinterpret_cast(&_tmem_load_1[54])), "=r"(*reinterpret_cast(&_tmem_load_1[55])), "=r"(*reinterpret_cast(&_tmem_load_1[56])), "=r"(*reinterpret_cast(&_tmem_load_1[57])), "=r"(*reinterpret_cast(&_tmem_load_1[58])), "=r"(*reinterpret_cast(&_tmem_load_1[59])), "=r"(*reinterpret_cast(&_tmem_load_1[60])), "=r"(*reinterpret_cast(&_tmem_load_1[61])), "=r"(*reinterpret_cast(&_tmem_load_1[62])), "=r"(*reinterpret_cast(&_tmem_load_1[63])) + : "r"(taddr + 256 + (unsigned int)row_addr + 32)); + float2 _f2_14 = make_float2(inv_sum * output_scale, inv_sum * output_scale); + #pragma unroll + for (int i_1 = 0; i_1 < 64; i_1 += 2) { + float2 _f2_15 = make_float2(_tmem_load_1[i_1], _tmem_load_1[i_1 + 1]); + float2 _mul_f32x2_0; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_0) : "l"(*(const unsigned long long*)&_f2_15), "l"(*(const unsigned long long*)&_f2_14)); + _tmem_load_1[i_1] = _mul_f32x2_0.x; + _tmem_load_1[i_1 + 1] = _mul_f32x2_0.y; + } + unsigned int packed_o[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_1[_lp*2 + 0], _tmem_load_1[_lp*2+1 + 0])); + packed_o[_lp] = *(uint32_t*)&_bf2; + } + bool stage_hi_1 = (lane & 1) != 0; + bool stage_hi_2 = (lane & 2) != 0; + bool stage_hi_4 = (lane & 4) != 0; + unsigned int _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[0] : packed_o[4]), 1); + packed_o[0] = ((stage_hi_1) ? _shfl_xor_2 : packed_o[0]); + packed_o[4] = ((stage_hi_1) ? packed_o[4] : _shfl_xor_2); + unsigned int _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[1] : packed_o[5]), 1); + packed_o[1] = ((stage_hi_1) ? _shfl_xor_3 : packed_o[1]); + packed_o[5] = ((stage_hi_1) ? packed_o[5] : _shfl_xor_3); + unsigned int _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[2] : packed_o[6]), 1); + packed_o[2] = ((stage_hi_1) ? _shfl_xor_4 : packed_o[2]); + packed_o[6] = ((stage_hi_1) ? packed_o[6] : _shfl_xor_4); + unsigned int _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[3] : packed_o[7]), 1); + packed_o[3] = ((stage_hi_1) ? _shfl_xor_5 : packed_o[3]); + packed_o[7] = ((stage_hi_1) ? packed_o[7] : _shfl_xor_5); + unsigned int _shfl_xor_6 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[8] : packed_o[12]), 1); + packed_o[8] = ((stage_hi_1) ? _shfl_xor_6 : packed_o[8]); + packed_o[12] = ((stage_hi_1) ? packed_o[12] : _shfl_xor_6); + unsigned int _shfl_xor_7 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[9] : packed_o[13]), 1); + packed_o[9] = ((stage_hi_1) ? _shfl_xor_7 : packed_o[9]); + packed_o[13] = ((stage_hi_1) ? packed_o[13] : _shfl_xor_7); + unsigned int _shfl_xor_8 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[10] : packed_o[14]), 1); + packed_o[10] = ((stage_hi_1) ? _shfl_xor_8 : packed_o[10]); + packed_o[14] = ((stage_hi_1) ? packed_o[14] : _shfl_xor_8); + unsigned int _shfl_xor_9 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[11] : packed_o[15]), 1); + packed_o[11] = ((stage_hi_1) ? _shfl_xor_9 : packed_o[11]); + packed_o[15] = ((stage_hi_1) ? packed_o[15] : _shfl_xor_9); + unsigned int _shfl_xor_10 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[16] : packed_o[20]), 1); + packed_o[16] = ((stage_hi_1) ? _shfl_xor_10 : packed_o[16]); + packed_o[20] = ((stage_hi_1) ? packed_o[20] : _shfl_xor_10); + unsigned int _shfl_xor_11 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[17] : packed_o[21]), 1); + packed_o[17] = ((stage_hi_1) ? _shfl_xor_11 : packed_o[17]); + packed_o[21] = ((stage_hi_1) ? packed_o[21] : _shfl_xor_11); + unsigned int _shfl_xor_12 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[18] : packed_o[22]), 1); + packed_o[18] = ((stage_hi_1) ? _shfl_xor_12 : packed_o[18]); + packed_o[22] = ((stage_hi_1) ? packed_o[22] : _shfl_xor_12); + unsigned int _shfl_xor_13 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[19] : packed_o[23]), 1); + packed_o[19] = ((stage_hi_1) ? _shfl_xor_13 : packed_o[19]); + packed_o[23] = ((stage_hi_1) ? packed_o[23] : _shfl_xor_13); + unsigned int _shfl_xor_14 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[24] : packed_o[28]), 1); + packed_o[24] = ((stage_hi_1) ? _shfl_xor_14 : packed_o[24]); + packed_o[28] = ((stage_hi_1) ? packed_o[28] : _shfl_xor_14); + unsigned int _shfl_xor_15 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[25] : packed_o[29]), 1); + packed_o[25] = ((stage_hi_1) ? _shfl_xor_15 : packed_o[25]); + packed_o[29] = ((stage_hi_1) ? packed_o[29] : _shfl_xor_15); + unsigned int _shfl_xor_16 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[26] : packed_o[30]), 1); + packed_o[26] = ((stage_hi_1) ? _shfl_xor_16 : packed_o[26]); + packed_o[30] = ((stage_hi_1) ? packed_o[30] : _shfl_xor_16); + unsigned int _shfl_xor_17 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[27] : packed_o[31]), 1); + packed_o[27] = ((stage_hi_1) ? _shfl_xor_17 : packed_o[27]); + packed_o[31] = ((stage_hi_1) ? packed_o[31] : _shfl_xor_17); + int lane_t = lane & 1; + unsigned int chunk_words[4]; + chunk_words[0] = packed_o[0]; + chunk_words[1] = packed_o[1]; + chunk_words[2] = packed_o[2]; + chunk_words[3] = packed_o[3]; + if (head_tile_1 * 64 + (my_row_1 & -2 | 0) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + head_tile_1 * 64 + (my_row_1 & -2 | 0)) * 512 + v_chunk * 128 + col_half_1 * 64 + (0 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[4]; + chunk_words[1] = packed_o[5]; + chunk_words[2] = packed_o[6]; + chunk_words[3] = packed_o[7]; + if (head_tile_1 * 64 + (my_row_1 & -2 | 1) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + head_tile_1 * 64 + (my_row_1 & -2 | 1)) * 512 + v_chunk * 128 + col_half_1 * 64 + (0 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[8]; + chunk_words[1] = packed_o[9]; + chunk_words[2] = packed_o[10]; + chunk_words[3] = packed_o[11]; + if (head_tile_1 * 64 + (my_row_1 & -2 | 0) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + head_tile_1 * 64 + (my_row_1 & -2 | 0)) * 512 + v_chunk * 128 + col_half_1 * 64 + (2 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[12]; + chunk_words[1] = packed_o[13]; + chunk_words[2] = packed_o[14]; + chunk_words[3] = packed_o[15]; + if (head_tile_1 * 64 + (my_row_1 & -2 | 1) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + head_tile_1 * 64 + (my_row_1 & -2 | 1)) * 512 + v_chunk * 128 + col_half_1 * 64 + (2 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[16]; + chunk_words[1] = packed_o[17]; + chunk_words[2] = packed_o[18]; + chunk_words[3] = packed_o[19]; + if (head_tile_1 * 64 + (my_row_1 & -2 | 0) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + head_tile_1 * 64 + (my_row_1 & -2 | 0)) * 512 + v_chunk * 128 + col_half_1 * 64 + (4 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[20]; + chunk_words[1] = packed_o[21]; + chunk_words[2] = packed_o[22]; + chunk_words[3] = packed_o[23]; + if (head_tile_1 * 64 + (my_row_1 & -2 | 1) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + head_tile_1 * 64 + (my_row_1 & -2 | 1)) * 512 + v_chunk * 128 + col_half_1 * 64 + (4 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[24]; + chunk_words[1] = packed_o[25]; + chunk_words[2] = packed_o[26]; + chunk_words[3] = packed_o[27]; + if (head_tile_1 * 64 + (my_row_1 & -2 | 0) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + head_tile_1 * 64 + (my_row_1 & -2 | 0)) * 512 + v_chunk * 128 + col_half_1 * 64 + (6 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[28]; + chunk_words[1] = packed_o[29]; + chunk_words[2] = packed_o[30]; + chunk_words[3] = packed_o[31]; + if (head_tile_1 * 64 + (my_row_1 & -2 | 1) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + head_tile_1 * 64 + (my_row_1 & -2 | 1)) * 512 + v_chunk * 128 + col_half_1 * 64 + (6 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_dealloc_addr); + } + } + // ---- Role: mma_warp ---- + if (warp == 8) { + { // mma_warp_main + unsigned int _phase_q_full_0 = 0; + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + unsigned int _phase_kv_full = 0; + #pragma unroll + for (int k_stage = 0; k_stage < 4; k_stage++) { + mbarrier_wait(kv_full_addr + (k_stage) * 8, _phase_kv_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_0 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (k_stage * 2) * 512); + int _mma_b_lo_0 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (k_stage) * 2048); + { + uint64_t _mma_ss_a_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_0); + uint64_t _mma_ss_b_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_0); + if (elect_sync()) { + tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, ((k_stage == 0) ? 0 : 1)); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); + } + } + int _mma_a_lo_1 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (k_stage * 2 + 1) * 512); + int _mma_b_lo_1 = make_warp_uniform((((smem_kv_addr + 16384) >> 4) & 0x3FFF) + (k_stage) * 2048); + { + uint64_t _mma_ss_a_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_1); + uint64_t _mma_ss_b_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_1); + if (elect_sync()) { + tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); + } + } + if (k_stage == 3) { + elect_commit(s_full_addr); + } + elect_commit(kv_empty_addr + (k_stage) * 8); + } + int _mma_prepared_b_lo_2 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (0) * 2048); + unsigned int _phase_p_full_0 = 0; + mbarrier_wait(p_full_addr, _phase_p_full_0); + _phase_p_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + mbarrier_wait(kv_full_addr, 1); + asm volatile("tcgen05.fence::after_thread_sync;"); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 dhi, blo, ta, id;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 69272720;\n\t" + "mov.b32 ta, %2;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p0;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "}\n" + :: "r"((tmem_tmem + (256))), "r"(_mma_prepared_b_lo_2), "r"(tmem_tmem + 128), "r"(0)); + elect_commit(o_done_addr); + elect_commit(kv_empty_addr); + unsigned int _phase_tmem_dealloc_0 = 0; + mbarrier_wait(tmem_dealloc_addr, _phase_tmem_dealloc_0); + _phase_tmem_dealloc_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + } + // ---- Role: empty ---- + if (warp >= 9 && warp <= 11) { + // idle — no tasks assigned + } + // ---- Role: load_warp ---- + if (warp >= 12 && warp <= 15) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 64;"); + { // load_warp_main + const int load_warp_rank = warp - 12; + int work_idx_1 = blockIdx.x; + int query_head_work_2 = work_idx_1 >> 2; + int query_idx_2 = query_head_work_2 / num_head_tiles; + int head_tile_2 = query_head_work_2 % num_head_tiles; + int v_chunk_1 = work_idx_1 & 3; + int* row_ptr = swa_indices + (query_idx_2 * swa_index_stride); + int index_offset = load_warp_rank * 32 + lane; + int sparse_row = -1; + if (index_offset < sparse_topk) { + sparse_row = row_ptr[index_offset]; + } + if (load_warp_rank == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(q_full_addr, 65536); + #pragma unroll + for (int q_stage = 0; q_stage < 8; q_stage++) { + tma_4d_gmem2smem(smem_q_addr + (unsigned int)(q_stage * 8192), (&tmap_q), 0, head_tile_2 * 64, q_stage, query_idx_2, q_full_addr); + } + } + } + smem_indices[index_offset] = sparse_row; + unsigned int _vote_0 = __ballot_sync(0xFFFFFFFF, sparse_row < 0); + unsigned int invalid_indices = _vote_0; + int valid_prefix = 32; + if (invalid_indices != 0) { + int _ffs_0 = __ffs(invalid_indices); + valid_prefix = _ffs_0 - 1; + } + if (elect_sync()) { + smem_index_flags[load_warp_rank] = valid_prefix; + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("barrier.sync 8, 128;" ::: "memory"); + if ((load_warp_rank & 1) != 0) { + if (elect_sync()) { + int lower_prefix = smem_index_flags[load_warp_rank - 1]; + int upper_prefix = smem_index_flags[load_warp_rank]; + smem_index_flags[load_warp_rank] = ((lower_prefix < 32) ? lower_prefix : 32 + upper_prefix); + } + } + unsigned int _phase_kv_empty = 1; + #pragma unroll + for (int k_stage_1 = 0; k_stage_1 < 4; k_stage_1++) { + mbarrier_wait(kv_empty_addr + (k_stage_1) * 8, _phase_kv_empty); + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (k_stage_1) * 8, 8192); + } + #pragma unroll + for (int local_group = 0; local_group < 8; local_group++) { + int dst_k = smem_kv_addr + (unsigned int)(k_stage_1 * 32768); + int group = load_warp_rank * 8 + local_group; + int group_offset = group * 4; + int raw_rows[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows[0])), "=r"(*reinterpret_cast(&raw_rows[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows[(0) + 3])) + : "r"(smem_indices_addr + (unsigned int)(group_offset * 4))); + int raw0 = raw_rows[0]; + int raw1 = raw_rows[1]; + int raw2 = raw_rows[2]; + int raw3 = raw_rows[3]; + int row0 = ((raw0 >= 0) ? raw0 : 0); + int row1 = ((raw1 >= 0) ? raw1 : 0); + int row2 = ((raw2 >= 0) ? raw2 : 0); + int row3 = ((raw3 >= 0) ? raw3 : 0); + if (elect_sync()) { + tma_gather4_gmem2smem(dst_k + group * 512, (&tmap_swa_kv), k_stage_1 * 128, row0, row1, row2, row3, kv_full_addr + (k_stage_1) * 8); + tma_gather4_gmem2smem(dst_k + 16384 + group * 512, (&tmap_swa_kv), k_stage_1 * 128 + 64, row0, row1, row2, row3, kv_full_addr + (k_stage_1) * 8); + } + } + } + #pragma unroll + for (int v_stage = 0; v_stage < 1; v_stage++) { + mbarrier_wait(kv_empty_addr + (v_stage) * 8, 0); + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (v_stage) * 8, 8192); + } + #pragma unroll + for (int local_group_1 = 0; local_group_1 < 8; local_group_1++) { + int dst_v = smem_v_addr + (unsigned int)(v_stage * 32768); + int v_col = v_chunk_1 * 128 + v_stage * 64; + int group_1 = load_warp_rank * 8 + local_group_1; + int group_offset_1 = group_1 * 4; + int raw_rows_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_1[0])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 3])) + : "r"(smem_indices_addr + (unsigned int)(group_offset_1 * 4))); + int raw0_1 = raw_rows_1[0]; + int raw1_1 = raw_rows_1[1]; + int raw2_1 = raw_rows_1[2]; + int raw3_1 = raw_rows_1[3]; + int row0_1 = ((raw0_1 >= 0) ? raw0_1 : 0); + int row1_1 = ((raw1_1 >= 0) ? raw1_1 : 0); + int row2_1 = ((raw2_1 >= 0) ? raw2_1 : 0); + int row3_1 = ((raw3_1 >= 0) ? raw3_1 : 0); + if (elect_sync()) { + tma_gather4_gmem2smem(dst_v + group_1 * 512, (&tmap_swa_kv), v_col, row0_1, row1_1, row2_1, row3_1, kv_full_addr + (v_stage) * 8); + tma_gather4_gmem2smem(dst_v + 16384 + group_1 * 512, (&tmap_swa_kv), v_col + 64, row0_1, row1_1, row2_1, row3_1, kv_full_addr + (v_stage) * 8); + } + } + } + } + } + + // Cleanup +} + +} // extern "C" diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_1e3cf4361de31c58b69b_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_1e3cf4361de31c58b69b_binding.cu new file mode 100644 index 00000000000..28868ff1e78 --- /dev/null +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_1e3cf4361de31c58b69b_binding.cu @@ -0,0 +1,436 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +// Generated by the Cake source exporter — do not edit. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_1e3cf4361de31c58b69b'. +#include +#include +#include +#include + +#include "tvm_ffi_utils.h" + +#include +#include +#include +#include +#include +#include +#include + +extern "C" __global__ void kernel_cake_dsv4_1e3cf4361de31c58b69b(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int num_query_tokens, int has_sinks); + + +namespace cake_host_shim_4f4d09577aaed6ed { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t CakeDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +// 4D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + int64_t d3 = t.size(t.ndim() - 3); + TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(d3)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 64u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_swa_kv' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_swa_kv' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_kv'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_compressed_kv' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_compressed_kv(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_compressed_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_compressed_kv' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_compressed_kv' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_compressed_kv"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_compressed_kv' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_kv'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_kv') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_compressed_kv, TensorView arg_O, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_num_query_tokens, int64_t arg_has_sinks, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckContiguous(arg_tmap_q, "tmap_q"); + CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); + CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 4, 16, 1); + CheckCudaTensor(arg_tmap_compressed_kv, "tmap_compressed_kv"); + CheckDtype(arg_tmap_compressed_kv, "tmap_compressed_kv", 4, 16, 1); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_swa_indices, "swa_indices"); + CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); + CheckContiguous(arg_swa_indices, "swa_indices"); + CheckCudaTensor(arg_compressed_indices, "compressed_indices"); + CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); + CheckContiguous(arg_compressed_indices, "compressed_indices"); + CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); + CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); + CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); + CheckContiguous(arg_bmm1_scale, "bmm1_scale"); + CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); + CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); + CheckContiguous(arg_bmm2_scale, "bmm2_scale"); + TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) + << "scalar 'num_heads' value " << arg_num_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) + << "scalar 'swa_index_stride' value " << arg_swa_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) + << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) + << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) + << "scalar 'sparse_topk' value " << arg_sparse_topk + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_query_tokens >= -2147483648LL && arg_num_query_tokens <= 2147483647LL, ValueError) + << "scalar 'num_query_tokens' value " << arg_num_query_tokens + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) + << "scalar 'has_sinks' value " << arg_has_sinks + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); + CheckSameCudaDevice(arg_tmap_compressed_kv, arg_tmap_q, "tmap_compressed_kv", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); + CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); + CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); + CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + + cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); + CUtensorMap p_tmap_compressed_kv = EncodeTma_tmap_compressed_kv(arg_tmap_compressed_kv); + __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); + int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); + int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); + int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); + float* p_sinks = static_cast(arg_sinks.data_ptr()); + float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); + float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); + int32_t v_num_heads = (int32_t)arg_num_heads; + int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; + int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; + int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; + int32_t v_sparse_topk = (int32_t)arg_sparse_topk; + int32_t v_num_query_tokens = (int32_t)arg_num_query_tokens; + int32_t v_has_sinks = (int32_t)arg_has_sinks; + void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_tmap_compressed_kv, &p_O, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_num_query_tokens, &v_has_sinks}; + + static std::mutex smem_mu; + static auto* smem_status_by_device = + new std::unordered_map(); + cudaError_t smem_status = cudaSuccess; + { + std::lock_guard lock(smem_mu); + auto it = smem_status_by_device->find(dev.device_id); + if (it == smem_status_by_device->end()) { + int current_device = -1; + smem_status = cudaGetDevice(¤t_device); + if (smem_status == cudaSuccess && current_device != dev.device_id) { + smem_status = cudaErrorInvalidDevice; + } + if (smem_status == cudaSuccess) { + smem_status = cudaFuncSetAttribute( + reinterpret_cast(kernel_cake_dsv4_1e3cf4361de31c58b69b), + cudaFuncAttributeMaxDynamicSharedMemorySize, + 203904); + } + smem_status_by_device->emplace(dev.device_id, smem_status); + } else { + smem_status = it->second; + } + } + TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) + << "cudaFuncSetAttribute for kernel_cake_dsv4_1e3cf4361de31c58b69b failed: " + << cudaGetErrorString(smem_status); + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(512u, 1u, 1u); + + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes = 203904u; + config.stream = stream; + cudaError_t launch_status = cudaLaunchKernelExC( + &config, reinterpret_cast(kernel_cake_dsv4_1e3cf4361de31c58b69b), kargs); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernelExC for kernel_cake_dsv4_1e3cf4361de31c58b69b failed: " + << cudaGetErrorString(launch_status); +} + +} // namespace cake_host_shim_4f4d09577aaed6ed + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_4f4d09577aaed6ed::Run); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_1e3cf4361de31c58b69b_kernel.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_1e3cf4361de31c58b69b_kernel.cu new file mode 100644 index 00000000000..8aebdd9da7c --- /dev/null +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_1e3cf4361de31c58b69b_kernel.cu @@ -0,0 +1,2556 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +#define CAKE_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_OFFSET 0 +#define NUM_KV_PIPE_STAGES 4 +#define NUM_PAGE_PIPE_STAGES 6 +#define NUM_WORK_PIPE_STAGES 2 +#define NUM_THROTTLE_PIPE_STAGES 2 +#define SMEM_SMEM_Q_OFF 1024 +#define SMEM_SMEM_Q_STAGE_BYTES 16384 +#define SMEM_SMEM_Q_STRIDE 16384 +#define SMEM_SMEM_KV_OFF 66560 +#define SMEM_SMEM_KV_STAGE_BYTES 32768 +#define SMEM_SMEM_KV_STRIDE 32768 +#define SMEM_SMEM_V_OFF 66560 +#define SMEM_SMEM_V_STAGE_BYTES 32768 +#define SMEM_SMEM_V_STRIDE 32768 +#define SMEM_SMEM_PAGE_OFFSETS_OFF 197632 +#define SMEM_SMEM_PAGE_OFFSETS_STAGE_BYTES 1024 +#define SMEM_SMEM_PAGE_OFFSETS_STRIDE 1024 +#define SMEM_WORK_RESPONSE_OFF 203776 +#define SMEM_WORK_RESPONSE_STAGE_BYTES 16 +#define SMEM_WORK_RESPONSE_STRIDE 16 +#define SMEM_TOTAL 203904 +#define THREADS 512 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count) : "memory"); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ float2 ex2_emulation_f32x2_value(float2 value) { + const float c0 = 1.0f, c1 = 0.695146143436431884765625f; + const float c2 = 0.227564394474029541015625f, c3 = 0.077119089663028717041015625f; + const float magic = 12582912.0f; + float x0 = max_noftz(value.x, -127.0f), x1 = max_noftz(value.y, -127.0f); + float2 xc2 = make_float2(x0, x1), magic2 = make_float2(magic, magic); + float2 xr2; + asm("add.rm.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xr2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&magic2)); + float2 c3_2 = make_float2(c3, c3), c2_2 = make_float2(c2, c2); + float2 c1_2 = make_float2(c1, c1), c0_2 = make_float2(c0, c0); + float2 xrb2, xfrac2; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xrb2) + : "l"(*(unsigned long long*)&xr2), "l"(*(unsigned long long*)&magic2)); + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xfrac2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&xrb2)); + float2 poly2; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&c3_2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c2_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c1_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c0_2)); + int x0r_i, x1r_i, p0_i, p1_i; + asm("mov.b64 {%0, %1}, %2;" : "=r"(x0r_i), "=r"(x1r_i) : "l"(*(unsigned long long*)&xr2)); + asm("mov.b64 {%0, %1}, %2;" : "=r"(p0_i), "=r"(p1_i) : "l"(*(unsigned long long*)&poly2)); + float r0, r1; + asm("mov.b32 %0, %1;" : "=f"(r0) : "r"((x0r_i << 23) + p0_i)); + asm("mov.b32 %0, %1;" : "=f"(r1) : "r"((x1r_i << 23) + p1_i)); + return make_float2(r0, r1); +} + +__device__ __forceinline__ void ex2_emulation_f32x2(float* x0_ptr, float* x1_ptr) { + float2 result = ex2_emulation_f32x2_value(make_float2(*x0_ptr, *x1_ptr)); + *x0_ptr = result.x; *x1_ptr = result.y; +} + +__device__ __forceinline__ void softmax_frag_exp2_cast( + float* sv, uint32_t* pv, int use_emu) +{ + #pragma unroll + for (int j = 0; j < 16; j++) { + if (use_emu && j >= 12) + ex2_emulation_f32x2(&sv[j*2], &sv[j*2+1]); + else { + sv[j*2] = approx_exp2(sv[j*2]); + sv[j*2+1] = approx_exp2(sv[j*2+1]); + } + } + #pragma unroll + for (int j = 0; j < 16; j++) { + __nv_bfloat162 bf = __float22bfloat162_rn({sv[j*2], sv[j*2+1]}); + pv[j] = reinterpret_cast(bf); + } +} + + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_2d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3}], [%4];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr) { + // Canonical .shared::cta form for non-multicast gather4, matching + // trtllm-gen / cuda_ptx and the PTX ISA qualifier order + // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, + // .shared::cluster is reserved for the multicast variant (ctaMask). + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(512, 1) void +kernel_cake_dsv4_1e3cf4361de31c58b69b(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int num_query_tokens, int has_sinks) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int mbar_base = smem; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define kv_full_addr (mbar_base + 16) + #define kv_empty_addr (mbar_base + 48) + #define s_full_addr (mbar_base + 80) + #define p_full_addr (mbar_base + 96) + #define stats_addr (mbar_base + 112) + #define stats_empty_addr (mbar_base + 128) + #define o_empty_addr (mbar_base + 144) + #define o_full_addr (mbar_base + 152) + #define page_full_addr (mbar_base + 160) + #define page_empty_addr (mbar_base + 208) + #define work_full_addr (mbar_base + 256) + #define work_empty_addr (mbar_base + 272) + #define throttle_full_addr (mbar_base + 288) + #define throttle_empty_addr (mbar_base + 304) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + + // Kernel setup ops + __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_q_addr = smem + 1024; + __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_kv_addr = smem + 66560; + __nv_bfloat16* smem_v = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_v_addr = smem + 66560; + int* smem_page_offsets = reinterpret_cast(smem_raw + 197632); + const int smem_page_offsets_addr = smem + 197632; + unsigned int* work_response = reinterpret_cast(smem_raw + 203776); + const int work_response_addr = smem + 203776; + + // Mbarrier init (16 pipeline groups, 0 ordered-sequence groups, 40 barriers) + // Mbarriers at smem_raw[0..320) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // q_full: 1 barriers, init_count=1 + mbarrier_init(smem + 0, 1); + // q_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 8, 1); + // --- pipeline 'kv_pipe' --- + // kv_full: 4 barriers, init_count=1 + mbarrier_init(smem + 16, 1); + mbarrier_init(smem + 24, 1); + mbarrier_init(smem + 32, 1); + mbarrier_init(smem + 40, 1); + // kv_empty: 4 barriers, init_count=1 + mbarrier_init(smem + 48, 1); + mbarrier_init(smem + 56, 1); + mbarrier_init(smem + 64, 1); + mbarrier_init(smem + 72, 1); + // s_full: 2 barriers, init_count=1 + mbarrier_init(smem + 80, 1); + mbarrier_init(smem + 88, 1); + // p_full: 2 barriers, init_count=128 + mbarrier_init(smem + 96, 128); + mbarrier_init(smem + 104, 128); + // stats: 2 barriers, init_count=128 + mbarrier_init(smem + 112, 128); + mbarrier_init(smem + 120, 128); + // stats_empty: 2 barriers, init_count=128 + mbarrier_init(smem + 128, 128); + mbarrier_init(smem + 136, 128); + // o_empty: 1 barriers, init_count=128 + mbarrier_init(smem + 144, 128); + // o_full: 1 barriers, init_count=1 + mbarrier_init(smem + 152, 1); + // --- pipeline 'page_pipe' --- + // page_full: 6 barriers, init_count=32 + mbarrier_init(smem + 160, 32); + mbarrier_init(smem + 168, 32); + mbarrier_init(smem + 176, 32); + mbarrier_init(smem + 184, 32); + mbarrier_init(smem + 192, 32); + mbarrier_init(smem + 200, 32); + // page_empty: 6 barriers, init_count=128 + mbarrier_init(smem + 208, 128); + mbarrier_init(smem + 216, 128); + mbarrier_init(smem + 224, 128); + mbarrier_init(smem + 232, 128); + mbarrier_init(smem + 240, 128); + mbarrier_init(smem + 248, 128); + // --- pipeline 'work_pipe' --- + // work_full: 2 barriers, init_count=1 + mbarrier_init(smem + 256, 1); + mbarrier_init(smem + 264, 1); + // work_empty: 2 barriers, init_count=512 + mbarrier_init(smem + 272, 512); + mbarrier_init(smem + 280, 512); + // --- pipeline 'throttle_pipe' --- + // throttle_full: 2 barriers, init_count=128 + mbarrier_init(smem + 288, 128); + mbarrier_init(smem + 296, 128); + // throttle_empty: 2 barriers, init_count=32 + mbarrier_init(smem + 304, 32); + mbarrier_init(smem + 312, 32); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncwarp(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 320); + if (warp == 0) { + int _tmem_hold = smem + 320; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_tmem = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Inc phase consumes the registers released above. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 136;"); + } + + // ---- Role: softmax ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 152;"); + { // softmax_main + asm volatile("griddepcontrol.wait;" ::: "memory"); + float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; + unsigned int work_stage = 0; + unsigned int query_idx = blockIdx.x; + unsigned int global_tile = 0; + unsigned int stats_cursor = 0; + const int warp_in_compute = warp; + const int tmem_row_origin = warp_in_compute * 32; + const int logical_row_origin = warp_in_compute * 16; + const int my_row = logical_row_origin + lane % 16; + const int col_half = lane / 16; + unsigned int _phase_work_full = 0; + #pragma unroll 1 + for (unsigned int work_iter = 0; work_iter < num_query_tokens; work_iter++) { + if (query_idx >= (unsigned int)num_query_tokens) { + break; + } + int _max_4 = ((sparse_topk_lens[query_idx] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx] + sparse_topk_lens_offset) : (0)); + int _min_2 = ((_max_4) < (sparse_topk) ? (_max_4) : (sparse_topk)); + int active_topk = _min_2; + int _max_5 = (((active_topk + 128 - 1) / 128) > (1) ? ((active_topk + 128 - 1) / 128) : (1)); + int num_loop_steps = _max_5; + float row_max = -CAKE_INF; + float row_sum = 0.0f; + if (has_sinks != 0 && my_row < num_heads) { + row_max = sinks[my_row] * 1.4426950408889634f / softmax_scale_log2; + row_sum = 1.0f; + } + if (num_loop_steps > 0) { + unsigned int first_stats_stage = stats_cursor & 1; + unsigned int first_stats_empty_phase = stats_cursor >> 1 & 1 ^ 1; + mbarrier_wait(stats_empty_addr + (first_stats_stage) * 8, first_stats_empty_phase); + } + #pragma unroll 1 + for (int tile = 0; tile < num_loop_steps; tile++) { + unsigned int sync_stage = global_tile & 1; + unsigned int sync_phase = global_tile >> 1 & 1; + mbarrier_wait(s_full_addr + (sync_stage) * 8, sync_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int score_col = ((sync_stage != 0) ? 128 : 0); + int score_addr = taddr + (unsigned int)score_col + (unsigned int)(tmem_row_origin << 16); + float _tmem_load_0[64]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_0[0])), "=r"(*reinterpret_cast(&_tmem_load_0[1])), "=r"(*reinterpret_cast(&_tmem_load_0[2])), "=r"(*reinterpret_cast(&_tmem_load_0[3])), "=r"(*reinterpret_cast(&_tmem_load_0[4])), "=r"(*reinterpret_cast(&_tmem_load_0[5])), "=r"(*reinterpret_cast(&_tmem_load_0[6])), "=r"(*reinterpret_cast(&_tmem_load_0[7])), "=r"(*reinterpret_cast(&_tmem_load_0[8])), "=r"(*reinterpret_cast(&_tmem_load_0[9])), "=r"(*reinterpret_cast(&_tmem_load_0[10])), "=r"(*reinterpret_cast(&_tmem_load_0[11])), "=r"(*reinterpret_cast(&_tmem_load_0[12])), "=r"(*reinterpret_cast(&_tmem_load_0[13])), "=r"(*reinterpret_cast(&_tmem_load_0[14])), "=r"(*reinterpret_cast(&_tmem_load_0[15])), "=r"(*reinterpret_cast(&_tmem_load_0[16])), "=r"(*reinterpret_cast(&_tmem_load_0[17])), "=r"(*reinterpret_cast(&_tmem_load_0[18])), "=r"(*reinterpret_cast(&_tmem_load_0[19])), "=r"(*reinterpret_cast(&_tmem_load_0[20])), "=r"(*reinterpret_cast(&_tmem_load_0[21])), "=r"(*reinterpret_cast(&_tmem_load_0[22])), "=r"(*reinterpret_cast(&_tmem_load_0[23])), "=r"(*reinterpret_cast(&_tmem_load_0[24])), "=r"(*reinterpret_cast(&_tmem_load_0[25])), "=r"(*reinterpret_cast(&_tmem_load_0[26])), "=r"(*reinterpret_cast(&_tmem_load_0[27])), "=r"(*reinterpret_cast(&_tmem_load_0[28])), "=r"(*reinterpret_cast(&_tmem_load_0[29])), "=r"(*reinterpret_cast(&_tmem_load_0[30])), "=r"(*reinterpret_cast(&_tmem_load_0[31])) + : "r"(score_addr)); + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_0[32])), "=r"(*reinterpret_cast(&_tmem_load_0[33])), "=r"(*reinterpret_cast(&_tmem_load_0[34])), "=r"(*reinterpret_cast(&_tmem_load_0[35])), "=r"(*reinterpret_cast(&_tmem_load_0[36])), "=r"(*reinterpret_cast(&_tmem_load_0[37])), "=r"(*reinterpret_cast(&_tmem_load_0[38])), "=r"(*reinterpret_cast(&_tmem_load_0[39])), "=r"(*reinterpret_cast(&_tmem_load_0[40])), "=r"(*reinterpret_cast(&_tmem_load_0[41])), "=r"(*reinterpret_cast(&_tmem_load_0[42])), "=r"(*reinterpret_cast(&_tmem_load_0[43])), "=r"(*reinterpret_cast(&_tmem_load_0[44])), "=r"(*reinterpret_cast(&_tmem_load_0[45])), "=r"(*reinterpret_cast(&_tmem_load_0[46])), "=r"(*reinterpret_cast(&_tmem_load_0[47])), "=r"(*reinterpret_cast(&_tmem_load_0[48])), "=r"(*reinterpret_cast(&_tmem_load_0[49])), "=r"(*reinterpret_cast(&_tmem_load_0[50])), "=r"(*reinterpret_cast(&_tmem_load_0[51])), "=r"(*reinterpret_cast(&_tmem_load_0[52])), "=r"(*reinterpret_cast(&_tmem_load_0[53])), "=r"(*reinterpret_cast(&_tmem_load_0[54])), "=r"(*reinterpret_cast(&_tmem_load_0[55])), "=r"(*reinterpret_cast(&_tmem_load_0[56])), "=r"(*reinterpret_cast(&_tmem_load_0[57])), "=r"(*reinterpret_cast(&_tmem_load_0[58])), "=r"(*reinterpret_cast(&_tmem_load_0[59])), "=r"(*reinterpret_cast(&_tmem_load_0[60])), "=r"(*reinterpret_cast(&_tmem_load_0[61])), "=r"(*reinterpret_cast(&_tmem_load_0[62])), "=r"(*reinterpret_cast(&_tmem_load_0[63])) + : "r"(score_addr + 32)); + int valid_cols = active_topk - tile * 128 - col_half * 64; + if (valid_cols < 0) { + valid_cols = 0; + } + if (valid_cols > 64) { + valid_cols = 64; + } + uint32_t _slice_lo_mask_0; + { + int _lim_0 = valid_cols; + if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_0 & (1u << 0))) _tmem_load_0[0] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 1))) _tmem_load_0[1] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 2))) _tmem_load_0[2] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 3))) _tmem_load_0[3] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 4))) _tmem_load_0[4] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 5))) _tmem_load_0[5] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 6))) _tmem_load_0[6] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 7))) _tmem_load_0[7] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 8))) _tmem_load_0[8] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 9))) _tmem_load_0[9] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 10))) _tmem_load_0[10] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 11))) _tmem_load_0[11] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 12))) _tmem_load_0[12] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 13))) _tmem_load_0[13] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 14))) _tmem_load_0[14] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 15))) _tmem_load_0[15] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 16))) _tmem_load_0[16] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 17))) _tmem_load_0[17] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 18))) _tmem_load_0[18] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 19))) _tmem_load_0[19] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 20))) _tmem_load_0[20] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 21))) _tmem_load_0[21] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 22))) _tmem_load_0[22] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 23))) _tmem_load_0[23] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 24))) _tmem_load_0[24] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 25))) _tmem_load_0[25] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 26))) _tmem_load_0[26] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 27))) _tmem_load_0[27] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 28))) _tmem_load_0[28] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 29))) _tmem_load_0[29] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 30))) _tmem_load_0[30] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 31))) _tmem_load_0[31] = -CAKE_INF; + uint32_t _slice_lo_mask_1; + { + int _lim_1 = valid_cols - 32; + if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) _tmem_load_0[32] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) _tmem_load_0[33] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) _tmem_load_0[34] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) _tmem_load_0[35] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) _tmem_load_0[36] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) _tmem_load_0[37] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) _tmem_load_0[38] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) _tmem_load_0[39] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) _tmem_load_0[40] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) _tmem_load_0[41] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) _tmem_load_0[42] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) _tmem_load_0[43] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) _tmem_load_0[44] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) _tmem_load_0[45] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) _tmem_load_0[46] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) _tmem_load_0[47] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) _tmem_load_0[48] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) _tmem_load_0[49] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) _tmem_load_0[50] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) _tmem_load_0[51] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) _tmem_load_0[52] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) _tmem_load_0[53] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) _tmem_load_0[54] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) _tmem_load_0[55] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) _tmem_load_0[56] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) _tmem_load_0[57] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) _tmem_load_0[58] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) _tmem_load_0[59] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) _tmem_load_0[60] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) _tmem_load_0[61] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) _tmem_load_0[62] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) _tmem_load_0[63] = -CAKE_INF; + float2 _reg_reduce_max2_2 = {-CAKE_INF, -CAKE_INF}; + row_max_x32_accum(&_tmem_load_0[0], _reg_reduce_max2_2); + row_max_x32_accum(&_tmem_load_0[32], _reg_reduce_max2_2); + float _tmem_load_0_max = row_max_reduce(_reg_reduce_max2_2); + float tile_max = _tmem_load_0_max; + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, tile_max, 16); + float _max_6 = max_noftz(tile_max, _shfl_xor_0); + tile_max = _max_6; + float old_max = row_max; + float _max_7 = max_noftz(row_max, tile_max); + float new_max = _max_7; + float _fma_0 = __fmaf_rn(old_max, softmax_scale_log2, (-new_max) * softmax_scale_log2); + float delta = _fma_0; + float _exp2_0 = approx_exp2(delta); + float acc_scale = ((old_max > -CAKE_INF) ? _exp2_0 : 1.0f); + unsigned int stats_stage = stats_cursor & 1; + int stats_col = ((stats_stage != 0) ? 128 : 0); + float max_stats[2]; + max_stats[0] = old_max; + max_stats[1] = new_max; + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(taddr + (unsigned int)stats_col), "f"(max_stats[0]), "f"(max_stats[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (stats_stage) * 8); + stats_cursor = stats_cursor + 1; + unsigned int next_stats_stage = stats_cursor & 1; + unsigned int next_stats_empty_phase = stats_cursor >> 1 & 1 ^ 1; + mbarrier_wait(stats_empty_addr + (next_stats_stage) * 8, next_stats_empty_phase); + bool empty_row = new_max == -CAKE_INF; + float safe_max = ((empty_row) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + const float2 _fma_b2_3 = {softmax_scale_log2, softmax_scale_log2}; + const float2 _fma_c2_4 = {-max_scaled, -max_scaled}; + #pragma unroll + for (int _lf = 0; _lf < 32; _lf++) + fma_f32x2_inplace(&reinterpret_cast(_tmem_load_0)[_lf], _fma_b2_3, _fma_c2_4); + #pragma unroll + for (int _le = 0; _le < 64; _le++) { + _tmem_load_0[_le] = approx_exp2(_tmem_load_0[_le]); + } + float2 _reg_reduce_sum2_5 = make_float2(0.0f, 0.0f); + softmax_block_sum(&_tmem_load_0[0], &_reg_reduce_sum2_5); + softmax_block_sum(&_tmem_load_0[32], &_reg_reduce_sum2_5); + float _tmem_load_0_sum = _reg_reduce_sum2_5.x + _reg_reduce_sum2_5.y; + float block_sum = _tmem_load_0_sum; + float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, block_sum, 16); + block_sum = block_sum + _shfl_xor_1; + float _fma_1 = __fmaf_rn(row_sum, acc_scale, block_sum); + row_sum = _fma_1; + row_max = new_max; + if (tile == num_loop_steps - 1) { + float final_stats[2]; + float _max_8 = max_noftz(row_sum, 1.1754943508222875e-38f); + final_stats[0] = _max_8; + final_stats[1] = row_max; + int final_stats_col = ((next_stats_stage != 0) ? 128 : 0); + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(taddr + (unsigned int)final_stats_col), "f"(final_stats[0]), "f"(final_stats[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (next_stats_stage) * 8); + stats_cursor = stats_cursor + 1; + unsigned int tail_stats_stage = stats_cursor & 1; + unsigned int tail_stats_empty_phase = stats_cursor >> 1 & 1 ^ 1; + mbarrier_wait(stats_empty_addr + (tail_stats_stage) * 8, tail_stats_empty_phase); + mbarrier_arrive(stats_addr + (tail_stats_stage) * 8); + stats_cursor = stats_cursor + 1; + } + unsigned int packed_p[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_0[_lp*2 + 0], _tmem_load_0[_lp*2+1 + 0])); + packed_p[_lp] = *(uint32_t*)&_bf2; + } + int p_addr = taddr + (unsigned int)score_col + 32 + (unsigned int)(tmem_row_origin << 16); + asm volatile( + "tcgen05.st.sync.aligned.16x32bx2.x16.b32" + " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(p_addr), "r"(*reinterpret_cast(&packed_p[0])), "r"(*reinterpret_cast(&packed_p[1])), "r"(*reinterpret_cast(&packed_p[2])), "r"(*reinterpret_cast(&packed_p[3])), "r"(*reinterpret_cast(&packed_p[4])), "r"(*reinterpret_cast(&packed_p[5])), "r"(*reinterpret_cast(&packed_p[6])), "r"(*reinterpret_cast(&packed_p[7])), "r"(*reinterpret_cast(&packed_p[8])), "r"(*reinterpret_cast(&packed_p[9])), "r"(*reinterpret_cast(&packed_p[10])), "r"(*reinterpret_cast(&packed_p[11])), "r"(*reinterpret_cast(&packed_p[12])), "r"(*reinterpret_cast(&packed_p[13])), "r"(*reinterpret_cast(&packed_p[14])), "r"(*reinterpret_cast(&packed_p[15]))); + asm volatile( + "tcgen05.st.sync.aligned.16x32bx2.x16.b32" + " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(p_addr + 16), "r"(*reinterpret_cast(&(packed_p + 16)[0])), "r"(*reinterpret_cast(&(packed_p + 16)[1])), "r"(*reinterpret_cast(&(packed_p + 16)[2])), "r"(*reinterpret_cast(&(packed_p + 16)[3])), "r"(*reinterpret_cast(&(packed_p + 16)[4])), "r"(*reinterpret_cast(&(packed_p + 16)[5])), "r"(*reinterpret_cast(&(packed_p + 16)[6])), "r"(*reinterpret_cast(&(packed_p + 16)[7])), "r"(*reinterpret_cast(&(packed_p + 16)[8])), "r"(*reinterpret_cast(&(packed_p + 16)[9])), "r"(*reinterpret_cast(&(packed_p + 16)[10])), "r"(*reinterpret_cast(&(packed_p + 16)[11])), "r"(*reinterpret_cast(&(packed_p + 16)[12])), "r"(*reinterpret_cast(&(packed_p + 16)[13])), "r"(*reinterpret_cast(&(packed_p + 16)[14])), "r"(*reinterpret_cast(&(packed_p + 16)[15]))); + int p_addr_high = p_addr + 1048576; + asm volatile( + "tcgen05.st.sync.aligned.16x32bx2.x16.b32" + " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(p_addr_high), "r"(*reinterpret_cast(&packed_p[0])), "r"(*reinterpret_cast(&packed_p[1])), "r"(*reinterpret_cast(&packed_p[2])), "r"(*reinterpret_cast(&packed_p[3])), "r"(*reinterpret_cast(&packed_p[4])), "r"(*reinterpret_cast(&packed_p[5])), "r"(*reinterpret_cast(&packed_p[6])), "r"(*reinterpret_cast(&packed_p[7])), "r"(*reinterpret_cast(&packed_p[8])), "r"(*reinterpret_cast(&packed_p[9])), "r"(*reinterpret_cast(&packed_p[10])), "r"(*reinterpret_cast(&packed_p[11])), "r"(*reinterpret_cast(&packed_p[12])), "r"(*reinterpret_cast(&packed_p[13])), "r"(*reinterpret_cast(&packed_p[14])), "r"(*reinterpret_cast(&packed_p[15]))); + asm volatile( + "tcgen05.st.sync.aligned.16x32bx2.x16.b32" + " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(p_addr_high + 16), "r"(*reinterpret_cast(&(packed_p + 16)[0])), "r"(*reinterpret_cast(&(packed_p + 16)[1])), "r"(*reinterpret_cast(&(packed_p + 16)[2])), "r"(*reinterpret_cast(&(packed_p + 16)[3])), "r"(*reinterpret_cast(&(packed_p + 16)[4])), "r"(*reinterpret_cast(&(packed_p + 16)[5])), "r"(*reinterpret_cast(&(packed_p + 16)[6])), "r"(*reinterpret_cast(&(packed_p + 16)[7])), "r"(*reinterpret_cast(&(packed_p + 16)[8])), "r"(*reinterpret_cast(&(packed_p + 16)[9])), "r"(*reinterpret_cast(&(packed_p + 16)[10])), "r"(*reinterpret_cast(&(packed_p + 16)[11])), "r"(*reinterpret_cast(&(packed_p + 16)[12])), "r"(*reinterpret_cast(&(packed_p + 16)[13])), "r"(*reinterpret_cast(&(packed_p + 16)[14])), "r"(*reinterpret_cast(&(packed_p + 16)[15]))); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(p_full_addr + (sync_stage) * 8); + global_tile = global_tile + 1; + } + mbarrier_wait(work_full_addr + (work_stage) * 8, _phase_work_full); + uint32_t _clc_valid_2 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_2) + : "r"(work_response_addr + work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_2 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_2) + : "r"(work_response_addr + work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(work_empty_addr + (work_stage) * 8); + work_stage += 1; + if (work_stage == 2) { work_stage = 0; _phase_work_full ^= 1; } + unsigned int valid = _clc_valid_2; + query_idx = _clc_ctaid_2; + if (valid == 0) { + break; + } + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + } + } + // ---- Role: correction ---- + if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 144;"); + { // correction_main + asm volatile("griddepcontrol.wait;" ::: "memory"); + float softmax_scale_log2_1 = bmm1_scale[0] * 1.4426950408889634f; + float output_scale = bmm2_scale[0]; + unsigned int work_stage_1 = 0; + unsigned int query_idx_1 = blockIdx.x; + unsigned int global_tile_1 = 0; + unsigned int stats_cursor_1 = 0; + const int warp_in_role = warp - 4; + const int tmem_row_origin_1 = warp_in_role * 32; + const int logical_row_origin_1 = warp_in_role * 16; + const int my_row_1 = logical_row_origin_1 + lane % 16; + const int col_half_1 = lane / 16; + unsigned int _phase_work_full_1 = 0; + #pragma unroll 1 + for (unsigned int work_iter_1 = 0; work_iter_1 < num_query_tokens; work_iter_1++) { + if (query_idx_1 >= (unsigned int)num_query_tokens) { + break; + } + int _max_9 = ((sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) : (0)); + int _min_3 = ((_max_9) < (sparse_topk) ? (_max_9) : (sparse_topk)); + int active_topk_1 = _min_3; + int _max_10 = (((active_topk_1 + 128 - 1) / 128) > (1) ? ((active_topk_1 + 128 - 1) / 128) : (1)); + int num_loop_steps_1 = _max_10; + unsigned int first_stats_stage_1 = stats_cursor_1 & 1; + unsigned int first_stats_phase = stats_cursor_1 >> 1 & 1; + mbarrier_wait(stats_addr + (first_stats_stage_1) * 8, first_stats_phase); + mbarrier_arrive(stats_empty_addr + (first_stats_stage_1) * 8); + stats_cursor_1 = stats_cursor_1 + 1; + global_tile_1 = global_tile_1 + 1; + #pragma unroll 1 + for (int tile_1 = 0; tile_1 < num_loop_steps_1 - 1; tile_1++) { + unsigned int stats_stage_1 = stats_cursor_1 & 1; + unsigned int stats_phase = stats_cursor_1 >> 1 & 1; + mbarrier_wait(stats_addr + (stats_stage_1) * 8, stats_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int stats_col_1 = ((stats_stage_1 != 0) ? 128 : 0); + float _tmem_load_1[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]) + : "r"(taddr + (unsigned int)stats_col_1)); + mbarrier_arrive(stats_empty_addr + (stats_stage_1) * 8); + stats_cursor_1 = stats_cursor_1 + 1; + float max_diff = _tmem_load_1[0] - _tmem_load_1[1]; + float _exp2_1 = approx_exp2(softmax_scale_log2_1 * max_diff); + float acc_scale_1 = ((max_diff != 0.0f) ? _exp2_1 : 1.0f); + mbarrier_wait(o_full_addr, global_tile_1 - 1 & 1); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _vote_0 = __any_sync(0xFFFFFFFF, acc_scale_1 < 1.0f); + int any_rescale = _vote_0; + if (any_rescale != 0) { + #pragma unroll + for (int v_stage = 0; v_stage < 4; v_stage++) { + int output_addr = taddr + 256 + (unsigned int)((v_stage & 1) * 128) + (unsigned int)((v_stage >> 1) * 16 << 16) + (unsigned int)(tmem_row_origin_1 << 16); + float _tmem_load_2[64]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_2[0])), "=r"(*reinterpret_cast(&_tmem_load_2[1])), "=r"(*reinterpret_cast(&_tmem_load_2[2])), "=r"(*reinterpret_cast(&_tmem_load_2[3])), "=r"(*reinterpret_cast(&_tmem_load_2[4])), "=r"(*reinterpret_cast(&_tmem_load_2[5])), "=r"(*reinterpret_cast(&_tmem_load_2[6])), "=r"(*reinterpret_cast(&_tmem_load_2[7])), "=r"(*reinterpret_cast(&_tmem_load_2[8])), "=r"(*reinterpret_cast(&_tmem_load_2[9])), "=r"(*reinterpret_cast(&_tmem_load_2[10])), "=r"(*reinterpret_cast(&_tmem_load_2[11])), "=r"(*reinterpret_cast(&_tmem_load_2[12])), "=r"(*reinterpret_cast(&_tmem_load_2[13])), "=r"(*reinterpret_cast(&_tmem_load_2[14])), "=r"(*reinterpret_cast(&_tmem_load_2[15])), "=r"(*reinterpret_cast(&_tmem_load_2[16])), "=r"(*reinterpret_cast(&_tmem_load_2[17])), "=r"(*reinterpret_cast(&_tmem_load_2[18])), "=r"(*reinterpret_cast(&_tmem_load_2[19])), "=r"(*reinterpret_cast(&_tmem_load_2[20])), "=r"(*reinterpret_cast(&_tmem_load_2[21])), "=r"(*reinterpret_cast(&_tmem_load_2[22])), "=r"(*reinterpret_cast(&_tmem_load_2[23])), "=r"(*reinterpret_cast(&_tmem_load_2[24])), "=r"(*reinterpret_cast(&_tmem_load_2[25])), "=r"(*reinterpret_cast(&_tmem_load_2[26])), "=r"(*reinterpret_cast(&_tmem_load_2[27])), "=r"(*reinterpret_cast(&_tmem_load_2[28])), "=r"(*reinterpret_cast(&_tmem_load_2[29])), "=r"(*reinterpret_cast(&_tmem_load_2[30])), "=r"(*reinterpret_cast(&_tmem_load_2[31])) + : "r"(output_addr)); + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_2[32])), "=r"(*reinterpret_cast(&_tmem_load_2[33])), "=r"(*reinterpret_cast(&_tmem_load_2[34])), "=r"(*reinterpret_cast(&_tmem_load_2[35])), "=r"(*reinterpret_cast(&_tmem_load_2[36])), "=r"(*reinterpret_cast(&_tmem_load_2[37])), "=r"(*reinterpret_cast(&_tmem_load_2[38])), "=r"(*reinterpret_cast(&_tmem_load_2[39])), "=r"(*reinterpret_cast(&_tmem_load_2[40])), "=r"(*reinterpret_cast(&_tmem_load_2[41])), "=r"(*reinterpret_cast(&_tmem_load_2[42])), "=r"(*reinterpret_cast(&_tmem_load_2[43])), "=r"(*reinterpret_cast(&_tmem_load_2[44])), "=r"(*reinterpret_cast(&_tmem_load_2[45])), "=r"(*reinterpret_cast(&_tmem_load_2[46])), "=r"(*reinterpret_cast(&_tmem_load_2[47])), "=r"(*reinterpret_cast(&_tmem_load_2[48])), "=r"(*reinterpret_cast(&_tmem_load_2[49])), "=r"(*reinterpret_cast(&_tmem_load_2[50])), "=r"(*reinterpret_cast(&_tmem_load_2[51])), "=r"(*reinterpret_cast(&_tmem_load_2[52])), "=r"(*reinterpret_cast(&_tmem_load_2[53])), "=r"(*reinterpret_cast(&_tmem_load_2[54])), "=r"(*reinterpret_cast(&_tmem_load_2[55])), "=r"(*reinterpret_cast(&_tmem_load_2[56])), "=r"(*reinterpret_cast(&_tmem_load_2[57])), "=r"(*reinterpret_cast(&_tmem_load_2[58])), "=r"(*reinterpret_cast(&_tmem_load_2[59])), "=r"(*reinterpret_cast(&_tmem_load_2[60])), "=r"(*reinterpret_cast(&_tmem_load_2[61])), "=r"(*reinterpret_cast(&_tmem_load_2[62])), "=r"(*reinterpret_cast(&_tmem_load_2[63])) + : "r"(output_addr + 32)); + const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 32; _ls++) + mul_f32x2_inplace(&reinterpret_cast(_tmem_load_2)[_ls], _scale2_0); + asm volatile( + "tcgen05.st.sync.aligned.16x32bx2.x64.b32" + " [%0], 64, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64};" + :: "r"(output_addr), "r"(*reinterpret_cast(&_tmem_load_2[0])), "r"(*reinterpret_cast(&_tmem_load_2[1])), "r"(*reinterpret_cast(&_tmem_load_2[2])), "r"(*reinterpret_cast(&_tmem_load_2[3])), "r"(*reinterpret_cast(&_tmem_load_2[4])), "r"(*reinterpret_cast(&_tmem_load_2[5])), "r"(*reinterpret_cast(&_tmem_load_2[6])), "r"(*reinterpret_cast(&_tmem_load_2[7])), "r"(*reinterpret_cast(&_tmem_load_2[8])), "r"(*reinterpret_cast(&_tmem_load_2[9])), "r"(*reinterpret_cast(&_tmem_load_2[10])), "r"(*reinterpret_cast(&_tmem_load_2[11])), "r"(*reinterpret_cast(&_tmem_load_2[12])), "r"(*reinterpret_cast(&_tmem_load_2[13])), "r"(*reinterpret_cast(&_tmem_load_2[14])), "r"(*reinterpret_cast(&_tmem_load_2[15])), "r"(*reinterpret_cast(&_tmem_load_2[16])), "r"(*reinterpret_cast(&_tmem_load_2[17])), "r"(*reinterpret_cast(&_tmem_load_2[18])), "r"(*reinterpret_cast(&_tmem_load_2[19])), "r"(*reinterpret_cast(&_tmem_load_2[20])), "r"(*reinterpret_cast(&_tmem_load_2[21])), "r"(*reinterpret_cast(&_tmem_load_2[22])), "r"(*reinterpret_cast(&_tmem_load_2[23])), "r"(*reinterpret_cast(&_tmem_load_2[24])), "r"(*reinterpret_cast(&_tmem_load_2[25])), "r"(*reinterpret_cast(&_tmem_load_2[26])), "r"(*reinterpret_cast(&_tmem_load_2[27])), "r"(*reinterpret_cast(&_tmem_load_2[28])), "r"(*reinterpret_cast(&_tmem_load_2[29])), "r"(*reinterpret_cast(&_tmem_load_2[30])), "r"(*reinterpret_cast(&_tmem_load_2[31])), "r"(*reinterpret_cast(&_tmem_load_2[32])), "r"(*reinterpret_cast(&_tmem_load_2[33])), "r"(*reinterpret_cast(&_tmem_load_2[34])), "r"(*reinterpret_cast(&_tmem_load_2[35])), "r"(*reinterpret_cast(&_tmem_load_2[36])), "r"(*reinterpret_cast(&_tmem_load_2[37])), "r"(*reinterpret_cast(&_tmem_load_2[38])), "r"(*reinterpret_cast(&_tmem_load_2[39])), "r"(*reinterpret_cast(&_tmem_load_2[40])), "r"(*reinterpret_cast(&_tmem_load_2[41])), "r"(*reinterpret_cast(&_tmem_load_2[42])), "r"(*reinterpret_cast(&_tmem_load_2[43])), "r"(*reinterpret_cast(&_tmem_load_2[44])), "r"(*reinterpret_cast(&_tmem_load_2[45])), "r"(*reinterpret_cast(&_tmem_load_2[46])), "r"(*reinterpret_cast(&_tmem_load_2[47])), "r"(*reinterpret_cast(&_tmem_load_2[48])), "r"(*reinterpret_cast(&_tmem_load_2[49])), "r"(*reinterpret_cast(&_tmem_load_2[50])), "r"(*reinterpret_cast(&_tmem_load_2[51])), "r"(*reinterpret_cast(&_tmem_load_2[52])), "r"(*reinterpret_cast(&_tmem_load_2[53])), "r"(*reinterpret_cast(&_tmem_load_2[54])), "r"(*reinterpret_cast(&_tmem_load_2[55])), "r"(*reinterpret_cast(&_tmem_load_2[56])), "r"(*reinterpret_cast(&_tmem_load_2[57])), "r"(*reinterpret_cast(&_tmem_load_2[58])), "r"(*reinterpret_cast(&_tmem_load_2[59])), "r"(*reinterpret_cast(&_tmem_load_2[60])), "r"(*reinterpret_cast(&_tmem_load_2[61])), "r"(*reinterpret_cast(&_tmem_load_2[62])), "r"(*reinterpret_cast(&_tmem_load_2[63]))); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + } + } + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(o_empty_addr); + global_tile_1 = global_tile_1 + 1; + } + unsigned int final_stats_stage = stats_cursor_1 & 1; + unsigned int final_stats_phase = stats_cursor_1 >> 1 & 1; + mbarrier_wait(stats_addr + (final_stats_stage) * 8, final_stats_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int final_stats_col_1 = ((final_stats_stage != 0) ? 128 : 0); + float _tmem_load_3[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]) + : "r"(taddr + (unsigned int)final_stats_col_1)); + mbarrier_arrive(stats_empty_addr + (final_stats_stage) * 8); + stats_cursor_1 = stats_cursor_1 + 1; + unsigned int tail_stats_stage_1 = stats_cursor_1 & 1; + unsigned int tail_stats_phase = stats_cursor_1 >> 1 & 1; + mbarrier_wait(stats_addr + (tail_stats_stage_1) * 8, tail_stats_phase); + mbarrier_arrive(stats_empty_addr + (tail_stats_stage_1) * 8); + stats_cursor_1 = stats_cursor_1 + 1; + mbarrier_wait(o_full_addr, global_tile_1 - 1 & 1); + asm volatile("tcgen05.fence::after_thread_sync;"); + float final_sum = _tmem_load_3[0]; + float _rcp_0 = approx_rcp(final_sum); + float inv_sum = _rcp_0; + int output_base = (query_idx_1 * (unsigned int)num_heads + (unsigned int)my_row_1) * 512; + #pragma unroll + for (int v_stage_1 = 0; v_stage_1 < 4; v_stage_1++) { + int output_addr_1 = taddr + 256 + (unsigned int)((v_stage_1 & 1) * 128) + (unsigned int)((v_stage_1 >> 1) * 16 << 16) + (unsigned int)(tmem_row_origin_1 << 16); + float _tmem_load_4[64]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_4[0])), "=r"(*reinterpret_cast(&_tmem_load_4[1])), "=r"(*reinterpret_cast(&_tmem_load_4[2])), "=r"(*reinterpret_cast(&_tmem_load_4[3])), "=r"(*reinterpret_cast(&_tmem_load_4[4])), "=r"(*reinterpret_cast(&_tmem_load_4[5])), "=r"(*reinterpret_cast(&_tmem_load_4[6])), "=r"(*reinterpret_cast(&_tmem_load_4[7])), "=r"(*reinterpret_cast(&_tmem_load_4[8])), "=r"(*reinterpret_cast(&_tmem_load_4[9])), "=r"(*reinterpret_cast(&_tmem_load_4[10])), "=r"(*reinterpret_cast(&_tmem_load_4[11])), "=r"(*reinterpret_cast(&_tmem_load_4[12])), "=r"(*reinterpret_cast(&_tmem_load_4[13])), "=r"(*reinterpret_cast(&_tmem_load_4[14])), "=r"(*reinterpret_cast(&_tmem_load_4[15])), "=r"(*reinterpret_cast(&_tmem_load_4[16])), "=r"(*reinterpret_cast(&_tmem_load_4[17])), "=r"(*reinterpret_cast(&_tmem_load_4[18])), "=r"(*reinterpret_cast(&_tmem_load_4[19])), "=r"(*reinterpret_cast(&_tmem_load_4[20])), "=r"(*reinterpret_cast(&_tmem_load_4[21])), "=r"(*reinterpret_cast(&_tmem_load_4[22])), "=r"(*reinterpret_cast(&_tmem_load_4[23])), "=r"(*reinterpret_cast(&_tmem_load_4[24])), "=r"(*reinterpret_cast(&_tmem_load_4[25])), "=r"(*reinterpret_cast(&_tmem_load_4[26])), "=r"(*reinterpret_cast(&_tmem_load_4[27])), "=r"(*reinterpret_cast(&_tmem_load_4[28])), "=r"(*reinterpret_cast(&_tmem_load_4[29])), "=r"(*reinterpret_cast(&_tmem_load_4[30])), "=r"(*reinterpret_cast(&_tmem_load_4[31])) + : "r"(output_addr_1)); + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_4[32])), "=r"(*reinterpret_cast(&_tmem_load_4[33])), "=r"(*reinterpret_cast(&_tmem_load_4[34])), "=r"(*reinterpret_cast(&_tmem_load_4[35])), "=r"(*reinterpret_cast(&_tmem_load_4[36])), "=r"(*reinterpret_cast(&_tmem_load_4[37])), "=r"(*reinterpret_cast(&_tmem_load_4[38])), "=r"(*reinterpret_cast(&_tmem_load_4[39])), "=r"(*reinterpret_cast(&_tmem_load_4[40])), "=r"(*reinterpret_cast(&_tmem_load_4[41])), "=r"(*reinterpret_cast(&_tmem_load_4[42])), "=r"(*reinterpret_cast(&_tmem_load_4[43])), "=r"(*reinterpret_cast(&_tmem_load_4[44])), "=r"(*reinterpret_cast(&_tmem_load_4[45])), "=r"(*reinterpret_cast(&_tmem_load_4[46])), "=r"(*reinterpret_cast(&_tmem_load_4[47])), "=r"(*reinterpret_cast(&_tmem_load_4[48])), "=r"(*reinterpret_cast(&_tmem_load_4[49])), "=r"(*reinterpret_cast(&_tmem_load_4[50])), "=r"(*reinterpret_cast(&_tmem_load_4[51])), "=r"(*reinterpret_cast(&_tmem_load_4[52])), "=r"(*reinterpret_cast(&_tmem_load_4[53])), "=r"(*reinterpret_cast(&_tmem_load_4[54])), "=r"(*reinterpret_cast(&_tmem_load_4[55])), "=r"(*reinterpret_cast(&_tmem_load_4[56])), "=r"(*reinterpret_cast(&_tmem_load_4[57])), "=r"(*reinterpret_cast(&_tmem_load_4[58])), "=r"(*reinterpret_cast(&_tmem_load_4[59])), "=r"(*reinterpret_cast(&_tmem_load_4[60])), "=r"(*reinterpret_cast(&_tmem_load_4[61])), "=r"(*reinterpret_cast(&_tmem_load_4[62])), "=r"(*reinterpret_cast(&_tmem_load_4[63])) + : "r"(output_addr_1 + 32)); + int col_base = v_stage_1 * 128 + col_half_1 * 64; + if (my_row_1 < num_heads) { + #pragma unroll + for (int offset = 0; offset < 64; offset += 16) { + { + const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; + #if __CUDA_ARCH__ >= 1000 + #pragma unroll + for (int _ps = 0; _ps < 8; _ps++) + mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_4[offset])[_ps], _prescale2_1); + #else + #pragma unroll + for (int _ps = 0; _ps < 16; _ps++) + _tmem_load_4[offset + _ps] *= inv_sum * output_scale; + #endif + { + __nv_bfloat162 _pk0 = __floats2bfloat162_rn(_tmem_load_4[offset + 0], _tmem_load_4[offset + 1]); + unsigned _pk_u0 = *reinterpret_cast(&_pk0); + __nv_bfloat162 _pk1 = __floats2bfloat162_rn(_tmem_load_4[offset + 2], _tmem_load_4[offset + 3]); + unsigned _pk_u1 = *reinterpret_cast(&_pk1); + __nv_bfloat162 _pk2 = __floats2bfloat162_rn(_tmem_load_4[offset + 4], _tmem_load_4[offset + 5]); + unsigned _pk_u2 = *reinterpret_cast(&_pk2); + __nv_bfloat162 _pk3 = __floats2bfloat162_rn(_tmem_load_4[offset + 6], _tmem_load_4[offset + 7]); + unsigned _pk_u3 = *reinterpret_cast(&_pk3); + __nv_bfloat162 _pk4 = __floats2bfloat162_rn(_tmem_load_4[offset + 8], _tmem_load_4[offset + 9]); + unsigned _pk_u4 = *reinterpret_cast(&_pk4); + __nv_bfloat162 _pk5 = __floats2bfloat162_rn(_tmem_load_4[offset + 10], _tmem_load_4[offset + 11]); + unsigned _pk_u5 = *reinterpret_cast(&_pk5); + __nv_bfloat162 _pk6 = __floats2bfloat162_rn(_tmem_load_4[offset + 12], _tmem_load_4[offset + 13]); + unsigned _pk_u6 = *reinterpret_cast(&_pk6); + __nv_bfloat162 _pk7 = __floats2bfloat162_rn(_tmem_load_4[offset + 14], _tmem_load_4[offset + 15]); + unsigned _pk_u7 = *reinterpret_cast(&_pk7); + asm volatile( + "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"((void*)(&((__nv_bfloat16*)(O + (output_base + col_base + offset)))[0])), "r"(_pk_u0), "r"(_pk_u1), "r"(_pk_u2), "r"(_pk_u3), "r"(_pk_u4), "r"(_pk_u5), "r"(_pk_u6), "r"(_pk_u7) : "memory"); + } + } + } + } + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(o_empty_addr); + mbarrier_wait(work_full_addr + (work_stage_1) * 8, _phase_work_full_1); + uint32_t _clc_valid_3 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_3) + : "r"(work_response_addr + work_stage_1 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_3 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_3) + : "r"(work_response_addr + work_stage_1 * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(work_empty_addr + (work_stage_1) * 8); + work_stage_1 += 1; + if (work_stage_1 == 2) { work_stage_1 = 0; _phase_work_full_1 ^= 1; } + unsigned int valid_1 = _clc_valid_3; + query_idx_1 = _clc_ctaid_3; + if (valid_1 == 0) { + break; + } + } + asm volatile("barrier.sync 7, 128;" ::: "memory"); + if (warp == 4) { + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + } + } + // ---- Role: mma_warp ---- + if (warp == 8) { + { // mma_warp_main + unsigned int work_stage_2 = 0; + unsigned int query_idx_2 = blockIdx.x; + unsigned int global_tile_2 = 0; + unsigned int _phase_work_full_2 = 0; + #pragma unroll 1 + for (unsigned int work_iter_2 = 0; work_iter_2 < num_query_tokens; work_iter_2++) { + if (query_idx_2 >= (unsigned int)num_query_tokens) { + break; + } + int _max_2 = ((sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) : (0)); + int _min_1 = ((_max_2) < (sparse_topk) ? (_max_2) : (sparse_topk)); + int active_topk_2 = _min_1; + int _max_3 = (((active_topk_2 + 128 - 1) / 128) > (1) ? ((active_topk_2 + 128 - 1) / 128) : (1)); + int num_loop_steps_2 = _max_3; + unsigned int query_phase = work_iter_2 & 1; + mbarrier_wait(q_full_addr, query_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + if (num_loop_steps_2 > 0) { + unsigned int score_stage = global_tile_2 & 1; + int score_col_1 = ((score_stage != 0) ? 128 : 0); + #pragma unroll + for (int k_stage = 0; k_stage < 4; k_stage++) { + mbarrier_wait(kv_full_addr + (k_stage) * 8, 0); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_0 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (k_stage) * 1024); + int _mma_b_lo_0 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (k_stage) * 2048); + { + uint64_t _mma_ss_a_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_0); + uint64_t _mma_ss_b_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_0); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, ((k_stage == 0) ? 0 : 1)); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 506U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 1018U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); + } + } + if (k_stage == 3) { + elect_commit(s_full_addr + (score_stage) * 8); + } + elect_commit(kv_empty_addr + (k_stage) * 8); + } + #pragma unroll 1 + for (int tile_2 = 0; tile_2 < num_loop_steps_2 - 1; tile_2++) { + unsigned int next_global_tile = global_tile_2 + 1; + unsigned int next_score_stage = next_global_tile & 1; + int next_score_col = ((next_score_stage != 0) ? 128 : 0); + #pragma unroll + for (int k_stage_1 = 0; k_stage_1 < 4; k_stage_1++) { + mbarrier_wait(kv_full_addr + (k_stage_1) * 8, 1); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_1 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (k_stage_1) * 1024); + int _mma_b_lo_1 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (k_stage_1) * 2048); + { + uint64_t _mma_ss_a_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_1); + uint64_t _mma_ss_b_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_1); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (next_score_col)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, ((k_stage_1 == 0) ? 0 : 1)); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (next_score_col)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (next_score_col)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (next_score_col)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 506U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 1018U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (next_score_col)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (next_score_col)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (next_score_col)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (next_score_col)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); + } + } + if (k_stage_1 == 3) { + elect_commit(s_full_addr + (next_score_stage) * 8); + } + elect_commit(kv_empty_addr + (k_stage_1) * 8); + } + unsigned int sync_stage_1 = global_tile_2 & 1; + unsigned int sync_phase_1 = global_tile_2 >> 1 & 1; + int current_score_col = ((sync_stage_1 != 0) ? 128 : 0); + mbarrier_wait(p_full_addr + (sync_stage_1) * 8, sync_phase_1); + mbarrier_wait(o_empty_addr, 1 - (global_tile_2 & 1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int v_stage_2 = 0; v_stage_2 < 4; v_stage_2++) { + mbarrier_wait(kv_full_addr + (v_stage_2) * 8, 0); + asm volatile("tcgen05.fence::after_thread_sync;"); + int output_addr_2 = 256 + (v_stage_2 & 1) * 128 + ((v_stage_2 >> 1) * 16 << 16); + int _mma_b_lo_2 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (v_stage_2) * 2048); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 dhi, blo, ta, id;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 69272720;\n\t" + "mov.b32 ta, %2;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p0;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "}\n" + :: "r"((tmem_tmem + (output_addr_2))), "r"(_mma_b_lo_2), "r"(tmem_tmem + (current_score_col + 32 + ((v_stage_2 >> 1) * 16 << 16))), "r"(((tile_2 == 0) ? 0 : 1))); + elect_commit(kv_empty_addr + (v_stage_2) * 8); + } + elect_commit(o_full_addr); + global_tile_2 = next_global_tile; + } + elect_commit(q_empty_addr); + int last_tile = num_loop_steps_2 - 1; + unsigned int last_sync_stage = global_tile_2 & 1; + unsigned int last_sync_phase = global_tile_2 >> 1 & 1; + int last_score_col = ((last_sync_stage != 0) ? 128 : 0); + mbarrier_wait(p_full_addr + (last_sync_stage) * 8, last_sync_phase); + mbarrier_wait(o_empty_addr, 1 - (global_tile_2 & 1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int v_stage_3 = 0; v_stage_3 < 4; v_stage_3++) { + mbarrier_wait(kv_full_addr + (v_stage_3) * 8, 1); + asm volatile("tcgen05.fence::after_thread_sync;"); + int output_addr_3 = 256 + (v_stage_3 & 1) * 128 + ((v_stage_3 >> 1) * 16 << 16); + int _mma_b_lo_3 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (v_stage_3) * 2048); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 dhi, blo, ta, id;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 69272720;\n\t" + "mov.b32 ta, %2;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p0;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "}\n" + :: "r"((tmem_tmem + (output_addr_3))), "r"(_mma_b_lo_3), "r"(tmem_tmem + (last_score_col + 32 + ((v_stage_3 >> 1) * 16 << 16))), "r"(((last_tile == 0) ? 0 : 1))); + elect_commit(kv_empty_addr + (v_stage_3) * 8); + } + elect_commit(o_full_addr); + global_tile_2 = global_tile_2 + 1; + } else { + elect_commit(q_empty_addr); + } + mbarrier_wait(work_full_addr + (work_stage_2) * 8, _phase_work_full_2); + uint32_t _clc_valid_1 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_1) + : "r"(work_response_addr + work_stage_2 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_1 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_1) + : "r"(work_response_addr + work_stage_2 * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(work_empty_addr + (work_stage_2) * 8); + work_stage_2 += 1; + if (work_stage_2 == 2) { work_stage_2 = 0; _phase_work_full_2 ^= 1; } + unsigned int valid_2 = _clc_valid_1; + query_idx_2 = _clc_ctaid_1; + if (valid_2 == 0) { + break; + } + } + } + } + // ---- Role: page_warp ---- + if (warp == 9) { + { // page_warp_main + unsigned int page_stage = 0; + unsigned int work_stage_3 = 0; + unsigned int query_idx_3 = blockIdx.x; + unsigned int _phase_page_empty = 1; + unsigned int _phase_work_full_3 = 0; + #pragma unroll 1 + for (unsigned int work_iter_3 = 0; work_iter_3 < num_query_tokens; work_iter_3++) { + if (query_idx_3 >= (unsigned int)num_query_tokens) { + break; + } + int* swa_row = swa_indices + (query_idx_3 * (unsigned int)swa_index_stride); + int* compressed_row = compressed_indices + (query_idx_3 * (unsigned int)compressed_index_stride); + int _max_11 = ((sparse_topk_lens[query_idx_3] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_3] + sparse_topk_lens_offset) : (0)); + int _min_4 = ((_max_11) < (sparse_topk) ? (_max_11) : (sparse_topk)); + int active_topk_3 = _min_4; + int _max_12 = (((active_topk_3 + 128 - 1) / 128) > (1) ? ((active_topk_3 + 128 - 1) / 128) : (1)); + int num_loop_steps_3 = _max_12; + int num_page_pairs = (num_loop_steps_3 + 1) / 2; + int _max_13 = ((active_topk_3 - 1 & -4) > (0) ? (active_topk_3 - 1 & -4) : (0)); + int last_aligned = _max_13; + #pragma unroll 1 + for (int pair = 0; pair < num_page_pairs; pair++) { + int pair_base = pair * 256; + #pragma unroll + for (int duplicate = 0; duplicate < 2; duplicate++) { + mbarrier_wait(page_empty_addr + (page_stage) * 8, _phase_page_empty); + int lane_base = lane * 4; + int page0 = pair_base + lane_base; + int page1 = page0 + 128; + int safe0 = ((page0 <= last_aligned) ? page0 : last_aligned); + int safe1 = ((page1 <= last_aligned) ? page1 : last_aligned); + int* src0 = ((safe0 < 128) ? swa_row + safe0 : compressed_row + (safe0 - 128)); + int* src1 = ((safe1 < 128) ? swa_row + safe1 : compressed_row + (safe1 - 128)); + asm volatile("cp.async.ca.shared::cta.global [%0], [%1], 16;" + :: "r"(smem_page_offsets_addr + page_stage * 1024 + (unsigned int)(lane_base * 4)), "l"(src0)); + asm volatile("cp.async.ca.shared::cta.global [%0], [%1], 16;" + :: "r"(smem_page_offsets_addr + page_stage * 1024 + (unsigned int)((128 + lane_base) * 4)), "l"(src1)); + asm volatile( + "{\n\t" + "cp.async.mbarrier.arrive.noinc.shared::cta.b64 [%0];\n\t" + "}" + :: "r"(page_full_addr + (page_stage) * 8) : "memory"); + page_stage += 1; + if (page_stage == 6) { page_stage = 0; _phase_page_empty ^= 1; } + } + } + mbarrier_wait(work_full_addr + (work_stage_3) * 8, _phase_work_full_3); + uint32_t _clc_valid_4 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_4) + : "r"(work_response_addr + work_stage_3 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_4 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_4) + : "r"(work_response_addr + work_stage_3 * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(work_empty_addr + (work_stage_3) * 8); + work_stage_3 += 1; + if (work_stage_3 == 2) { work_stage_3 = 0; _phase_work_full_3 ^= 1; } + unsigned int valid_3 = _clc_valid_4; + query_idx_3 = _clc_ctaid_4; + if (valid_3 == 0) { + break; + } + } + } + } + // ---- Role: scheduler_warp ---- + if (warp == 10) { + { // scheduler_warp_main + unsigned int work_stage_4 = 0; + unsigned int throttle_stage = 0; + unsigned int query_idx_4 = blockIdx.x; + unsigned int _phase_throttle_full = 0; + unsigned int _phase_work_empty = 1; + unsigned int _phase_work_full_4 = 0; + #pragma unroll 1 + for (unsigned int work_iter_4 = 0; work_iter_4 < num_query_tokens; work_iter_4++) { + if (query_idx_4 >= (unsigned int)num_query_tokens) { + break; + } + mbarrier_wait(throttle_full_addr + (throttle_stage) * 8, _phase_throttle_full); + mbarrier_arrive(throttle_empty_addr + (throttle_stage) * 8); + throttle_stage += 1; + if (throttle_stage == 2) { throttle_stage = 0; _phase_throttle_full ^= 1; } + if (elect_sync()) { + mbarrier_wait(work_empty_addr + (work_stage_4) * 8, _phase_work_empty); + mbarrier_arrive_expect_tx(work_full_addr + (work_stage_4) * 8, 16); + asm volatile( + "fence.proxy.async.shared::cta;\n\t" + "clusterlaunchcontrol.try_cancel.async.shared::cta" + ".mbarrier::complete_tx::bytes.b128" + " [%0], [%1];" + :: "r"(work_response_addr + work_stage_4 * 16 + 0 * 16), "r"(work_full_addr + work_stage_4 * 8) + : "memory"); + } + mbarrier_wait(work_full_addr + (work_stage_4) * 8, _phase_work_full_4); + uint32_t _clc_valid_5 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_5) + : "r"(work_response_addr + work_stage_4 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_5 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_5) + : "r"(work_response_addr + work_stage_4 * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(work_empty_addr + (work_stage_4) * 8); + work_stage_4 += 1; + if (work_stage_4 == 2) { work_stage_4 = 0; _phase_work_empty ^= 1; _phase_work_full_4 ^= 1; } + unsigned int valid_4 = _clc_valid_5; + query_idx_4 = _clc_ctaid_5; + if (valid_4 == 0) { + break; + } + } + } + } + // ---- Role: padding_warp ---- + if (warp == 11) { + { // padding_warp_main + unsigned int work_stage_5 = 0; + unsigned int query_idx_5 = blockIdx.x; + unsigned int _phase_work_full_5 = 0; + #pragma unroll 1 + for (unsigned int work_iter_5 = 0; work_iter_5 < num_query_tokens; work_iter_5++) { + if (query_idx_5 >= (unsigned int)num_query_tokens) { + break; + } + mbarrier_wait(work_full_addr + (work_stage_5) * 8, _phase_work_full_5); + uint32_t _clc_valid_6 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_6) + : "r"(work_response_addr + work_stage_5 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_6 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_6) + : "r"(work_response_addr + work_stage_5 * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(work_empty_addr + (work_stage_5) * 8); + work_stage_5 += 1; + if (work_stage_5 == 2) { work_stage_5 = 0; _phase_work_full_5 ^= 1; } + unsigned int valid_5 = _clc_valid_6; + query_idx_5 = _clc_ctaid_6; + if (valid_5 == 0) { + break; + } + } + } + } + // ---- Role: load_warp ---- + if (warp >= 12 && warp <= 15) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 72;"); + { // load_warp_main + asm volatile("griddepcontrol.wait;" ::: "memory"); + const int load_warp_rank = warp - 12; + unsigned int load_page_stage = 0; + unsigned int page_k_stage = 0; + unsigned int page_v_stage = 0; + int page_k_addr = 0; + int page_v_addr = 0; + unsigned int work_stage_6 = 0; + unsigned int throttle_stage_1 = 0; + unsigned int query_idx_6 = blockIdx.x; + unsigned int _phase_throttle_empty = 1; + unsigned int _phase_page_full = 0; + unsigned int _phase_work_full_6 = 0; + #pragma unroll 1 + for (unsigned int work_iter_6 = 0; work_iter_6 < num_query_tokens; work_iter_6++) { + if (query_idx_6 >= (unsigned int)num_query_tokens) { + break; + } + int _max_0 = ((sparse_topk_lens[query_idx_6] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_6] + sparse_topk_lens_offset) : (0)); + int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); + int active_topk_4 = _min_0; + int _max_1 = (((active_topk_4 + 128 - 1) / 128) > (1) ? ((active_topk_4 + 128 - 1) / 128) : (1)); + int num_loop_steps_4 = _max_1; + mbarrier_wait(throttle_empty_addr + (throttle_stage_1) * 8, _phase_throttle_empty); + mbarrier_arrive(throttle_full_addr + (throttle_stage_1) * 8); + throttle_stage_1 += 1; + if (throttle_stage_1 == 2) { throttle_stage_1 = 0; _phase_throttle_empty ^= 1; } + unsigned int query_phase_1 = work_iter_6 & 1; + mbarrier_wait(q_empty_addr, 1 - query_phase_1); + if (warp == 12) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(q_full_addr, 65536); + #pragma unroll + for (int q_subload = 0; q_subload < 8; q_subload++) { + tma_4d_gmem2smem(smem_q_addr + (unsigned int)(q_subload * 8192), (&tmap_q), 0, 0, q_subload, query_idx_6, q_full_addr); + } + } + } + if (num_loop_steps_4 > 0) { + mbarrier_wait(page_full_addr + (load_page_stage) * 8, _phase_page_full); + page_k_stage = load_page_stage; + page_k_addr = smem_page_offsets_addr + page_k_stage * 1024; + load_page_stage += 1; + if (load_page_stage == 6) { load_page_stage = 0; _phase_page_full ^= 1; } + #pragma unroll + for (int kv_stage = 0; kv_stage < 4; kv_stage++) { + mbarrier_wait(kv_empty_addr + (kv_stage) * 8, 1); + if (warp == 12) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage) * 8, 32768); + } + } + #pragma unroll + for (int group = 0; group < 8; group++) { + int row_offset = load_warp_rank * 4 + group * 16; + int raw_rows[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows[0])), "=r"(*reinterpret_cast(&raw_rows[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows[(0) + 3])) + : "r"(page_k_addr + row_offset * 4)); + int row0 = raw_rows[0]; + int row1 = raw_rows[1]; + int row2 = raw_rows[2]; + int row3 = raw_rows[3]; + int dst = smem_kv_addr + (unsigned int)(kv_stage * 32768) + (unsigned int)(row_offset * 64 * 2); + if (elect_sync()) { + { + tma_gather4_gmem2smem(dst, (&tmap_swa_kv), kv_stage * 128, row0, row1, row2, row3, kv_full_addr + (kv_stage) * 8); + tma_gather4_gmem2smem(dst + 16384, (&tmap_swa_kv), kv_stage * 128 + 64, row0, row1, row2, row3, kv_full_addr + (kv_stage) * 8); + } + } + } + } + if (num_loop_steps_4 > 1) { + #pragma unroll + for (int kv_stage_1 = 0; kv_stage_1 < 4; kv_stage_1++) { + mbarrier_wait(kv_empty_addr + (kv_stage_1) * 8, 0); + if (warp == 12) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_1) * 8, 32768); + } + } + #pragma unroll + for (int group_1 = 0; group_1 < 8; group_1++) { + int row_offset_1 = load_warp_rank * 4 + group_1 * 16; + int raw_rows_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_1[0])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 3])) + : "r"(page_k_addr + 512 + row_offset_1 * 4)); + int row0_1 = raw_rows_1[0]; + int row1_1 = raw_rows_1[1]; + int row2_1 = raw_rows_1[2]; + int row3_1 = raw_rows_1[3]; + int dst_1 = smem_kv_addr + (unsigned int)(kv_stage_1 * 32768) + (unsigned int)(row_offset_1 * 64 * 2); + if (elect_sync()) { + { + tma_gather4_gmem2smem(dst_1, (&tmap_compressed_kv), kv_stage_1 * 128, row0_1, row1_1, row2_1, row3_1, kv_full_addr + (kv_stage_1) * 8); + tma_gather4_gmem2smem(dst_1 + 16384, (&tmap_compressed_kv), kv_stage_1 * 128 + 64, row0_1, row1_1, row2_1, row3_1, kv_full_addr + (kv_stage_1) * 8); + } + } + } + } + if (num_loop_steps_4 > 2) { + mbarrier_arrive(page_empty_addr + (page_k_stage) * 8); + } + mbarrier_wait(page_full_addr + (load_page_stage) * 8, _phase_page_full); + page_v_stage = load_page_stage; + page_v_addr = smem_page_offsets_addr + page_v_stage * 1024; + load_page_stage += 1; + if (load_page_stage == 6) { load_page_stage = 0; _phase_page_full ^= 1; } + #pragma unroll + for (int kv_stage_2 = 0; kv_stage_2 < 4; kv_stage_2++) { + mbarrier_wait(kv_empty_addr + (kv_stage_2) * 8, 1); + if (warp == 12) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_2) * 8, 32768); + } + } + #pragma unroll + for (int group_2 = 0; group_2 < 8; group_2++) { + int row_offset_2 = load_warp_rank * 4 + group_2 * 16; + int raw_rows_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_2[0])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 3])) + : "r"(page_v_addr + row_offset_2 * 4)); + int row0_2 = raw_rows_2[0]; + int row1_2 = raw_rows_2[1]; + int row2_2 = raw_rows_2[2]; + int row3_2 = raw_rows_2[3]; + int dst_2 = smem_v_addr + (unsigned int)(kv_stage_2 * 32768) + (unsigned int)(row_offset_2 * 64 * 2); + if (elect_sync()) { + { + tma_gather4_gmem2smem(dst_2, (&tmap_swa_kv), kv_stage_2 * 128, row0_2, row1_2, row2_2, row3_2, kv_full_addr + (kv_stage_2) * 8); + tma_gather4_gmem2smem(dst_2 + 16384, (&tmap_swa_kv), kv_stage_2 * 128 + 64, row0_2, row1_2, row2_2, row3_2, kv_full_addr + (kv_stage_2) * 8); + } + } + } + } + } + #pragma unroll 1 + for (int tile_3 = 1; tile_3 < num_loop_steps_4 - 1; tile_3++) { + int next_tile = tile_3 + 1; + if ((next_tile & 1) == 0) { + mbarrier_wait(page_full_addr + (load_page_stage) * 8, _phase_page_full); + page_k_stage = load_page_stage; + page_k_addr = smem_page_offsets_addr + page_k_stage * 1024; + load_page_stage += 1; + if (load_page_stage == 6) { load_page_stage = 0; _phase_page_full ^= 1; } + } + #pragma unroll + for (int kv_stage_3 = 0; kv_stage_3 < 4; kv_stage_3++) { + mbarrier_wait(kv_empty_addr + (kv_stage_3) * 8, 0); + if (warp == 12) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_3) * 8, 32768); + } + } + #pragma unroll + for (int group_3 = 0; group_3 < 8; group_3++) { + int row_offset_3 = load_warp_rank * 4 + group_3 * 16; + int raw_rows_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_3[0])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 3])) + : "r"(page_k_addr + (next_tile & 1) * 128 * 4 + row_offset_3 * 4)); + int row0_3 = raw_rows_3[0]; + int row1_3 = raw_rows_3[1]; + int row2_3 = raw_rows_3[2]; + int row3_3 = raw_rows_3[3]; + int dst_3 = smem_kv_addr + (unsigned int)(kv_stage_3 * 32768) + (unsigned int)(row_offset_3 * 64 * 2); + if (elect_sync()) { + { + tma_gather4_gmem2smem(dst_3, (&tmap_compressed_kv), kv_stage_3 * 128, row0_3, row1_3, row2_3, row3_3, kv_full_addr + (kv_stage_3) * 8); + tma_gather4_gmem2smem(dst_3 + 16384, (&tmap_compressed_kv), kv_stage_3 * 128 + 64, row0_3, row1_3, row2_3, row3_3, kv_full_addr + (kv_stage_3) * 8); + } + } + } + } + if ((next_tile & 1) != 0 && tile_3 < num_loop_steps_4 - 2) { + mbarrier_arrive(page_empty_addr + (page_k_stage) * 8); + } + if ((tile_3 & 1) == 0) { + mbarrier_wait(page_full_addr + (load_page_stage) * 8, _phase_page_full); + page_v_stage = load_page_stage; + page_v_addr = smem_page_offsets_addr + page_v_stage * 1024; + load_page_stage += 1; + if (load_page_stage == 6) { load_page_stage = 0; _phase_page_full ^= 1; } + } + #pragma unroll + for (int kv_stage_4 = 0; kv_stage_4 < 4; kv_stage_4++) { + mbarrier_wait(kv_empty_addr + (kv_stage_4) * 8, 1); + if (warp == 12) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_4) * 8, 32768); + } + } + #pragma unroll + for (int group_4 = 0; group_4 < 8; group_4++) { + int row_offset_4 = load_warp_rank * 4 + group_4 * 16; + int raw_rows_4[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_4[0])), "=r"(*reinterpret_cast(&raw_rows_4[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_4[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_4[(0) + 3])) + : "r"(page_v_addr + (tile_3 & 1) * 128 * 4 + row_offset_4 * 4)); + int row0_4 = raw_rows_4[0]; + int row1_4 = raw_rows_4[1]; + int row2_4 = raw_rows_4[2]; + int row3_4 = raw_rows_4[3]; + int dst_4 = smem_v_addr + (unsigned int)(kv_stage_4 * 32768) + (unsigned int)(row_offset_4 * 64 * 2); + if (elect_sync()) { + { + tma_gather4_gmem2smem(dst_4, (&tmap_compressed_kv), kv_stage_4 * 128, row0_4, row1_4, row2_4, row3_4, kv_full_addr + (kv_stage_4) * 8); + tma_gather4_gmem2smem(dst_4 + 16384, (&tmap_compressed_kv), kv_stage_4 * 128 + 64, row0_4, row1_4, row2_4, row3_4, kv_full_addr + (kv_stage_4) * 8); + } + } + } + } + if ((tile_3 & 1) != 0) { + mbarrier_arrive(page_empty_addr + (page_v_stage) * 8); + } + } + int last_tile_1 = num_loop_steps_4 - 1; + mbarrier_arrive(page_empty_addr + (page_k_stage) * 8); + if ((last_tile_1 & 1) == 0) { + mbarrier_wait(page_full_addr + (load_page_stage) * 8, _phase_page_full); + page_v_stage = load_page_stage; + page_v_addr = smem_page_offsets_addr + page_v_stage * 1024; + load_page_stage += 1; + if (load_page_stage == 6) { load_page_stage = 0; _phase_page_full ^= 1; } + } + if (num_loop_steps_4 == 1) { + #pragma unroll + for (int kv_stage_5 = 0; kv_stage_5 < 4; kv_stage_5++) { + mbarrier_wait(kv_empty_addr + (kv_stage_5) * 8, 0); + if (warp == 12) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_5) * 8, 32768); + } + } + #pragma unroll + for (int group_5 = 0; group_5 < 8; group_5++) { + int row_offset_5 = load_warp_rank * 4 + group_5 * 16; + int raw_rows_5[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_5[0])), "=r"(*reinterpret_cast(&raw_rows_5[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_5[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_5[(0) + 3])) + : "r"(page_v_addr + row_offset_5 * 4)); + int row0_5 = raw_rows_5[0]; + int row1_5 = raw_rows_5[1]; + int row2_5 = raw_rows_5[2]; + int row3_5 = raw_rows_5[3]; + int dst_5 = smem_v_addr + (unsigned int)(kv_stage_5 * 32768) + (unsigned int)(row_offset_5 * 64 * 2); + if (elect_sync()) { + { + tma_gather4_gmem2smem(dst_5, (&tmap_swa_kv), kv_stage_5 * 128, row0_5, row1_5, row2_5, row3_5, kv_full_addr + (kv_stage_5) * 8); + tma_gather4_gmem2smem(dst_5 + 16384, (&tmap_swa_kv), kv_stage_5 * 128 + 64, row0_5, row1_5, row2_5, row3_5, kv_full_addr + (kv_stage_5) * 8); + } + } + } + } + } else { + #pragma unroll + for (int kv_stage_6 = 0; kv_stage_6 < 4; kv_stage_6++) { + mbarrier_wait(kv_empty_addr + (kv_stage_6) * 8, 0); + if (warp == 12) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_6) * 8, 32768); + } + } + #pragma unroll + for (int group_6 = 0; group_6 < 8; group_6++) { + int row_offset_6 = load_warp_rank * 4 + group_6 * 16; + int raw_rows_6[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_6[0])), "=r"(*reinterpret_cast(&raw_rows_6[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_6[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_6[(0) + 3])) + : "r"(page_v_addr + (last_tile_1 & 1) * 128 * 4 + row_offset_6 * 4)); + int row0_6 = raw_rows_6[0]; + int row1_6 = raw_rows_6[1]; + int row2_6 = raw_rows_6[2]; + int row3_6 = raw_rows_6[3]; + int dst_6 = smem_v_addr + (unsigned int)(kv_stage_6 * 32768) + (unsigned int)(row_offset_6 * 64 * 2); + if (elect_sync()) { + { + tma_gather4_gmem2smem(dst_6, (&tmap_compressed_kv), kv_stage_6 * 128, row0_6, row1_6, row2_6, row3_6, kv_full_addr + (kv_stage_6) * 8); + tma_gather4_gmem2smem(dst_6 + 16384, (&tmap_compressed_kv), kv_stage_6 * 128 + 64, row0_6, row1_6, row2_6, row3_6, kv_full_addr + (kv_stage_6) * 8); + } + } + } + } + } + mbarrier_arrive(page_empty_addr + (page_v_stage) * 8); + } + mbarrier_wait(work_full_addr + (work_stage_6) * 8, _phase_work_full_6); + uint32_t _clc_valid_0 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_0) + : "r"(work_response_addr + work_stage_6 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_0 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_0) + : "r"(work_response_addr + work_stage_6 * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(work_empty_addr + (work_stage_6) * 8); + work_stage_6 += 1; + if (work_stage_6 == 2) { work_stage_6 = 0; _phase_work_full_6 ^= 1; } + unsigned int valid_6 = _clc_valid_0; + query_idx_6 = _clc_ctaid_0; + if (valid_6 == 0) { + break; + } + } + } + } + + // Cleanup +} + +} // extern "C" diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_1ff52fece183b851b441_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_1ff52fece183b851b441_binding.cu new file mode 100644 index 00000000000..d59944bacd5 --- /dev/null +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_1ff52fece183b851b441_binding.cu @@ -0,0 +1,534 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +// Generated by the Cake source exporter — do not edit. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_1ff52fece183b851b441'. +#include +#include +#include +#include + +#include "tvm_ffi_utils.h" + +#include +#include +#include +#include +#include +#include +#include + +struct CakeTensorMap; +extern "C" __global__ void kernel_cake_dsv4_1ff52fece183b851b441(CakeTensorMap const* tmap_q, CakeTensorMap const* tmap_swa_kv, CakeTensorMap const* tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int batch_size, int max_q_len, int ragged_query, int has_sinks); + + +namespace cake_host_shim_831b6e81c0a7ccc0 { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t CakeDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +// Explicit immutable caller storage: initialize once, outside graph capture, +// exactly as the production TmaDeviceSlot path. The owner must retain the +// allocation and must never expose it as mutable scratch or rewrite its bytes. +template +static inline void PrepareImmutableCallerTmaWorkspace( + const tvm::ffi::Tensor& retained_workspace, + const CUtensorMap (&maps)[N], cudaStream_t stream) { + void* workspace = retained_workspace.data_ptr(); + CUcontext context = nullptr; + CUresult result = cuCtxGetCurrent(&context); + TVM_FFI_CHECK(result == CUDA_SUCCESS && context != nullptr, RuntimeError) + << "immutable TMA workspace requires an active CUDA context"; + CUcontext allocation_context = nullptr; + result = cuPointerGetAttribute( + &allocation_context, CU_POINTER_ATTRIBUTE_CONTEXT, + reinterpret_cast(workspace)); + TVM_FFI_CHECK(result == CUDA_SUCCESS && allocation_context == context, RuntimeError) + << "immutable TMA workspace must belong to the active CUDA context"; + unsigned long long buffer_id = 0; + result = cuPointerGetAttribute( + &buffer_id, CU_POINTER_ATTRIBUTE_BUFFER_ID, + reinterpret_cast(workspace)); + TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) + << "querying immutable TMA allocation identity failed"; + std::string key = std::to_string(reinterpret_cast(context)); + key += ":" + std::to_string(reinterpret_cast(workspace)); + key += ":" + std::to_string(buffer_id); + std::string descriptor_bytes(reinterpret_cast(maps), sizeof(maps)); + struct InitializedWorkspace { + tvm::ffi::Tensor owner; + std::string descriptor_bytes; + }; + static std::mutex mu; + // Keep the actual Tensor allocation alive: a caching allocator may reuse a + // freed suballocation without changing its enclosing CUDA buffer identity. + static auto* initialized = new std::unordered_map(); + std::lock_guard lock(mu); + auto it = initialized->find(key); + if (it != initialized->end()) { + TVM_FFI_CHECK(it->second.descriptor_bytes == descriptor_bytes, ValueError) + << "immutable TMA descriptors changed; provide new private descriptor storage"; + return; + } + CUstreamCaptureStatus capture_status = CU_STREAM_CAPTURE_STATUS_NONE; + result = cuStreamIsCapturing(reinterpret_cast(stream), &capture_status); + TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) + << "querying immutable TMA initialization capture state failed"; + TVM_FFI_CHECK(capture_status == CU_STREAM_CAPTURE_STATUS_NONE, RuntimeError) + << "immutable TMA storage must be initialized before CUDA Graph capture"; + // Synchronous host-to-device initialization completes before publishing the + // slot to any thread/stream, matching TmaDeviceSlot. The consumer retains its + // generated tensor-map acquire; the mutable upload's release fence is unchanged. + result = cuMemcpyHtoD(reinterpret_cast(workspace), maps, sizeof(maps)); + TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) + << "initializing immutable caller TMA descriptors failed"; + initialized->emplace( + std::move(key), InitializedWorkspace{retained_workspace, std::move(descriptor_bytes)}); +} + +static inline void* CallerTmaWorkspaceSlot( + const tvm::ffi::TensorView& workspace, size_t slot) { + return static_cast(workspace.data_ptr()) + slot * sizeof(CUtensorMap); +} + +// 4D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + int64_t d3 = t.size(t.ndim() - 3); + TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(d3)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 8u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 64, 8, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 64u, 8u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_swa_kv' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_swa_kv' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_kv'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_compressed_kv' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_compressed_kv(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_compressed_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_compressed_kv' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_compressed_kv' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_compressed_kv"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_compressed_kv' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_kv'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_kv') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_compressed_kv, TensorView arg_O, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_seq_lens, TensorView arg_cum_seq_lens_q, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_batch_size, int64_t arg_max_q_len, int64_t arg_ragged_query, int64_t arg_has_sinks, tvm::ffi::Tensor arg_tma_descriptor_workspace, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckContiguous(arg_tmap_q, "tmap_q"); + CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); + CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 4, 16, 1); + CheckCudaTensor(arg_tmap_compressed_kv, "tmap_compressed_kv"); + CheckDtype(arg_tmap_compressed_kv, "tmap_compressed_kv", 4, 16, 1); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_swa_indices, "swa_indices"); + CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); + CheckContiguous(arg_swa_indices, "swa_indices"); + CheckCudaTensor(arg_compressed_indices, "compressed_indices"); + CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); + CheckContiguous(arg_compressed_indices, "compressed_indices"); + CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); + CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckCudaTensor(arg_seq_lens, "seq_lens"); + CheckDtype(arg_seq_lens, "seq_lens", 0, 32, 1); + CheckContiguous(arg_seq_lens, "seq_lens"); + CheckCudaTensor(arg_cum_seq_lens_q, "cum_seq_lens_q"); + CheckDtype(arg_cum_seq_lens_q, "cum_seq_lens_q", 0, 32, 1); + CheckContiguous(arg_cum_seq_lens_q, "cum_seq_lens_q"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); + CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); + CheckContiguous(arg_bmm1_scale, "bmm1_scale"); + CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); + CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); + CheckContiguous(arg_bmm2_scale, "bmm2_scale"); + TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) + << "scalar 'num_heads' value " << arg_num_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) + << "scalar 'swa_index_stride' value " << arg_swa_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) + << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) + << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) + << "scalar 'sparse_topk' value " << arg_sparse_topk + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_batch_size >= -2147483648LL && arg_batch_size <= 2147483647LL, ValueError) + << "scalar 'batch_size' value " << arg_batch_size + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_q_len >= -2147483648LL && arg_max_q_len <= 2147483647LL, ValueError) + << "scalar 'max_q_len' value " << arg_max_q_len + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_ragged_query >= -2147483648LL && arg_ragged_query <= 2147483647LL, ValueError) + << "scalar 'ragged_query' value " << arg_ragged_query + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) + << "scalar 'has_sinks' value " << arg_has_sinks + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); + CheckSameCudaDevice(arg_tmap_compressed_kv, arg_tmap_q, "tmap_compressed_kv", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); + CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); + CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); + CheckSameCudaDevice(arg_seq_lens, arg_tmap_q, "seq_lens", "tmap_q"); + CheckSameCudaDevice(arg_cum_seq_lens_q, arg_tmap_q, "cum_seq_lens_q", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); + CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); + CheckCudaTensor(arg_tma_descriptor_workspace, "tma_descriptor_workspace"); + CheckSameCudaDevice(arg_tma_descriptor_workspace, arg_tmap_q, "tma_descriptor_workspace", "tmap_q"); + CheckContiguous(arg_tma_descriptor_workspace, "tma_descriptor_workspace"); + CheckDtype(arg_tma_descriptor_workspace, "tma_descriptor_workspace", 1, 8, 1); + TVM_FFI_CHECK(arg_tma_descriptor_workspace.numel() >= 384, ValueError) + << "tma_descriptor_workspace requires at least 384 bytes"; + TVM_FFI_CHECK( + reinterpret_cast(arg_tma_descriptor_workspace.data_ptr()) % 128 == 0, + ValueError) + << "tma_descriptor_workspace must be 128-byte aligned"; + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + + cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); + CUtensorMap h_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CakeTensorMap const* p_tmap_q = static_cast(CallerTmaWorkspaceSlot(arg_tma_descriptor_workspace, 0u)); + CUtensorMap h_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); + CakeTensorMap const* p_tmap_swa_kv = static_cast(CallerTmaWorkspaceSlot(arg_tma_descriptor_workspace, 1u)); + CUtensorMap h_tmap_compressed_kv = EncodeTma_tmap_compressed_kv(arg_tmap_compressed_kv); + CakeTensorMap const* p_tmap_compressed_kv = static_cast(CallerTmaWorkspaceSlot(arg_tma_descriptor_workspace, 2u)); + __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); + int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); + int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); + int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); + int* p_seq_lens = static_cast(arg_seq_lens.data_ptr()); + int* p_cum_seq_lens_q = static_cast(arg_cum_seq_lens_q.data_ptr()); + float* p_sinks = static_cast(arg_sinks.data_ptr()); + float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); + float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); + int32_t v_num_heads = (int32_t)arg_num_heads; + int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; + int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; + int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; + int32_t v_sparse_topk = (int32_t)arg_sparse_topk; + int32_t v_batch_size = (int32_t)arg_batch_size; + int32_t v_max_q_len = (int32_t)arg_max_q_len; + int32_t v_ragged_query = (int32_t)arg_ragged_query; + int32_t v_has_sinks = (int32_t)arg_has_sinks; + void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_tmap_compressed_kv, &p_O, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_seq_lens, &p_cum_seq_lens_q, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_batch_size, &v_max_q_len, &v_ragged_query, &v_has_sinks}; + + static std::mutex smem_mu; + static auto* smem_status_by_device = + new std::unordered_map(); + cudaError_t smem_status = cudaSuccess; + { + std::lock_guard lock(smem_mu); + auto it = smem_status_by_device->find(dev.device_id); + if (it == smem_status_by_device->end()) { + int current_device = -1; + smem_status = cudaGetDevice(¤t_device); + if (smem_status == cudaSuccess && current_device != dev.device_id) { + smem_status = cudaErrorInvalidDevice; + } + if (smem_status == cudaSuccess) { + smem_status = cudaFuncSetAttribute( + reinterpret_cast(kernel_cake_dsv4_1ff52fece183b851b441), + cudaFuncAttributeMaxDynamicSharedMemorySize, + 198144); + } + smem_status_by_device->emplace(dev.device_id, smem_status); + } else { + smem_status = it->second; + } + } + TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) + << "cudaFuncSetAttribute for kernel_cake_dsv4_1ff52fece183b851b441 failed: " + << cudaGetErrorString(smem_status); + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(256u, 1u, 1u); + const CUtensorMap workspace_maps[] = {h_tmap_q, h_tmap_swa_kv, h_tmap_compressed_kv}; + PrepareImmutableCallerTmaWorkspace(arg_tma_descriptor_workspace, workspace_maps, stream); + + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes = 198144u; + config.stream = stream; + cudaError_t launch_status = cudaLaunchKernelExC( + &config, reinterpret_cast(kernel_cake_dsv4_1ff52fece183b851b441), kargs); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernelExC for kernel_cake_dsv4_1ff52fece183b851b441 failed: " + << cudaGetErrorString(launch_status); +} + +} // namespace cake_host_shim_831b6e81c0a7ccc0 + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_831b6e81c0a7ccc0::Run); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_1ff52fece183b851b441_kernel.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_1ff52fece183b851b441_kernel.cu new file mode 100644 index 00000000000..47cf58b610b --- /dev/null +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_1ff52fece183b851b441_kernel.cu @@ -0,0 +1,1686 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); +template +struct __align__(128) CakeTensorMapPack { CakeTensorMap maps[N]; }; + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +#define CAKE_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_OFFSET 0 +#define NUM_KV_PIPE_STAGES 4 +#define SMEM_SMEM_Q_OFF 1024 +#define SMEM_SMEM_Q_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_STRIDE 8192 +#define SMEM_SMEM_KV_OFF 66560 +#define SMEM_SMEM_KV_STAGE_BYTES 32768 +#define SMEM_SMEM_KV_STRIDE 32768 +#define SMEM_SMEM_V_OFF 66560 +#define SMEM_SMEM_V_STAGE_BYTES 32768 +#define SMEM_SMEM_V_STRIDE 32768 +#define SMEM_TOTAL 198144 +#define THREADS 256 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count) : "memory"); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ float2 ex2_emulation_f32x2_value(float2 value) { + const float c0 = 1.0f, c1 = 0.695146143436431884765625f; + const float c2 = 0.227564394474029541015625f, c3 = 0.077119089663028717041015625f; + const float magic = 12582912.0f; + float x0 = max_noftz(value.x, -127.0f), x1 = max_noftz(value.y, -127.0f); + float2 xc2 = make_float2(x0, x1), magic2 = make_float2(magic, magic); + float2 xr2; + asm("add.rm.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xr2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&magic2)); + float2 c3_2 = make_float2(c3, c3), c2_2 = make_float2(c2, c2); + float2 c1_2 = make_float2(c1, c1), c0_2 = make_float2(c0, c0); + float2 xrb2, xfrac2; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xrb2) + : "l"(*(unsigned long long*)&xr2), "l"(*(unsigned long long*)&magic2)); + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xfrac2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&xrb2)); + float2 poly2; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&c3_2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c2_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c1_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c0_2)); + int x0r_i, x1r_i, p0_i, p1_i; + asm("mov.b64 {%0, %1}, %2;" : "=r"(x0r_i), "=r"(x1r_i) : "l"(*(unsigned long long*)&xr2)); + asm("mov.b64 {%0, %1}, %2;" : "=r"(p0_i), "=r"(p1_i) : "l"(*(unsigned long long*)&poly2)); + float r0, r1; + asm("mov.b32 %0, %1;" : "=f"(r0) : "r"((x0r_i << 23) + p0_i)); + asm("mov.b32 %0, %1;" : "=f"(r1) : "r"((x1r_i << 23) + p1_i)); + return make_float2(r0, r1); +} + +__device__ __forceinline__ void ex2_emulation_f32x2(float* x0_ptr, float* x1_ptr) { + float2 result = ex2_emulation_f32x2_value(make_float2(*x0_ptr, *x1_ptr)); + *x0_ptr = result.x; *x1_ptr = result.y; +} + +__device__ __forceinline__ void softmax_frag_exp2_cast( + float* sv, uint32_t* pv, int use_emu) +{ + #pragma unroll + for (int j = 0; j < 16; j++) { + if (use_emu && j >= 12) + ex2_emulation_f32x2(&sv[j*2], &sv[j*2+1]); + else { + sv[j*2] = approx_exp2(sv[j*2]); + sv[j*2+1] = approx_exp2(sv[j*2+1]); + } + } + #pragma unroll + for (int j = 0; j < 16; j++) { + __nv_bfloat162 bf = __float22bfloat162_rn({sv[j*2], sv[j*2+1]}); + pv[j] = reinterpret_cast(bf); + } +} + + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_2d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3}], [%4];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr) { + // Canonical .shared::cta form for non-multicast gather4, matching + // trtllm-gen / cuda_ptx and the PTX ISA qualifier order + // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, + // .shared::cluster is reserved for the multicast variant (ctaMask). + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(256, 1) void +kernel_cake_dsv4_1ff52fece183b851b441(CakeTensorMap const* tmap_q, CakeTensorMap const* tmap_swa_kv, CakeTensorMap const* tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int batch_size, int max_q_len, int ragged_query, int has_sinks) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int mbar_base = smem; + #define q_full_addr (mbar_base + 0) + #define kv_full_addr (mbar_base + 8) + #define s_full_addr (mbar_base + 40) + #define p_full_addr (mbar_base + 56) + #define pv_done_addr (mbar_base + 72) + #define tmem_dealloc_addr (mbar_base + 80) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + if (tid == 0) { + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_q)) : "memory"); + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_swa_kv)) : "memory"); + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_compressed_kv)) : "memory"); + } + __syncthreads(); + + + // Kernel setup ops + __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_q_addr = smem + 1024; + __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_kv_addr = smem + 66560; + __nv_bfloat16* smem_v = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_v_addr = smem + 66560; + + // Mbarrier init (6 pipeline groups, 0 ordered-sequence groups, 11 barriers) + // Mbarriers at smem_raw[0..88) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // q_full: 1 barriers, init_count=1 + mbarrier_init(smem + 0, 1); + // --- pipeline 'kv_pipe' --- + // kv_full: 4 barriers, init_count=3 + mbarrier_init(smem + 8, 3); + mbarrier_init(smem + 16, 3); + mbarrier_init(smem + 24, 3); + mbarrier_init(smem + 32, 3); + // s_full: 2 barriers, init_count=1 + mbarrier_init(smem + 40, 1); + mbarrier_init(smem + 48, 1); + // p_full: 2 barriers, init_count=4 + mbarrier_init(smem + 56, 4); + mbarrier_init(smem + 64, 4); + // pv_done: 1 barriers, init_count=1 + mbarrier_init(smem + 72, 1); + // tmem_dealloc: 1 barriers, init_count=7 + mbarrier_init(smem + 80, 7); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncwarp(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 88); + if (warp == 0) { + int _tmem_hold = smem + 88; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_tmem = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 64;"); + } + + // ---- Role: softmax ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 176;"); + { // softmax_main + float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; + int query_idx = blockIdx.x; + int _max_2 = ((sparse_topk_lens[query_idx] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx] + sparse_topk_lens_offset) : (0)); + int _min_2 = ((_max_2) < (sparse_topk) ? (_max_2) : (sparse_topk)); + int active_topk = _min_2; + int tile_count = (active_topk + 128 - 1) / 128; + int query_batch = query_idx / max_q_len; + int query_offset = query_idx - query_batch * max_q_len; + int query_length = max_q_len; + if (ragged_query != 0) { + query_batch = 0; + #pragma unroll 1 + for (int batch = 0; batch < batch_size; batch++) { + if (query_idx >= cum_seq_lens_q[batch + 1]) { + query_batch = batch + 1; + } + } + int query_begin = cum_seq_lens_q[query_batch]; + query_length = cum_seq_lens_q[query_batch + 1] - query_begin; + query_offset = query_idx - query_begin; + } + int swa_valid_columns = seq_lens[query_batch] - query_length + query_offset + 1; + if (swa_valid_columns < 0) { + swa_valid_columns = 0; + } + if (swa_valid_columns > 128) { + swa_valid_columns = 128; + } + const int warp_in_compute = warp; + const int tmem_row_origin = warp_in_compute * 32; + const int logical_row_origin = warp_in_compute * 16; + const int my_row = logical_row_origin + lane % 16; + const int col_half = lane / 16; + float row_max = -CAKE_INF; + float row_sum = 0.0f; + if (has_sinks != 0 && my_row < num_heads) { + row_max = sinks[my_row] * 1.4426950408889634f / softmax_scale_log2; + row_sum = 1.0f; + } + #pragma unroll + for (int tile = 0; tile < 1; tile++) { + if (tile_count > tile) { + int phase = tile & 1; + int wait_phase = tile >> 1 & 1; + mbarrier_wait(s_full_addr + (phase) * 8, wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int score_col = ((phase != 0) ? 128 : 0); + int score_addr = taddr + (unsigned int)score_col + (unsigned int)(tmem_row_origin << 16); + float _tmem_load_0[64]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_0[0])), "=r"(*reinterpret_cast(&_tmem_load_0[1])), "=r"(*reinterpret_cast(&_tmem_load_0[2])), "=r"(*reinterpret_cast(&_tmem_load_0[3])), "=r"(*reinterpret_cast(&_tmem_load_0[4])), "=r"(*reinterpret_cast(&_tmem_load_0[5])), "=r"(*reinterpret_cast(&_tmem_load_0[6])), "=r"(*reinterpret_cast(&_tmem_load_0[7])), "=r"(*reinterpret_cast(&_tmem_load_0[8])), "=r"(*reinterpret_cast(&_tmem_load_0[9])), "=r"(*reinterpret_cast(&_tmem_load_0[10])), "=r"(*reinterpret_cast(&_tmem_load_0[11])), "=r"(*reinterpret_cast(&_tmem_load_0[12])), "=r"(*reinterpret_cast(&_tmem_load_0[13])), "=r"(*reinterpret_cast(&_tmem_load_0[14])), "=r"(*reinterpret_cast(&_tmem_load_0[15])), "=r"(*reinterpret_cast(&_tmem_load_0[16])), "=r"(*reinterpret_cast(&_tmem_load_0[17])), "=r"(*reinterpret_cast(&_tmem_load_0[18])), "=r"(*reinterpret_cast(&_tmem_load_0[19])), "=r"(*reinterpret_cast(&_tmem_load_0[20])), "=r"(*reinterpret_cast(&_tmem_load_0[21])), "=r"(*reinterpret_cast(&_tmem_load_0[22])), "=r"(*reinterpret_cast(&_tmem_load_0[23])), "=r"(*reinterpret_cast(&_tmem_load_0[24])), "=r"(*reinterpret_cast(&_tmem_load_0[25])), "=r"(*reinterpret_cast(&_tmem_load_0[26])), "=r"(*reinterpret_cast(&_tmem_load_0[27])), "=r"(*reinterpret_cast(&_tmem_load_0[28])), "=r"(*reinterpret_cast(&_tmem_load_0[29])), "=r"(*reinterpret_cast(&_tmem_load_0[30])), "=r"(*reinterpret_cast(&_tmem_load_0[31])) + : "r"(score_addr)); + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_0[32])), "=r"(*reinterpret_cast(&_tmem_load_0[33])), "=r"(*reinterpret_cast(&_tmem_load_0[34])), "=r"(*reinterpret_cast(&_tmem_load_0[35])), "=r"(*reinterpret_cast(&_tmem_load_0[36])), "=r"(*reinterpret_cast(&_tmem_load_0[37])), "=r"(*reinterpret_cast(&_tmem_load_0[38])), "=r"(*reinterpret_cast(&_tmem_load_0[39])), "=r"(*reinterpret_cast(&_tmem_load_0[40])), "=r"(*reinterpret_cast(&_tmem_load_0[41])), "=r"(*reinterpret_cast(&_tmem_load_0[42])), "=r"(*reinterpret_cast(&_tmem_load_0[43])), "=r"(*reinterpret_cast(&_tmem_load_0[44])), "=r"(*reinterpret_cast(&_tmem_load_0[45])), "=r"(*reinterpret_cast(&_tmem_load_0[46])), "=r"(*reinterpret_cast(&_tmem_load_0[47])), "=r"(*reinterpret_cast(&_tmem_load_0[48])), "=r"(*reinterpret_cast(&_tmem_load_0[49])), "=r"(*reinterpret_cast(&_tmem_load_0[50])), "=r"(*reinterpret_cast(&_tmem_load_0[51])), "=r"(*reinterpret_cast(&_tmem_load_0[52])), "=r"(*reinterpret_cast(&_tmem_load_0[53])), "=r"(*reinterpret_cast(&_tmem_load_0[54])), "=r"(*reinterpret_cast(&_tmem_load_0[55])), "=r"(*reinterpret_cast(&_tmem_load_0[56])), "=r"(*reinterpret_cast(&_tmem_load_0[57])), "=r"(*reinterpret_cast(&_tmem_load_0[58])), "=r"(*reinterpret_cast(&_tmem_load_0[59])), "=r"(*reinterpret_cast(&_tmem_load_0[60])), "=r"(*reinterpret_cast(&_tmem_load_0[61])), "=r"(*reinterpret_cast(&_tmem_load_0[62])), "=r"(*reinterpret_cast(&_tmem_load_0[63])) + : "r"(score_addr + 32)); + int valid_cols = active_topk - tile * 128 - col_half * 64; + if (valid_cols < 0) { + valid_cols = 0; + } + if (valid_cols > 64) { + valid_cols = 64; + } + if (tile == 0) { + int valid_swa_cols = swa_valid_columns - col_half * 64; + if (valid_swa_cols < 0) { + valid_swa_cols = 0; + } + if (valid_swa_cols < valid_cols) { + valid_cols = valid_swa_cols; + } + } + uint32_t _slice_lo_mask_0; + { + int _lim_0 = valid_cols; + if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_0 & (1u << 0))) _tmem_load_0[0] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 1))) _tmem_load_0[1] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 2))) _tmem_load_0[2] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 3))) _tmem_load_0[3] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 4))) _tmem_load_0[4] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 5))) _tmem_load_0[5] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 6))) _tmem_load_0[6] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 7))) _tmem_load_0[7] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 8))) _tmem_load_0[8] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 9))) _tmem_load_0[9] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 10))) _tmem_load_0[10] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 11))) _tmem_load_0[11] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 12))) _tmem_load_0[12] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 13))) _tmem_load_0[13] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 14))) _tmem_load_0[14] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 15))) _tmem_load_0[15] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 16))) _tmem_load_0[16] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 17))) _tmem_load_0[17] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 18))) _tmem_load_0[18] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 19))) _tmem_load_0[19] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 20))) _tmem_load_0[20] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 21))) _tmem_load_0[21] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 22))) _tmem_load_0[22] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 23))) _tmem_load_0[23] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 24))) _tmem_load_0[24] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 25))) _tmem_load_0[25] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 26))) _tmem_load_0[26] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 27))) _tmem_load_0[27] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 28))) _tmem_load_0[28] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 29))) _tmem_load_0[29] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 30))) _tmem_load_0[30] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 31))) _tmem_load_0[31] = -CAKE_INF; + uint32_t _slice_lo_mask_1; + { + int _lim_1 = valid_cols - 32; + if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) _tmem_load_0[32] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) _tmem_load_0[33] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) _tmem_load_0[34] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) _tmem_load_0[35] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) _tmem_load_0[36] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) _tmem_load_0[37] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) _tmem_load_0[38] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) _tmem_load_0[39] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) _tmem_load_0[40] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) _tmem_load_0[41] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) _tmem_load_0[42] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) _tmem_load_0[43] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) _tmem_load_0[44] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) _tmem_load_0[45] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) _tmem_load_0[46] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) _tmem_load_0[47] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) _tmem_load_0[48] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) _tmem_load_0[49] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) _tmem_load_0[50] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) _tmem_load_0[51] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) _tmem_load_0[52] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) _tmem_load_0[53] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) _tmem_load_0[54] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) _tmem_load_0[55] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) _tmem_load_0[56] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) _tmem_load_0[57] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) _tmem_load_0[58] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) _tmem_load_0[59] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) _tmem_load_0[60] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) _tmem_load_0[61] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) _tmem_load_0[62] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) _tmem_load_0[63] = -CAKE_INF; + int* index_row = swa_indices + (query_idx * swa_index_stride + tile * 128 + col_half * 64); + #pragma unroll + for (int chunk = 0; chunk < 64; chunk += 4) { + if (valid_cols > chunk) { + int _vec_load_1[4]; + { + const int4* _ivptr_2 = reinterpret_cast(index_row + chunk + 0); + int4 _ivld_2; + _ivld_2 = *_ivptr_2; + _vec_load_1[0 + 0] = _ivld_2.x; + _vec_load_1[0 + 1] = _ivld_2.y; + _vec_load_1[0 + 2] = _ivld_2.z; + _vec_load_1[0 + 3] = _ivld_2.w; + } + #pragma unroll + for (int i = 0; i < 4; i++) { + if (_vec_load_1[i] < 0) { + _tmem_load_0[chunk + i] = -CAKE_INF; + } + } + } + } + float2 _reg_reduce_max2_3 = {-CAKE_INF, -CAKE_INF}; + row_max_x32_accum(&_tmem_load_0[0], _reg_reduce_max2_3); + row_max_x32_accum(&_tmem_load_0[32], _reg_reduce_max2_3); + float _tmem_load_0_max = row_max_reduce(_reg_reduce_max2_3); + float tile_max = _tmem_load_0_max; + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, tile_max, 16); + float _max_3 = max_noftz(tile_max, _shfl_xor_0); + tile_max = _max_3; + float _max_4 = max_noftz(row_max, tile_max); + float new_max = _max_4; + float _fma_0 = __fmaf_rn(row_max, softmax_scale_log2, (-new_max) * softmax_scale_log2); + float delta = _fma_0; + float _exp2_0 = approx_exp2(delta); + float acc_scale = ((row_max > -CAKE_INF) ? _exp2_0 : 1.0f); + float safe_max = ((new_max == -CAKE_INF) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + const float2 _fma_b2_4 = {softmax_scale_log2, softmax_scale_log2}; + const float2 _fma_c2_5 = {-max_scaled, -max_scaled}; + #pragma unroll + for (int _lf = 0; _lf < 32; _lf++) + fma_f32x2_inplace(&reinterpret_cast(_tmem_load_0)[_lf], _fma_b2_4, _fma_c2_5); + #pragma unroll + for (int _le = 0; _le < 64; _le++) { + _tmem_load_0[_le] = approx_exp2(_tmem_load_0[_le]); + } + unsigned int packed_p[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_0[_lp*2 + 0], _tmem_load_0[_lp*2+1 + 0])); + packed_p[_lp] = *(uint32_t*)&_bf2; + } + int p_addr = taddr + (unsigned int)score_col + 64 + (unsigned int)(tmem_row_origin << 16); + asm volatile( + "tcgen05.st.sync.aligned.16x32bx2.x16.b32" + " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(p_addr), "r"(*reinterpret_cast(&packed_p[0])), "r"(*reinterpret_cast(&packed_p[1])), "r"(*reinterpret_cast(&packed_p[2])), "r"(*reinterpret_cast(&packed_p[3])), "r"(*reinterpret_cast(&packed_p[4])), "r"(*reinterpret_cast(&packed_p[5])), "r"(*reinterpret_cast(&packed_p[6])), "r"(*reinterpret_cast(&packed_p[7])), "r"(*reinterpret_cast(&packed_p[8])), "r"(*reinterpret_cast(&packed_p[9])), "r"(*reinterpret_cast(&packed_p[10])), "r"(*reinterpret_cast(&packed_p[11])), "r"(*reinterpret_cast(&packed_p[12])), "r"(*reinterpret_cast(&packed_p[13])), "r"(*reinterpret_cast(&packed_p[14])), "r"(*reinterpret_cast(&packed_p[15]))); + asm volatile( + "tcgen05.st.sync.aligned.16x32bx2.x16.b32" + " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(p_addr + 16), "r"(*reinterpret_cast(&(packed_p + 16)[0])), "r"(*reinterpret_cast(&(packed_p + 16)[1])), "r"(*reinterpret_cast(&(packed_p + 16)[2])), "r"(*reinterpret_cast(&(packed_p + 16)[3])), "r"(*reinterpret_cast(&(packed_p + 16)[4])), "r"(*reinterpret_cast(&(packed_p + 16)[5])), "r"(*reinterpret_cast(&(packed_p + 16)[6])), "r"(*reinterpret_cast(&(packed_p + 16)[7])), "r"(*reinterpret_cast(&(packed_p + 16)[8])), "r"(*reinterpret_cast(&(packed_p + 16)[9])), "r"(*reinterpret_cast(&(packed_p + 16)[10])), "r"(*reinterpret_cast(&(packed_p + 16)[11])), "r"(*reinterpret_cast(&(packed_p + 16)[12])), "r"(*reinterpret_cast(&(packed_p + 16)[13])), "r"(*reinterpret_cast(&(packed_p + 16)[14])), "r"(*reinterpret_cast(&(packed_p + 16)[15]))); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + if (elect_sync()) { + mbarrier_arrive(p_full_addr + (phase) * 8); + } + float2 _reg_reduce_sum2_6 = make_float2(0.0f, 0.0f); + softmax_block_sum(&_tmem_load_0[0], &_reg_reduce_sum2_6); + softmax_block_sum(&_tmem_load_0[32], &_reg_reduce_sum2_6); + float _tmem_load_0_sum = _reg_reduce_sum2_6.x + _reg_reduce_sum2_6.y; + float block_sum = _tmem_load_0_sum; + float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, block_sum, 16); + block_sum = block_sum + _shfl_xor_1; + float _fma_1 = __fmaf_rn(row_sum, acc_scale, block_sum); + row_sum = _fma_1; + row_max = new_max; + } + } + if (sparse_topk == 128) { + mbarrier_wait(pv_done_addr, 0); + asm volatile("tcgen05.fence::after_thread_sync;"); + float output_scale = bmm2_scale[0]; + float _rcp_0 = approx_rcp(row_sum); + float inv_sum = ((row_sum > 0.0f) ? _rcp_0 : 0.0f); + int output_base = (query_idx * num_heads + my_row) * 512 + blockIdx.y * 256; + #pragma unroll + for (int v_stage = 0; v_stage < 2; v_stage++) { + int output_col = 256 + v_stage * 128; + int output_addr = taddr + (unsigned int)output_col + (unsigned int)(tmem_row_origin << 16); + float _tmem_load_1[64]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_1[0])), "=r"(*reinterpret_cast(&_tmem_load_1[1])), "=r"(*reinterpret_cast(&_tmem_load_1[2])), "=r"(*reinterpret_cast(&_tmem_load_1[3])), "=r"(*reinterpret_cast(&_tmem_load_1[4])), "=r"(*reinterpret_cast(&_tmem_load_1[5])), "=r"(*reinterpret_cast(&_tmem_load_1[6])), "=r"(*reinterpret_cast(&_tmem_load_1[7])), "=r"(*reinterpret_cast(&_tmem_load_1[8])), "=r"(*reinterpret_cast(&_tmem_load_1[9])), "=r"(*reinterpret_cast(&_tmem_load_1[10])), "=r"(*reinterpret_cast(&_tmem_load_1[11])), "=r"(*reinterpret_cast(&_tmem_load_1[12])), "=r"(*reinterpret_cast(&_tmem_load_1[13])), "=r"(*reinterpret_cast(&_tmem_load_1[14])), "=r"(*reinterpret_cast(&_tmem_load_1[15])), "=r"(*reinterpret_cast(&_tmem_load_1[16])), "=r"(*reinterpret_cast(&_tmem_load_1[17])), "=r"(*reinterpret_cast(&_tmem_load_1[18])), "=r"(*reinterpret_cast(&_tmem_load_1[19])), "=r"(*reinterpret_cast(&_tmem_load_1[20])), "=r"(*reinterpret_cast(&_tmem_load_1[21])), "=r"(*reinterpret_cast(&_tmem_load_1[22])), "=r"(*reinterpret_cast(&_tmem_load_1[23])), "=r"(*reinterpret_cast(&_tmem_load_1[24])), "=r"(*reinterpret_cast(&_tmem_load_1[25])), "=r"(*reinterpret_cast(&_tmem_load_1[26])), "=r"(*reinterpret_cast(&_tmem_load_1[27])), "=r"(*reinterpret_cast(&_tmem_load_1[28])), "=r"(*reinterpret_cast(&_tmem_load_1[29])), "=r"(*reinterpret_cast(&_tmem_load_1[30])), "=r"(*reinterpret_cast(&_tmem_load_1[31])) + : "r"(output_addr)); + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_1[32])), "=r"(*reinterpret_cast(&_tmem_load_1[33])), "=r"(*reinterpret_cast(&_tmem_load_1[34])), "=r"(*reinterpret_cast(&_tmem_load_1[35])), "=r"(*reinterpret_cast(&_tmem_load_1[36])), "=r"(*reinterpret_cast(&_tmem_load_1[37])), "=r"(*reinterpret_cast(&_tmem_load_1[38])), "=r"(*reinterpret_cast(&_tmem_load_1[39])), "=r"(*reinterpret_cast(&_tmem_load_1[40])), "=r"(*reinterpret_cast(&_tmem_load_1[41])), "=r"(*reinterpret_cast(&_tmem_load_1[42])), "=r"(*reinterpret_cast(&_tmem_load_1[43])), "=r"(*reinterpret_cast(&_tmem_load_1[44])), "=r"(*reinterpret_cast(&_tmem_load_1[45])), "=r"(*reinterpret_cast(&_tmem_load_1[46])), "=r"(*reinterpret_cast(&_tmem_load_1[47])), "=r"(*reinterpret_cast(&_tmem_load_1[48])), "=r"(*reinterpret_cast(&_tmem_load_1[49])), "=r"(*reinterpret_cast(&_tmem_load_1[50])), "=r"(*reinterpret_cast(&_tmem_load_1[51])), "=r"(*reinterpret_cast(&_tmem_load_1[52])), "=r"(*reinterpret_cast(&_tmem_load_1[53])), "=r"(*reinterpret_cast(&_tmem_load_1[54])), "=r"(*reinterpret_cast(&_tmem_load_1[55])), "=r"(*reinterpret_cast(&_tmem_load_1[56])), "=r"(*reinterpret_cast(&_tmem_load_1[57])), "=r"(*reinterpret_cast(&_tmem_load_1[58])), "=r"(*reinterpret_cast(&_tmem_load_1[59])), "=r"(*reinterpret_cast(&_tmem_load_1[60])), "=r"(*reinterpret_cast(&_tmem_load_1[61])), "=r"(*reinterpret_cast(&_tmem_load_1[62])), "=r"(*reinterpret_cast(&_tmem_load_1[63])) + : "r"(output_addr + 32)); + int col_base = v_stage * 128 + col_half * 64; + if (my_row < num_heads) { + #pragma unroll + for (int offset = 0; offset < 64; offset += 8) { + { + const float2 _prescale2_7 = {inv_sum * output_scale, inv_sum * output_scale}; + #if __CUDA_ARCH__ >= 1000 + #pragma unroll + for (int _ps = 0; _ps < 4; _ps++) + mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_1[offset])[_ps], _prescale2_7); + #else + #pragma unroll + for (int _ps = 0; _ps < 8; _ps++) + _tmem_load_1[offset + _ps] *= inv_sum * output_scale; + #endif + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_tmem_load_1[offset + 0], _tmem_load_1[offset + 1]); + _pk[1] = __floats2bfloat162_rn(_tmem_load_1[offset + 2], _tmem_load_1[offset + 3]); + _pk[2] = __floats2bfloat162_rn(_tmem_load_1[offset + 4], _tmem_load_1[offset + 5]); + _pk[3] = __floats2bfloat162_rn(_tmem_load_1[offset + 6], _tmem_load_1[offset + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(O + (output_base + col_base + offset)))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + } + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + if (elect_sync()) { + mbarrier_arrive(tmem_dealloc_addr); + } + } + } + // ---- Role: mma_warp ---- + if (warp == 4) { + { // mma_warp_main + int query_idx_1 = blockIdx.x; + int _max_1 = ((sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) : (0)); + int _min_1 = ((_max_1) < (sparse_topk) ? (_max_1) : (sparse_topk)); + int active_topk_1 = _min_1; + int tile_count_1 = (active_topk_1 + 128 - 1) / 128; + unsigned int _phase_q_full_0 = 0; + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int tile_1 = 0; tile_1 < 1; tile_1++) { + if (tile_count_1 > tile_1) { + int phase_1 = tile_1 & 1; + int score_col_1 = ((phase_1 != 0) ? 128 : 0); + #pragma unroll + for (int k_stage = 0; k_stage < 4; k_stage++) { + mbarrier_wait(kv_full_addr + (k_stage) * 8, phase_1); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_0 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (k_stage * 2) * 512); + int _mma_b_lo_0 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (k_stage) * 2048); + { + uint64_t _mma_ss_a_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_0); + uint64_t _mma_ss_b_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_0); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, ((k_stage == 0) ? 0 : 1)); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); + } + } + int _mma_a_lo_1 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (k_stage * 2 + 1) * 512); + int _mma_b_lo_1 = make_warp_uniform((((smem_kv_addr + 16384) >> 4) & 0x3FFF) + (k_stage) * 2048); + { + uint64_t _mma_ss_a_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_1); + uint64_t _mma_ss_b_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_1); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); + } + } + if (k_stage == 3) { + elect_commit(s_full_addr + (phase_1) * 8); + } + } + int wait_phase_1 = tile_1 >> 1 & 1; + mbarrier_wait(p_full_addr + (phase_1) * 8, wait_phase_1); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int v_stage_1 = 0; v_stage_1 < 2; v_stage_1++) { + int _mma_b_lo_2 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (blockIdx.y * 2 + v_stage_1) * 2048); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 dhi, blo, ta, id;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 69272720;\n\t" + "mov.b32 ta, %2;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p0;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "}\n" + :: "r"((tmem_tmem + (256 + v_stage_1 * 128))), "r"(_mma_b_lo_2), "r"(tmem_tmem + (score_col_1 + 64)), "r"(0)); + } + elect_commit(pv_done_addr); + } + } + unsigned int _phase_tmem_dealloc_0 = 0; + mbarrier_wait(tmem_dealloc_addr, _phase_tmem_dealloc_0); + _phase_tmem_dealloc_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + } + // ---- Role: load_warp ---- + if (warp >= 5 && warp <= 7) { + { // load_warp_main + const int load_warp_rank = warp - 5; + int query_idx_2 = blockIdx.x; + int _max_0 = ((sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) : (0)); + int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); + int active_topk_2 = _min_0; + int tile_count_2 = (active_topk_2 + 128 - 1) / 128; + unsigned int load_kv_stage = 0; + unsigned int load_kv_phase = 1; + int group = load_warp_rank + lane * 3; + int row0 = 0; + int row1 = 0; + int row2 = 0; + int row3 = 0; + if (group < 32) { + int index_base = group * 4; + int* index_row_1 = swa_indices + (query_idx_2 * swa_index_stride); + int raw_rows[4]; + if (index_base + 3 < sparse_topk) { + int _vec_load_0[4]; + { + const int4* _ivptr_0 = reinterpret_cast(index_row_1 + index_base + 0); + int4 _ivld_0; + _ivld_0 = *_ivptr_0; + _vec_load_0[0 + 0] = _ivld_0.x; + _vec_load_0[0 + 1] = _ivld_0.y; + _vec_load_0[0 + 2] = _ivld_0.z; + _vec_load_0[0 + 3] = _ivld_0.w; + } + #pragma unroll + for (int i_1 = 0; i_1 < 4; i_1++) { + raw_rows[i_1] = _vec_load_0[i_1]; + } + } else { + #pragma unroll + for (int i_2 = 0; i_2 < 4; i_2++) { + raw_rows[i_2] = -1; + if (index_base + i_2 < sparse_topk) { + raw_rows[i_2] = index_row_1[index_base + i_2]; + } + } + } + row0 = ((raw_rows[0] >= 0) ? raw_rows[0] : 0); + row1 = ((raw_rows[1] >= 0) ? raw_rows[1] : 0); + row2 = ((raw_rows[2] >= 0) ? raw_rows[2] : 0); + row3 = ((raw_rows[3] >= 0) ? raw_rows[3] : 0); + } + if (load_warp_rank == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(q_full_addr, 65536); + #pragma unroll + for (int q_stage = 0; q_stage < 8; q_stage += 8) { + tma_4d_gmem2smem(smem_q_addr + (unsigned int)(q_stage * 8192), tmap_q, 0, 0, q_stage, query_idx_2, q_full_addr); + } + } + } + #pragma unroll + for (int tile_2 = 0; tile_2 < 1; tile_2++) { + if (tile_count_2 > tile_2) { + #pragma unroll + for (int k_stage_1 = 0; k_stage_1 < 4; k_stage_1++) { + if (elect_sync()) { + int producer_groups = ((load_warp_rank < 2) ? 11 : 10); + mbarrier_arrive_expect_tx(kv_full_addr + (load_kv_stage) * 8, producer_groups * 512 * 2); + } + if (group < 32) { + int dst_k = smem_kv_addr + load_kv_stage * 32768; + if (tile_2 == 0) { + tma_gather4_gmem2smem(dst_k + group * 512, tmap_swa_kv, k_stage_1 * 128, row0, row1, row2, row3, kv_full_addr + (load_kv_stage) * 8); + tma_gather4_gmem2smem(dst_k + 16384 + group * 512, tmap_swa_kv, k_stage_1 * 128 + 64, row0, row1, row2, row3, kv_full_addr + (load_kv_stage) * 8); + } else { + tma_gather4_gmem2smem(dst_k + group * 512, tmap_compressed_kv, k_stage_1 * 128, row0, row1, row2, row3, kv_full_addr + (load_kv_stage) * 8); + tma_gather4_gmem2smem(dst_k + 16384 + group * 512, tmap_compressed_kv, k_stage_1 * 128 + 64, row0, row1, row2, row3, kv_full_addr + (load_kv_stage) * 8); + } + } + load_kv_stage += 1; + if (load_kv_stage == 4) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + } + if (elect_sync()) { + mbarrier_arrive(tmem_dealloc_addr); + } + } + } + + // Cleanup +} + +} // extern "C" diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_4ab0487d5fde70d4569e_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_4ab0487d5fde70d4569e_binding.cu new file mode 100644 index 00000000000..9e594768617 --- /dev/null +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_4ab0487d5fde70d4569e_binding.cu @@ -0,0 +1,606 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +// Generated by the Cake source exporter — do not edit. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_4ab0487d5fde70d4569e'. +#include +#include +#include +#include + +#include "tvm_ffi_utils.h" + +#include +#include +#include +#include +#include +#include +#include + +extern "C" __global__ void kernel_cake_dsv4_4ab0487d5fde70d4569e(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, const __grid_constant__ CUtensorMap O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); + + +namespace cake_host_shim_df4c5e4bdb73dbc3 { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t CakeDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +// 4D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + int64_t d3 = t.size(t.ndim() - 3); + TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(d3)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 64u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_swa_k' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_swa_k(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_swa_k' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_swa_k' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_swa_k' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_swa_k"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_swa_k' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_swa_k' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_k'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_swa_k') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_compressed_k' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_compressed_k(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_compressed_k' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_compressed_k' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_compressed_k' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_compressed_k"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_compressed_k' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_compressed_k' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_k'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_k') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_swa_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_swa_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_swa_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_swa_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_swa_v' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_swa_v"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_swa_v' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_swa_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_v'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_swa_v') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_compressed_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_compressed_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_compressed_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_compressed_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_compressed_v' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_compressed_v"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_compressed_v' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_compressed_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_v'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_v') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'O' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_O(const TensorView& t, int64_t tp_num_heads, int64_t tp_num_query_tokens) { + TVM_FFI_CHECK(t.ndim() >= 1, ValueError) + << "TMA source 'O' must have at least 1 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'O' must have unit innermost stride, got " << t.stride(-1); + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(8), (uint64_t)(5), (uint64_t)((tp_num_query_tokens * tp_num_heads))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'O' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1] && 1u <= global_dim[2] && 32u <= global_dim[3], ValueError) + << "TMA box (64, 1, 1, 32) exceeds resolved global dims for 'O'"; + int64_t carrier_stride_0 = 64; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 512; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = 2560; + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 1u, 1u, 32u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'O') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_compressed_k, TensorView arg_tmap_swa_v, TensorView arg_tmap_compressed_v, TensorView arg_O, TensorView arg_partial_lse, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_num_query_tokens, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckContiguous(arg_tmap_q, "tmap_q"); + CheckCudaTensor(arg_tmap_swa_k, "tmap_swa_k"); + CheckDtype(arg_tmap_swa_k, "tmap_swa_k", 4, 16, 1); + CheckCudaTensor(arg_tmap_compressed_k, "tmap_compressed_k"); + CheckDtype(arg_tmap_compressed_k, "tmap_compressed_k", 4, 16, 1); + CheckCudaTensor(arg_tmap_swa_v, "tmap_swa_v"); + CheckDtype(arg_tmap_swa_v, "tmap_swa_v", 4, 16, 1); + CheckCudaTensor(arg_tmap_compressed_v, "tmap_compressed_v"); + CheckDtype(arg_tmap_compressed_v, "tmap_compressed_v", 4, 16, 1); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_partial_lse, "partial_lse"); + CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); + CheckContiguous(arg_partial_lse, "partial_lse"); + CheckCudaTensor(arg_swa_indices, "swa_indices"); + CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); + CheckContiguous(arg_swa_indices, "swa_indices"); + CheckCudaTensor(arg_compressed_indices, "compressed_indices"); + CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); + CheckContiguous(arg_compressed_indices, "compressed_indices"); + CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); + CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); + CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); + CheckContiguous(arg_bmm1_scale, "bmm1_scale"); + CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); + CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); + CheckContiguous(arg_bmm2_scale, "bmm2_scale"); + TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) + << "scalar 'num_heads' value " << arg_num_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_query_tokens >= -2147483648LL && arg_num_query_tokens <= 2147483647LL, ValueError) + << "scalar 'num_query_tokens' value " << arg_num_query_tokens + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) + << "scalar 'swa_index_stride' value " << arg_swa_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) + << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) + << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) + << "scalar 'sparse_topk' value " << arg_sparse_topk + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) + << "scalar 'has_sinks' value " << arg_has_sinks + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_swa_k, arg_tmap_q, "tmap_swa_k", "tmap_q"); + CheckSameCudaDevice(arg_tmap_compressed_k, arg_tmap_q, "tmap_compressed_k", "tmap_q"); + CheckSameCudaDevice(arg_tmap_swa_v, arg_tmap_q, "tmap_swa_v", "tmap_q"); + CheckSameCudaDevice(arg_tmap_compressed_v, arg_tmap_q, "tmap_compressed_v", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_partial_lse, arg_tmap_q, "partial_lse", "tmap_q"); + CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); + CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); + CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); + CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") must be divisible by cluster dims (2, 1, 1)"; + + cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_swa_k = EncodeTma_tmap_swa_k(arg_tmap_swa_k); + CUtensorMap p_tmap_compressed_k = EncodeTma_tmap_compressed_k(arg_tmap_compressed_k); + CUtensorMap p_tmap_swa_v = EncodeTma_tmap_swa_v(arg_tmap_swa_v); + CUtensorMap p_tmap_compressed_v = EncodeTma_tmap_compressed_v(arg_tmap_compressed_v); + CUtensorMap p_O = EncodeTma_O(arg_O, (int64_t)arg_num_heads, (int64_t)arg_num_query_tokens); + float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); + int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); + int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); + int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); + float* p_sinks = static_cast(arg_sinks.data_ptr()); + float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); + float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); + int32_t v_num_heads = (int32_t)arg_num_heads; + int32_t v_num_query_tokens = (int32_t)arg_num_query_tokens; + int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; + int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; + int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; + int32_t v_sparse_topk = (int32_t)arg_sparse_topk; + int32_t v_has_sinks = (int32_t)arg_has_sinks; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + void* kargs[] = {&p_tmap_q, &p_tmap_swa_k, &p_tmap_compressed_k, &p_tmap_swa_v, &p_tmap_compressed_v, &p_O, &p_partial_lse, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_num_query_tokens, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_has_sinks, &v_total_work_items}; + + static std::mutex smem_mu; + static auto* smem_status_by_device = + new std::unordered_map(); + cudaError_t smem_status = cudaSuccess; + { + std::lock_guard lock(smem_mu); + auto it = smem_status_by_device->find(dev.device_id); + if (it == smem_status_by_device->end()) { + int current_device = -1; + smem_status = cudaGetDevice(¤t_device); + if (smem_status == cudaSuccess && current_device != dev.device_id) { + smem_status = cudaErrorInvalidDevice; + } + if (smem_status == cudaSuccess) { + smem_status = cudaFuncSetAttribute( + reinterpret_cast(kernel_cake_dsv4_4ab0487d5fde70d4569e), + cudaFuncAttributeMaxDynamicSharedMemorySize, + 231936); + } + smem_status_by_device->emplace(dev.device_id, smem_status); + } else { + smem_status = it->second; + } + } + TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) + << "cudaFuncSetAttribute for kernel_cake_dsv4_4ab0487d5fde70d4569e failed: " + << cudaGetErrorString(smem_status); + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(512u, 1u, 1u); + + // Extended direct-source launch for a linked CUDA kernel symbol. + cudaLaunchAttribute attrs[2]{}; + int n = 0; + attrs[n].id = cudaLaunchAttributeClusterDimension; + attrs[n].val.clusterDim.x = 2u; + attrs[n].val.clusterDim.y = 1u; + attrs[n].val.clusterDim.z = 1u; + ++n; + attrs[n].id = cudaLaunchAttributeClusterSchedulingPolicyPreference; + attrs[n].val.clusterSchedulingPolicyPreference = cudaClusterSchedulingPolicySpread; + ++n; + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes = 231936u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; + cudaError_t launch_status = cudaLaunchKernelExC( + &config, reinterpret_cast(kernel_cake_dsv4_4ab0487d5fde70d4569e), kargs); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernelExC for kernel_cake_dsv4_4ab0487d5fde70d4569e failed: " + << cudaGetErrorString(launch_status); +} + +} // namespace cake_host_shim_df4c5e4bdb73dbc3 + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_df4c5e4bdb73dbc3::Run); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_4ab0487d5fde70d4569e_kernel.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_4ab0487d5fde70d4569e_kernel.cu new file mode 100644 index 00000000000..17aca3f4231 --- /dev/null +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_4ab0487d5fde70d4569e_kernel.cu @@ -0,0 +1,2418 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define CAKE_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_SCRATCH_OFFSET 0 +#define NUM_K_PIPE_STAGES 8 +#define NUM_V_PIPE_STAGES 2 +#define SMEM_SMEM_Q_OFF 0 +#define SMEM_SMEM_Q_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_STRIDE 8192 +#define SMEM_SMEM_K_OFF 65536 +#define SMEM_SMEM_K_STAGE_BYTES 16384 +#define SMEM_SMEM_K_STRIDE 16384 +#define SMEM_SMEM_V_OFF 65536 +#define SMEM_SMEM_V_STAGE_BYTES 16384 +#define SMEM_SMEM_V_STRIDE 16384 +#define SMEM_SMEM_P_OFF 196608 +#define SMEM_SMEM_P_STAGE_BYTES 8192 +#define SMEM_SMEM_P_STRIDE 8192 +#define SMEM_SMEM_INDICES_OFF 229376 +#define SMEM_SMEM_INDICES_STAGE_BYTES 1024 +#define SMEM_SMEM_INDICES_STRIDE 1024 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_OFF 230400 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STRIDE 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_OFF 230912 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STRIDE 512 +#define SMEM_SMEM_INDEX_VALID_OFF 231424 +#define SMEM_SMEM_INDEX_VALID_STAGE_BYTES 64 +#define SMEM_SMEM_INDEX_VALID_STRIDE 64 +#define SMEM_TOTAL 231936 +#define THREADS 512 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count) : "memory"); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16_cta2( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" + "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" + "}\n" + :: "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x4(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x4.b32" + " {%0, %1, %2, %3}, [%4];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); +} + + +__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ float2 ex2_emulation_f32x2_value(float2 value) { + const float c0 = 1.0f, c1 = 0.695146143436431884765625f; + const float c2 = 0.227564394474029541015625f, c3 = 0.077119089663028717041015625f; + const float magic = 12582912.0f; + float x0 = max_noftz(value.x, -127.0f), x1 = max_noftz(value.y, -127.0f); + float2 xc2 = make_float2(x0, x1), magic2 = make_float2(magic, magic); + float2 xr2; + asm("add.rm.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xr2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&magic2)); + float2 c3_2 = make_float2(c3, c3), c2_2 = make_float2(c2, c2); + float2 c1_2 = make_float2(c1, c1), c0_2 = make_float2(c0, c0); + float2 xrb2, xfrac2; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xrb2) + : "l"(*(unsigned long long*)&xr2), "l"(*(unsigned long long*)&magic2)); + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xfrac2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&xrb2)); + float2 poly2; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&c3_2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c2_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c1_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c0_2)); + int x0r_i, x1r_i, p0_i, p1_i; + asm("mov.b64 {%0, %1}, %2;" : "=r"(x0r_i), "=r"(x1r_i) : "l"(*(unsigned long long*)&xr2)); + asm("mov.b64 {%0, %1}, %2;" : "=r"(p0_i), "=r"(p1_i) : "l"(*(unsigned long long*)&poly2)); + float r0, r1; + asm("mov.b32 %0, %1;" : "=f"(r0) : "r"((x0r_i << 23) + p0_i)); + asm("mov.b32 %0, %1;" : "=f"(r1) : "r"((x1r_i << 23) + p1_i)); + return make_float2(r0, r1); +} + +__device__ __forceinline__ void ex2_emulation_f32x2(float* x0_ptr, float* x1_ptr) { + float2 result = ex2_emulation_f32x2_value(make_float2(*x0_ptr, *x1_ptr)); + *x0_ptr = result.x; *x1_ptr = result.y; +} + +__device__ __forceinline__ void softmax_frag_exp2_cast( + float* sv, uint32_t* pv, int use_emu) +{ + #pragma unroll + for (int j = 0; j < 16; j++) { + if (use_emu && j >= 12) + ex2_emulation_f32x2(&sv[j*2], &sv[j*2+1]); + else { + sv[j*2] = approx_exp2(sv[j*2]); + sv[j*2+1] = approx_exp2(sv[j*2+1]); + } + } + #pragma unroll + for (int j = 0; j < 16; j++) { + __nv_bfloat162 bf = __float22bfloat162_rn({sv[j*2], sv[j*2+1]}); + pv[j] = reinterpret_cast(bf); + } +} + + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_2d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3}], [%4];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr) { + // Canonical .shared::cta form for non-multicast gather4, matching + // trtllm-gen / cuda_ptx and the PTX ISA qualifier order + // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, + // .shared::cluster is reserved for the multicast variant (ctaMask). + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem_cta2( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr) { + // Canonical .shared::cta form; see tma_gather4_gmem2smem above. + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem_mc( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr, unsigned short cta_mask) { + // Multicast variant: the PTX grammar ties the .shared::cluster + // destination to .multicast::cluster + ctaMask (cf. cuda_ptx / + // SM100_TMA_LOAD_MULTICAST_2D_GATHER4 in CUTLASS). + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" + ".mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem_mc_cta2( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr, unsigned short cta_mask) { + // Multicast + cta_group::2 variant; see tma_gather4_gmem2smem_mc. + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" + ".mbarrier::complete_tx::bytes.multicast::cluster.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void tma_store_4d( + const void *tmap, int x, int y, int z, int w, unsigned smem_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.global.shared::cta.tile.bulk_group" + " [%0, {%1, %2, %3, %4}], [%5];" + :: "l"(tmap), "r"(x), "r"(y), "r"(z), "r"(w), "r"(smem_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %1;\n\t" + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" + "}\n" + :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); +} + +extern "C" { + +__global__ __launch_bounds__(512, 1) __cluster_dims__(2,1,1) void +kernel_cake_dsv4_4ab0487d5fde70d4569e(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, const __grid_constant__ CUtensorMap O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int mbar_base = smem + 231488; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define k_full_addr (mbar_base + 16) + #define k_empty_addr (mbar_base + 80) + #define slot_empty_addr (mbar_base + 144) + #define v_full_addr (mbar_base + 208) + #define v_empty_addr (mbar_base + 224) + #define s_full_addr (mbar_base + 240) + #define s_empty_addr (mbar_base + 256) + #define p_full_addr (mbar_base + 272) + #define p_empty_addr (mbar_base + 288) + #define o_empty_addr (mbar_base + 304) + #define stats_addr (mbar_base + 312) + #define stats_empty_addr (mbar_base + 328) + #define o_full_addr (mbar_base + 344) + #define epi_q_free_addr (mbar_base + 352) + #define o_first_slice_seeded_addr (mbar_base + 360) + #define tmem_dealloc_peer_addr (mbar_base + 368) + #define index_valid_full_addr (mbar_base + 376) + #define index_valid_empty_addr (mbar_base + 392) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 2; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + + // Kernel setup ops + __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 0); + const int smem_q_addr = smem + 0; + __nv_bfloat16* smem_k = reinterpret_cast<__nv_bfloat16*>(smem_raw + 65536); + const int smem_k_addr = smem + 65536; + __nv_bfloat16* smem_v = reinterpret_cast<__nv_bfloat16*>(smem_raw + 65536); + const int smem_v_addr = smem + 65536; + __nv_bfloat16* smem_p = reinterpret_cast<__nv_bfloat16*>(smem_raw + 196608); + const int smem_p_addr = smem + 196608; + int* smem_indices = reinterpret_cast(smem_raw + 229376); + const int smem_indices_addr = smem + 229376; + float* smem_softmax_exchange = reinterpret_cast(smem_raw + 230400); + const int smem_softmax_exchange_addr = smem + 230400; + float* smem_epilogue_exchange = reinterpret_cast(smem_raw + 230912); + const int smem_epilogue_exchange_addr = smem + 230912; + unsigned int* smem_index_valid = reinterpret_cast(smem_raw + 231424); + const int smem_index_valid_addr = smem + 231424; + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // Mbarrier init (20 pipeline groups, 0 ordered-sequence groups, 51 barriers) + // Mbarriers at smem_raw[231488..231896) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // q_full: 1 barriers, init_count=1 + mbarrier_init(smem + 231488, 1); + // q_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 231496, 1); + // --- pipeline 'k_pipe' --- + // k_full: 8 barriers, init_count=1 + mbarrier_init(smem + 231504, 1); + mbarrier_init(smem + 231512, 1); + mbarrier_init(smem + 231520, 1); + mbarrier_init(smem + 231528, 1); + mbarrier_init(smem + 231536, 1); + mbarrier_init(smem + 231544, 1); + mbarrier_init(smem + 231552, 1); + mbarrier_init(smem + 231560, 1); + // k_empty: 8 barriers, init_count=1 + mbarrier_init(smem + 231568, 1); + mbarrier_init(smem + 231576, 1); + mbarrier_init(smem + 231584, 1); + mbarrier_init(smem + 231592, 1); + mbarrier_init(smem + 231600, 1); + mbarrier_init(smem + 231608, 1); + mbarrier_init(smem + 231616, 1); + mbarrier_init(smem + 231624, 1); + // slot_empty: 8 barriers, init_count=1 + mbarrier_init(smem + 231632, 1); + mbarrier_init(smem + 231640, 1); + mbarrier_init(smem + 231648, 1); + mbarrier_init(smem + 231656, 1); + mbarrier_init(smem + 231664, 1); + mbarrier_init(smem + 231672, 1); + mbarrier_init(smem + 231680, 1); + mbarrier_init(smem + 231688, 1); + // --- pipeline 'v_pipe' --- + // v_full: 2 barriers, init_count=1 + mbarrier_init(smem + 231696, 1); + mbarrier_init(smem + 231704, 1); + // v_empty: 2 barriers, init_count=1 + mbarrier_init(smem + 231712, 1); + mbarrier_init(smem + 231720, 1); + // s_full: 2 barriers, init_count=1 + mbarrier_init(smem + 231728, 1); + mbarrier_init(smem + 231736, 1); + // s_empty: 2 barriers, init_count=256 + mbarrier_init(smem + 231744, 256); + mbarrier_init(smem + 231752, 256); + // p_full: 2 barriers, init_count=256 + mbarrier_init(smem + 231760, 256); + mbarrier_init(smem + 231768, 256); + // p_empty: 2 barriers, init_count=1 + mbarrier_init(smem + 231776, 1); + mbarrier_init(smem + 231784, 1); + // o_empty: 1 barriers, init_count=256 + mbarrier_init(smem + 231792, 256); + // stats: 2 barriers, init_count=128 + mbarrier_init(smem + 231800, 128); + mbarrier_init(smem + 231808, 128); + // stats_empty: 2 barriers, init_count=128 + mbarrier_init(smem + 231816, 128); + mbarrier_init(smem + 231824, 128); + // o_full: 1 barriers, init_count=1 + mbarrier_init(smem + 231832, 1); + // epi_q_free: 1 barriers, init_count=4 + mbarrier_init(smem + 231840, 4); + // o_first_slice_seeded: 1 barriers, init_count=256 + mbarrier_init(smem + 231848, 256); + // tmem_dealloc_peer: 1 barriers, init_count=32 + mbarrier_init(smem + 231856, 32); + // index_valid_full: 2 barriers, init_count=32 + mbarrier_init(smem + 231864, 32); + mbarrier_init(smem + 231872, 32); + // index_valid_empty: 2 barriers, init_count=128 + mbarrier_init(smem + 231880, 128); + mbarrier_init(smem + 231888, 128); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 231896); + if (warp == 0) { + int _tmem_hold = smem + 231896; + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); + } + + __syncthreads(); + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_tmem_scratch = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 15) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 56;"); + } + + // ---- Role: softmax_wg ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); + { // softmax_wg_main + const int softmax_dummy = 0; + float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; + const int head_row_base = warp % 2 * 32; + const int score_row_base = warp % 4 * 32; + const int n_half = warp % 4 / 2; + const int my_row = head_row_base + lane; + const int exchange_idx = n_half * 64 + my_row; + int tile_cursor = 0; + int valid_iter = 0; + float seed_zero[4]; + #pragma unroll + for (int seed_i = 0; seed_i < 4; seed_i++) { + seed_zero[seed_i] = 0.0f; + } + #pragma unroll + for (int seed_stage = 0; seed_stage < 2; seed_stage++) { + #pragma unroll + for (int seed_col = 128 + seed_stage * 128; seed_col < 128 + (seed_stage + 1) * 128; seed_col += 4) { + int seed_addr = taddr + (unsigned int)seed_col + (unsigned int)(score_row_base << 16); + tmem_st_x4_f32(seed_addr, seed_zero); + } + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(o_first_slice_seeded_addr); + int peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_first_slice_seeded_addr), "r"(peer_rank) : "memory"); + #pragma unroll 1 + for (unsigned int work_idx = cluster_id; work_idx < total_work_items; work_idx += num_clusters) { + int split_idx = work_idx % (unsigned int)(((1) ? 5 : 1)); + int query_idx = work_idx / (unsigned int)(((1) ? 5 : 1)); + int _max_0 = ((sparse_topk_lens[query_idx] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx] + sparse_topk_lens_offset) : (0)); + int _min_2 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); + int active_topk = _min_2; + int owner_capacity = ((sparse_topk + 128 - 1) / 128 + ((1) ? 5 : 1) - 1) / ((1) ? 5 : 1); + int first_owner_tile = work_idx % (unsigned int)(((1) ? 5 : 1)) * (unsigned int)owner_capacity; + int _min_3 = ((owner_capacity) < ((sparse_topk + 128 - 1) / 128 - first_owner_tile) ? (owner_capacity) : ((sparse_topk + 128 - 1) / 128 - first_owner_tile)); + int num_kv_tiles = _min_3; + int first_kv_tile = first_owner_tile; + int head_idx = cta_rank * 64 + my_row; + float row_max_val = -CAKE_INF; + float row_sum_val = 0.0f; + if (has_sinks != 0 && head_idx < num_heads && split_idx == 0) { + row_max_val = sinks[head_idx] * 1.4426950408889634f / softmax_scale_log2; + if (n_half == 0) { + row_sum_val = 1.0f; + } + } + int valid_slot = valid_iter & 1; + int valid_full_phase = valid_iter >> 1 & 1; + mbarrier_wait(index_valid_full_addr + (valid_slot) * 8, valid_full_phase); + int valid_word_base = valid_slot * 8 + n_half * 2; + #pragma unroll 1 + for (int tile = 0; tile < num_kv_tiles; tile++) { + int pipeline_tile = tile_cursor + tile; + int score_phase = pipeline_tile & 1; + int s_full_phase = pipeline_tile >> 1 & 1; + mbarrier_wait(s_full_addr + (score_phase) * 8, s_full_phase); + int score_offset = ((score_phase != 0) ? 64 : 0); + int score_base = taddr + (unsigned int)score_offset + (unsigned int)(score_row_base << 16); + float _tmem_load_0[64]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]) + : "r"(score_base)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]) + : "r"(score_base + 32)); + int valid_cols = active_topk - (first_kv_tile + tile) * 128 - n_half * 64; + if (valid_cols < 0) { + valid_cols = 0; + } + if (valid_cols > 64) { + valid_cols = 64; + } + unsigned int keep_lo = smem_index_valid[valid_word_base + tile * 4]; + unsigned int keep_hi = smem_index_valid[valid_word_base + tile * 4 + 1]; + uint32_t _thresh_mask_0; + { + int _lim_0 = valid_cols; + if (_lim_0 <= 0) { _thresh_mask_0 = 0u; } + else if (_lim_0 >= 32) { _thresh_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_thresh_mask_0) : "r"(_lim_0)); + } + } + uint32_t _thresh_mask_1; + { + int _lim_1 = valid_cols - 32; + if (_lim_1 <= 0) { _thresh_mask_1 = 0u; } + else if (_lim_1 >= 32) { _thresh_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_thresh_mask_1) : "r"(_lim_1)); + } + } + if (!(_thresh_mask_0 & keep_lo & (1u << 0))) _tmem_load_0[0] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 1))) _tmem_load_0[1] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 2))) _tmem_load_0[2] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 3))) _tmem_load_0[3] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 4))) _tmem_load_0[4] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 5))) _tmem_load_0[5] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 6))) _tmem_load_0[6] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 7))) _tmem_load_0[7] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 8))) _tmem_load_0[8] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 9))) _tmem_load_0[9] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 10))) _tmem_load_0[10] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 11))) _tmem_load_0[11] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 12))) _tmem_load_0[12] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 13))) _tmem_load_0[13] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 14))) _tmem_load_0[14] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 15))) _tmem_load_0[15] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 16))) _tmem_load_0[16] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 17))) _tmem_load_0[17] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 18))) _tmem_load_0[18] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 19))) _tmem_load_0[19] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 20))) _tmem_load_0[20] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 21))) _tmem_load_0[21] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 22))) _tmem_load_0[22] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 23))) _tmem_load_0[23] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 24))) _tmem_load_0[24] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 25))) _tmem_load_0[25] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 26))) _tmem_load_0[26] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 27))) _tmem_load_0[27] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 28))) _tmem_load_0[28] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 29))) _tmem_load_0[29] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 30))) _tmem_load_0[30] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 31))) _tmem_load_0[31] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 0))) _tmem_load_0[32] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 1))) _tmem_load_0[33] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 2))) _tmem_load_0[34] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 3))) _tmem_load_0[35] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 4))) _tmem_load_0[36] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 5))) _tmem_load_0[37] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 6))) _tmem_load_0[38] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 7))) _tmem_load_0[39] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 8))) _tmem_load_0[40] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 9))) _tmem_load_0[41] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 10))) _tmem_load_0[42] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 11))) _tmem_load_0[43] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 12))) _tmem_load_0[44] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 13))) _tmem_load_0[45] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 14))) _tmem_load_0[46] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 15))) _tmem_load_0[47] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 16))) _tmem_load_0[48] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 17))) _tmem_load_0[49] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 18))) _tmem_load_0[50] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 19))) _tmem_load_0[51] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 20))) _tmem_load_0[52] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 21))) _tmem_load_0[53] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 22))) _tmem_load_0[54] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 23))) _tmem_load_0[55] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 24))) _tmem_load_0[56] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 25))) _tmem_load_0[57] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 26))) _tmem_load_0[58] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 27))) _tmem_load_0[59] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 28))) _tmem_load_0[60] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 29))) _tmem_load_0[61] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 30))) _tmem_load_0[62] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 31))) _tmem_load_0[63] = -CAKE_INF; + float2 _reg_reduce_max2_2 = {-CAKE_INF, -CAKE_INF}; + row_max_x32_accum(&_tmem_load_0[0], _reg_reduce_max2_2); + row_max_x32_accum(&_tmem_load_0[32], _reg_reduce_max2_2); + float _tmem_load_0_max = row_max_reduce(_reg_reduce_max2_2); + float _max_1 = max_noftz(_tmem_load_0_max, row_max_val); + float new_max = _max_1; + smem_softmax_exchange[exchange_idx] = new_max; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + float _max_2 = max_noftz(new_max, smem_softmax_exchange[exchange_idx ^ 64]); + new_max = _max_2; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + float no_correction = (((new_max - row_max_val) * softmax_scale_log2 <= 0.0f) ? 1.0f : 0.0f); + float delta = softmax_scale_log2 * (row_max_val - new_max); + float _exp2_0 = approx_exp2(delta); + float exp_delta = _exp2_0; + float acc_scale = ((row_max_val > -CAKE_INF) ? exp_delta : 1.0f); + row_max_val = new_max; + float safe_max = ((new_max == -CAKE_INF) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + const float2 _fma_b2_3 = {softmax_scale_log2, softmax_scale_log2}; + const float2 _fma_c2_4 = {-max_scaled, -max_scaled}; + #pragma unroll + for (int _lf = 0; _lf < 32; _lf++) + fma_f32x2_inplace(&reinterpret_cast(_tmem_load_0)[_lf], _fma_b2_3, _fma_c2_4); + #pragma unroll + for (int _le = 0; _le < 64; _le++) { + _tmem_load_0[_le] = approx_exp2(_tmem_load_0[_le]); + } + float2 _reg_reduce_sum2_5 = make_float2(0.0f, 0.0f); + softmax_block_sum(&_tmem_load_0[0], &_reg_reduce_sum2_5); + softmax_block_sum(&_tmem_load_0[32], &_reg_reduce_sum2_5); + float _tmem_load_0_sum = _reg_reduce_sum2_5.x + _reg_reduce_sum2_5.y; + row_sum_val = row_sum_val * acc_scale + _tmem_load_0_sum; + int stats_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + mbarrier_wait(stats_empty_addr + (score_phase) * 8, stats_empty_phase); + float meta[4]; + meta[0] = row_sum_val; + meta[1] = row_max_val; + meta[2] = acc_scale; + meta[3] = no_correction; + int meta_addr = taddr + 384 + (unsigned int)(score_phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(head_row_base << 16); + tmem_st_x4_f32(meta_addr, meta); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (score_phase) * 8); + int p_stage = pipeline_tile % 2; + int p_empty_phase = pipeline_tile / 2 & 1 ^ 1; + mbarrier_wait(p_empty_addr + (p_stage) * 8, p_empty_phase); + uint32_t _tmem_load_0_bf16[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_0[_lp*2 + 0], _tmem_load_0[_lp*2+1 + 0])); + _tmem_load_0_bf16[_lp] = *(uint32_t*)&_bf2; + } + int smem_p_stage = p_stage * 2 + n_half; + int p_base_smem = smem_p_addr + (unsigned int)(smem_p_stage * 8192); + #pragma unroll + for (int vec = 0; vec < 8; vec++) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_base_smem + (my_row * 128 + vec * 16 ^ (my_row * 128 + vec * 16 >> 7 & 7) << 4))), "r"(_tmem_load_0_bf16[vec * 4]), "r"(_tmem_load_0_bf16[vec * 4 + 1]), "r"(_tmem_load_0_bf16[vec * 4 + 2]), "r"(_tmem_load_0_bf16[vec * 4 + 3]) : "memory"); + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(p_full_addr + (unsigned int)(p_stage * 8)), "r"(0) : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(s_empty_addr + (unsigned int)(score_phase * 8)), "r"(0) : "memory"); + } + mbarrier_arrive(index_valid_empty_addr + (valid_slot) * 8); + valid_iter = valid_iter + 1; + tile_cursor = tile_cursor + num_kv_tiles; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + } + } + // ---- Role: correction_wg ---- + if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); + { // correction_wg_main + const int correction_dummy = 0; + float softmax_scale_log2_1 = bmm1_scale[0] * 1.4426950408889634f; + float output_scale = bmm2_scale[0]; + const int head_row_base_1 = warp % 2 * 32; + const int output_row_base = warp % 4 * 32; + const int n_half_1 = warp % 4 / 2; + const int my_row_1 = head_row_base_1 + lane; + const int corr_row = head_row_base_1 << 16; + int tile_cursor_1 = 0; + unsigned int _phase_o_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = cluster_id; work_idx_1 < total_work_items; work_idx_1 += num_clusters) { + int split_idx_1 = work_idx_1 % (unsigned int)(((1) ? 5 : 1)); + int query_idx_1 = work_idx_1 / (unsigned int)(((1) ? 5 : 1)); + int owner_capacity_1 = ((sparse_topk + 128 - 1) / 128 + ((1) ? 5 : 1) - 1) / ((1) ? 5 : 1); + int first_owner_tile_1 = work_idx_1 % (unsigned int)(((1) ? 5 : 1)) * (unsigned int)owner_capacity_1; + int _min_4 = ((owner_capacity_1) < ((sparse_topk + 128 - 1) / 128 - first_owner_tile_1) ? (owner_capacity_1) : ((sparse_topk + 128 - 1) / 128 - first_owner_tile_1)); + int num_kv_tiles_1 = _min_4; + float final_sum = 0.0f; + float final_max = -CAKE_INF; + #pragma unroll 1 + for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { + int pipeline_tile_1 = tile_cursor_1 + tile_1; + int score_phase_1 = pipeline_tile_1 & 1; + int stats_phase = pipeline_tile_1 >> 1 & 1; + mbarrier_wait(stats_addr + (score_phase_1) * 8, stats_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int meta_addr_1 = taddr + 384 + (unsigned int)(score_phase_1 * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row; + float _tmem_load_1[4]; + tmem_ld_x4(&_tmem_load_1[0], meta_addr_1); + final_sum = _tmem_load_1[0]; + final_max = _tmem_load_1[1]; + float acc_scale_1 = _tmem_load_1[2]; + float no_correction_1 = _tmem_load_1[3]; + if (tile_1 > 0) { + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _vote_0 = __all_sync(0xFFFFFFFF, no_correction_1 == 1.0f); + int skip_correction = _vote_0; + if (skip_correction == 0) { + #pragma unroll + for (int local_slice = 0; local_slice < 2; local_slice++) { + #pragma unroll + for (int acc_stage = 0; acc_stage < 2; acc_stage++) { + int o_base = taddr + 128 + (unsigned int)(acc_stage * 128) + (unsigned int)(local_slice * 64) + (unsigned int)(output_row_base << 16); + #pragma unroll + for (int c = 0; c < 64; c += 16) { + float _tmem_load_2[16]; + tmem_ld_x16(&_tmem_load_2[0], o_base + c); + const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + mul_f32x2_inplace(&reinterpret_cast(_tmem_load_2)[_ls], _scale2_0); + tmem_st_x16_f32(o_base + c, _tmem_load_2); + } + } + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + } + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + } + mbarrier_arrive(stats_empty_addr + (score_phase_1) * 8); + } + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + const int exchange_idx_1 = n_half_1 * 64 + my_row_1; + smem_epilogue_exchange[exchange_idx_1] = final_sum; + asm volatile("barrier.sync 4, 128;" ::: "memory"); + final_sum = final_sum + smem_epilogue_exchange[exchange_idx_1 ^ 64]; + asm volatile("barrier.sync 4, 128;" ::: "memory"); + float _rcp_0 = approx_rcp(final_sum); + float inv_sum = ((final_sum > 0.0f) ? _rcp_0 : 0.0f); + int head_idx_1 = cta_rank * 64 + my_row_1; + int output_offset = ((1) ? ((query_idx_1 * num_heads + head_idx_1) * ((1) ? 5 : 1) + split_idx_1) * 512 : (query_idx_1 * num_heads + head_idx_1) * 512); + if (n_half_1 == 0 && head_idx_1 < num_heads) { + int stat_offset = (query_idx_1 * num_heads + head_idx_1) * ((1) ? 5 : 1) + split_idx_1; + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(final_sum)); + partial_lse[stat_offset] = ((final_sum > 0.0f) ? final_max * softmax_scale_log2_1 + _log2_0 : -CAKE_INF); + } + float epi_scale = inv_sum * output_scale; + int epi_box_base = smem_q_addr + (unsigned int)(warp % 4 * 16384); + int epi_row0 = query_idx_1 * num_heads + cta_rank * 64 + head_row_base_1; + #pragma unroll + for (int acc_stage_1 = 0; acc_stage_1 < 2; acc_stage_1++) { + #pragma unroll + for (int local_slice_1 = 0; local_slice_1 < 2; local_slice_1++) { + int logical_slice = acc_stage_1 * 2 * 2 + n_half_1 * 2 + local_slice_1; + int epi_box = epi_box_base + (acc_stage_1 * 2 + local_slice_1) * 4096; + int o_base_1 = taddr + 128 + (unsigned int)(acc_stage_1 * 128) + (unsigned int)(local_slice_1 * 64) + (unsigned int)(output_row_base << 16); + #pragma unroll + for (int c_1 = 0; c_1 < 64; c_1 += 32) { + float _tmem_load_3[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]) + : "r"(o_base_1 + c_1)); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + float epi_scaled[32]; + #pragma unroll + for (int e = 0; e < 32; e++) { + epi_scaled[e] = _tmem_load_3[e] * epi_scale; + } + unsigned int epi_packed[16]; + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(epi_scaled[_lp*2 + 0], epi_scaled[_lp*2+1 + 0])); + epi_packed[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((epi_box + (lane * 128 + c_1 * 2 ^ (lane * 128 + c_1 * 2 >> 7 & 7) << 4))), "r"(epi_packed[0]), "r"(epi_packed[1]), "r"(epi_packed[2]), "r"(epi_packed[3]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((epi_box + (lane * 128 + (c_1 * 2 + 16) ^ (lane * 128 + (c_1 * 2 + 16) >> 7 & 7) << 4))), "r"(epi_packed[4]), "r"(epi_packed[5]), "r"(epi_packed[6]), "r"(epi_packed[7]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((epi_box + (lane * 128 + (c_1 * 2 + 32) ^ (lane * 128 + (c_1 * 2 + 32) >> 7 & 7) << 4))), "r"(epi_packed[8]), "r"(epi_packed[9]), "r"(epi_packed[10]), "r"(epi_packed[11]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((epi_box + (lane * 128 + (c_1 * 2 + 48) ^ (lane * 128 + (c_1 * 2 + 48) >> 7 & 7) << 4))), "r"(epi_packed[12]), "r"(epi_packed[13]), "r"(epi_packed[14]), "r"(epi_packed[15]) : "memory"); + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + __syncwarp(); + if (elect_sync()) { + tma_store_4d((&O), 0, logical_slice, split_idx_1, epi_row0, epi_box); + asm volatile("cp.async.bulk.commit_group;"); + } + } + } + if (elect_sync()) { + asm volatile("cp.async.bulk.wait_group.read 0;"); + mbarrier_arrive(epi_q_free_addr); + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + tile_cursor_1 = tile_cursor_1 + num_kv_tiles_1; + } + if (elect_sync()) { + asm volatile("cp.async.bulk.wait_group 0;"); + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + } + } + // ---- Role: mma_warp ---- + if (warp == 8) { + { // mma_warp_main + const int mma_dummy = 0; + unsigned int k_stage = 0; + unsigned int v_stage = 0; + int tile_cursor_2 = 0; + unsigned int _phase_o_first_slice_seeded_0 = 0; + mbarrier_wait(o_first_slice_seeded_addr, _phase_o_first_slice_seeded_0); + _phase_o_first_slice_seeded_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + unsigned int _phase_q_full_0 = 0; + unsigned int _phase_k_full = 0; + unsigned int _phase_o_empty_0 = 1; + unsigned int _phase_v_full = 0; + #pragma unroll 1 + for (unsigned int work_idx_2 = cluster_id; work_idx_2 < total_work_items; work_idx_2 += num_clusters) { + int owner_capacity_2 = ((sparse_topk + 128 - 1) / 128 + ((1) ? 5 : 1) - 1) / ((1) ? 5 : 1); + int first_owner_tile_2 = work_idx_2 % (unsigned int)(((1) ? 5 : 1)) * (unsigned int)owner_capacity_2; + int _min_1 = ((owner_capacity_2) < ((sparse_topk + 128 - 1) / 128 - first_owner_tile_2) ? (owner_capacity_2) : ((sparse_topk + 128 - 1) / 128 - first_owner_tile_2)); + int num_kv_tiles_2 = _min_1; + if (cta_rank == 0) { + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + int first_pv = 1; + #pragma unroll 1 + for (int tile_2 = 0; tile_2 < num_kv_tiles_2; tile_2++) { + int pipeline_tile_2 = tile_cursor_2 + tile_2; + int score_phase_2 = pipeline_tile_2 & 1; + int s_empty_phase = pipeline_tile_2 >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (score_phase_2) * 8, s_empty_phase); + int score_col = ((score_phase_2 != 0) ? 64 : 0); + #pragma unroll + for (int qk_stage = 0; qk_stage < 4; qk_stage++) { + mbarrier_wait(k_full_addr + (k_stage) * 8, _phase_k_full); + int _mma_a_lo_0 = (((smem_q_addr) >> 4) & 0x3FFF) + (qk_stage * 2) * 512; + int _mma_b_lo_0 = (((smem_k_addr) >> 4) & 0x3FFF) + (k_stage) * 1024; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (score_col))), "r"(((qk_stage == 0) ? 0 : 1))); + int _mma_a_lo_1 = (((smem_q_addr) >> 4) & 0x3FFF) + (qk_stage * 2 + 1) * 512; + int _mma_b_lo_1 = (((smem_k_addr + 8192) >> 4) & 0x3FFF) + (k_stage) * 1024; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + elect_commit_cg2_multicast(k_empty_addr + (k_stage) * 8, (uint16_t)(3)); + k_stage += 1; + if (k_stage == 8) { k_stage = 0; _phase_k_full ^= 1; } + } + elect_commit_cg2_multicast(s_full_addr + (score_phase_2) * 8, (uint16_t)(3)); + if (tile_2 > 0) { + int p_stage_1 = (pipeline_tile_2 - 1) % 2; + int p_full_phase = (pipeline_tile_2 - 1) / 2 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (p_stage_1) * 8, p_full_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps = 0; ps < 2; ps++) { + int smem_p_stage_1 = p_stage_1 * 2 + ps; + mbarrier_wait(v_full_addr + (v_stage) * 8, _phase_v_full); + #pragma unroll + for (int acc_stage_2 = 0; acc_stage_2 < 2; acc_stage_2++) { + int output_col = 128 + acc_stage_2 * 128; + unsigned int v_slot = (k_stage + (unsigned int)(acc_stage_2 * 2 + ps)) % 8; + int _mma_a_lo_2 = (((smem_p_addr) >> 4) & 0x3FFF) + (smem_p_stage_1) * 512; + int _mma_b_lo_2 = ((((smem_v_addr) >> 4) & 0x3FFF) | 0x2000000) + (v_slot) * 1024; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (output_col))), "r"(((((ps == 0) ? first_pv : 0)) ? 0 : 1))); + } + elect_commit_cg2_multicast(v_empty_addr + (v_stage) * 8, (uint16_t)(3)); + #pragma unroll + for (int acc_stage_3 = 0; acc_stage_3 < 2; acc_stage_3++) { + unsigned int free_slot = (k_stage + (unsigned int)(acc_stage_3 * 2 + ps)) % 8; + elect_commit_cg2_multicast(slot_empty_addr + (free_slot) * 8, (uint16_t)(3)); + } + v_stage += 1; + if (v_stage == 2) { v_stage = 0; _phase_v_full ^= 1; } + } + elect_commit_cg2_multicast(p_empty_addr + (p_stage_1) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + first_pv = 0; + } + } + int p_stage_2 = (tile_cursor_2 + num_kv_tiles_2 - 1) % 2; + int p_full_phase_1 = (tile_cursor_2 + num_kv_tiles_2 - 1) / 2 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (p_stage_2) * 8, p_full_phase_1); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_1 = 0; ps_1 < 2; ps_1++) { + int smem_p_stage_2 = p_stage_2 * 2 + ps_1; + mbarrier_wait(v_full_addr + (v_stage) * 8, _phase_v_full); + #pragma unroll + for (int acc_stage_4 = 0; acc_stage_4 < 2; acc_stage_4++) { + int output_col_1 = 128 + acc_stage_4 * 128; + unsigned int v_slot_1 = (k_stage + (unsigned int)(4 + acc_stage_4 * 2 + ps_1)) % 8; + int _mma_a_lo_3 = (((smem_p_addr) >> 4) & 0x3FFF) + (smem_p_stage_2) * 512; + int _mma_b_lo_3 = ((((smem_v_addr) >> 4) & 0x3FFF) | 0x2000000) + (v_slot_1) * 1024; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"((tmem_tmem_scratch + (output_col_1))), "r"(((((ps_1 == 0) ? first_pv : 0)) ? 0 : 1))); + } + elect_commit_cg2_multicast(v_empty_addr + (v_stage) * 8, (uint16_t)(3)); + #pragma unroll + for (int acc_stage_5 = 0; acc_stage_5 < 2; acc_stage_5++) { + unsigned int free_slot_1 = (k_stage + (unsigned int)(4 + acc_stage_5 * 2 + ps_1)) % 8; + elect_commit_cg2_multicast(slot_empty_addr + (free_slot_1) * 8, (uint16_t)(3)); + } + v_stage += 1; + if (v_stage == 2) { v_stage = 0; _phase_v_full ^= 1; } + } + elect_commit_cg2_multicast(p_empty_addr + (p_stage_2) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + } + tile_cursor_2 = tile_cursor_2 + num_kv_tiles_2; + } + if (cta_rank == 0) { + #pragma unroll + for (int tail_offset = 0; tail_offset < 2; tail_offset++) { + int tail_tile = tile_cursor_2 + tail_offset; + int tail_stage = tail_tile & 1; + int tail_phase = tail_tile >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (tail_stage) * 8, tail_phase); + } + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + } + asm volatile("tcgen05.fence::before_thread_sync;"); + int peer_rank_1 = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_dealloc_peer_addr), "r"(peer_rank_1) : "memory"); + mbarrier_wait(tmem_dealloc_peer_addr, 0); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + } + // ---- Role: load_wg ---- + if (warp >= 9 && warp <= 12) { + { // load_wg_main + const int load_dummy = 0; + const int load_warp_rank = warp - 9; + unsigned int k_stage_1 = 0; + unsigned int v_stage_1 = 0; + int load_tile_cursor = 0; + unsigned int _phase_q_empty_0 = 1; + unsigned int _phase_epi_q_free_0 = 1; + unsigned int _phase_slot_empty = 1; + unsigned int _phase_v_empty = 1; + #pragma unroll 1 + for (unsigned int work_idx_3 = cluster_id; work_idx_3 < total_work_items; work_idx_3 += num_clusters) { + int split_idx_2 = work_idx_3 % (unsigned int)(((1) ? 5 : 1)); + int query_idx_2 = work_idx_3 / (unsigned int)(((1) ? 5 : 1)); + int owner_capacity_3 = ((sparse_topk + 128 - 1) / 128 + ((1) ? 5 : 1) - 1) / ((1) ? 5 : 1); + int first_owner_tile_3 = work_idx_3 % (unsigned int)(((1) ? 5 : 1)) * (unsigned int)owner_capacity_3; + int _min_0 = ((owner_capacity_3) < ((sparse_topk + 128 - 1) / 128 - first_owner_tile_3) ? (owner_capacity_3) : ((sparse_topk + 128 - 1) / 128 - first_owner_tile_3)); + int num_kv_tiles_3 = _min_0; + int first_kv_tile_1 = first_owner_tile_3; + int group = load_warp_rank * 4 + lane; + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + mbarrier_wait(epi_q_free_addr, _phase_epi_q_free_0); + _phase_epi_q_free_0 ^= 1; + if (load_warp_rank == 0) { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(131072)) : "memory"); + } + } + if (elect_sync()) { + #pragma unroll + for (int q_stage = 0; q_stage < 8; q_stage++) { + tma_4d_gmem2smem_cta2(smem_q_addr + (unsigned int)(q_stage * 8192), (&tmap_q), 0, cta_rank * 64, q_stage, query_idx_2, ((q_full_addr) & 0xFEFFFFFF)); + } + } + } + asm volatile("barrier.sync 2, 160;" ::: "memory"); + #pragma unroll + for (int qk_stage_1 = 0; qk_stage_1 < 4; qk_stage_1++) { + if (first_kv_tile_1 == 0) { + mbarrier_wait(slot_empty_addr + (k_stage_1) * 8, _phase_slot_empty); + if (load_warp_rank == 0) { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + } + } + } + if (lane < 4) { + int index_offset = cta_rank * 64 + group * 4; + int raw_rows[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows[0])), "=r"(*reinterpret_cast(&raw_rows[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows[(0) + 3])) + : "r"(smem_indices_addr + (unsigned int)(index_offset * 4))); + tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group * 512), (&tmap_swa_k), qk_stage_1 * 128, ((raw_rows[0] >= 0) ? raw_rows[0] : 0), ((raw_rows[1] >= 0) ? raw_rows[1] : 0), ((raw_rows[2] >= 0) ? raw_rows[2] : 0), ((raw_rows[3] >= 0) ? raw_rows[3] : 0), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group * 512) + 8192, (&tmap_swa_k), qk_stage_1 * 128 + 64, ((raw_rows[0] >= 0) ? raw_rows[0] : 0), ((raw_rows[1] >= 0) ? raw_rows[1] : 0), ((raw_rows[2] >= 0) ? raw_rows[2] : 0), ((raw_rows[3] >= 0) ? raw_rows[3] : 0), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + } + k_stage_1 += 1; + if (k_stage_1 == 8) { k_stage_1 = 0; _phase_slot_empty ^= 1; } + } else { + mbarrier_wait(slot_empty_addr + (k_stage_1) * 8, _phase_slot_empty); + if (load_warp_rank == 0) { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + } + } + } + if (lane < 4) { + int index_offset_1 = cta_rank * 64 + group * 4; + int raw_rows_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_1[0])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 3])) + : "r"(smem_indices_addr + (unsigned int)(index_offset_1 * 4))); + tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group * 512), (&tmap_compressed_k), qk_stage_1 * 128, ((raw_rows_1[0] >= 0) ? raw_rows_1[0] : 0), ((raw_rows_1[1] >= 0) ? raw_rows_1[1] : 0), ((raw_rows_1[2] >= 0) ? raw_rows_1[2] : 0), ((raw_rows_1[3] >= 0) ? raw_rows_1[3] : 0), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group * 512) + 8192, (&tmap_compressed_k), qk_stage_1 * 128 + 64, ((raw_rows_1[0] >= 0) ? raw_rows_1[0] : 0), ((raw_rows_1[1] >= 0) ? raw_rows_1[1] : 0), ((raw_rows_1[2] >= 0) ? raw_rows_1[2] : 0), ((raw_rows_1[3] >= 0) ? raw_rows_1[3] : 0), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + } + k_stage_1 += 1; + if (k_stage_1 == 8) { k_stage_1 = 0; _phase_slot_empty ^= 1; } + } + } + #pragma unroll 1 + for (int tile_3 = 1; tile_3 < num_kv_tiles_3; tile_3++) { + #pragma unroll + for (int qk_stage_2 = 0; qk_stage_2 < 4; qk_stage_2++) { + mbarrier_wait(slot_empty_addr + (k_stage_1) * 8, _phase_slot_empty); + if (load_warp_rank == 0) { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + } + } + } + if (lane < 4) { + int index_offset_2 = tile_3 * 128 + cta_rank * 64 + group * 4; + int raw_rows_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_2[0])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 3])) + : "r"(smem_indices_addr + (unsigned int)(index_offset_2 * 4))); + tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group * 512), (&tmap_compressed_k), qk_stage_2 * 128, ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : 0), ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : 0), ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : 0), ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : 0), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group * 512) + 8192, (&tmap_compressed_k), qk_stage_2 * 128 + 64, ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : 0), ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : 0), ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : 0), ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : 0), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + } + k_stage_1 += 1; + if (k_stage_1 == 8) { k_stage_1 = 0; _phase_slot_empty ^= 1; } + } + int prev_tile = tile_3 - 1; + #pragma unroll + for (int ps_2 = 0; ps_2 < 2; ps_2++) { + if (first_kv_tile_1 + prev_tile == 0) { + mbarrier_wait(v_empty_addr + (v_stage_1) * 8, _phase_v_empty); + if (load_warp_rank == 0) { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + } + } + } + if (cta_rank != ps_2) { + int fill_base = (load_tile_cursor + prev_tile) * 4; + #pragma unroll + for (int acc_stage_6 = 0; acc_stage_6 < 2; acc_stage_6++) { + int fill = fill_base + acc_stage_6 * 2 + ps_2; + mbarrier_wait(k_empty_addr + ((k_stage_1 + (unsigned int)(acc_stage_6 * 2 + ps_2)) % 8) * 8, fill / 8 & 1); + } + if (lane < 4) { + int index_offset_3 = prev_tile * 128 + ps_2 * 64 + group * 4; + int raw_rows_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_3[0])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 3])) + : "r"(smem_indices_addr + (unsigned int)(index_offset_3 * 4))); + #pragma unroll + for (int acc_stage_7 = 0; acc_stage_7 < 2; acc_stage_7++) { + unsigned int copy_slot = (k_stage_1 + (unsigned int)(acc_stage_7 * 2 + ps_2)) % 8; + tma_gather4_gmem2smem_mc_cta2(smem_k_addr + copy_slot * 16384 + (unsigned int)(group * 512), (&tmap_swa_v), (acc_stage_7 * 2 + (1 - ps_2)) * 128, ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : 0), ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : 0), ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : 0), ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : 0), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + tma_gather4_gmem2smem_mc_cta2(smem_k_addr + copy_slot * 16384 + (unsigned int)(group * 512) + 8192, (&tmap_swa_v), (acc_stage_7 * 2 + (1 - ps_2)) * 128 + 64, ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : 0), ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : 0), ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : 0), ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : 0), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + } + } + } + v_stage_1 += 1; + if (v_stage_1 == 2) { v_stage_1 = 0; _phase_v_empty ^= 1; } + } else { + mbarrier_wait(v_empty_addr + (v_stage_1) * 8, _phase_v_empty); + if (load_warp_rank == 0) { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + } + } + } + if (cta_rank != ps_2) { + int fill_base_1 = (load_tile_cursor + prev_tile) * 4; + #pragma unroll + for (int acc_stage_8 = 0; acc_stage_8 < 2; acc_stage_8++) { + int fill_1 = fill_base_1 + acc_stage_8 * 2 + ps_2; + mbarrier_wait(k_empty_addr + ((k_stage_1 + (unsigned int)(acc_stage_8 * 2 + ps_2)) % 8) * 8, fill_1 / 8 & 1); + } + if (lane < 4) { + int index_offset_4 = prev_tile * 128 + ps_2 * 64 + group * 4; + int raw_rows_4[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_4[0])), "=r"(*reinterpret_cast(&raw_rows_4[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_4[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_4[(0) + 3])) + : "r"(smem_indices_addr + (unsigned int)(index_offset_4 * 4))); + #pragma unroll + for (int acc_stage_9 = 0; acc_stage_9 < 2; acc_stage_9++) { + unsigned int copy_slot_1 = (k_stage_1 + (unsigned int)(acc_stage_9 * 2 + ps_2)) % 8; + tma_gather4_gmem2smem_mc_cta2(smem_k_addr + copy_slot_1 * 16384 + (unsigned int)(group * 512), (&tmap_compressed_v), (acc_stage_9 * 2 + (1 - ps_2)) * 128, ((raw_rows_4[0] >= 0) ? raw_rows_4[0] : 0), ((raw_rows_4[1] >= 0) ? raw_rows_4[1] : 0), ((raw_rows_4[2] >= 0) ? raw_rows_4[2] : 0), ((raw_rows_4[3] >= 0) ? raw_rows_4[3] : 0), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + tma_gather4_gmem2smem_mc_cta2(smem_k_addr + copy_slot_1 * 16384 + (unsigned int)(group * 512) + 8192, (&tmap_compressed_v), (acc_stage_9 * 2 + (1 - ps_2)) * 128 + 64, ((raw_rows_4[0] >= 0) ? raw_rows_4[0] : 0), ((raw_rows_4[1] >= 0) ? raw_rows_4[1] : 0), ((raw_rows_4[2] >= 0) ? raw_rows_4[2] : 0), ((raw_rows_4[3] >= 0) ? raw_rows_4[3] : 0), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + } + } + } + v_stage_1 += 1; + if (v_stage_1 == 2) { v_stage_1 = 0; _phase_v_empty ^= 1; } + } + } + } + int last_tile = num_kv_tiles_3 - 1; + #pragma unroll + for (int ps_3 = 0; ps_3 < 2; ps_3++) { + if (first_kv_tile_1 + last_tile == 0) { + mbarrier_wait(v_empty_addr + (v_stage_1) * 8, _phase_v_empty); + if (load_warp_rank == 0) { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + } + } + } + if (cta_rank != ps_3) { + int fill_base_2 = (load_tile_cursor + last_tile) * 4; + #pragma unroll + for (int acc_stage_10 = 0; acc_stage_10 < 2; acc_stage_10++) { + int fill_2 = fill_base_2 + acc_stage_10 * 2 + ps_3; + mbarrier_wait(k_empty_addr + ((k_stage_1 + (unsigned int)(4 + acc_stage_10 * 2 + ps_3)) % 8) * 8, fill_2 / 8 & 1); + } + if (lane < 4) { + int index_offset_5 = last_tile * 128 + ps_3 * 64 + group * 4; + int raw_rows_5[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_5[0])), "=r"(*reinterpret_cast(&raw_rows_5[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_5[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_5[(0) + 3])) + : "r"(smem_indices_addr + (unsigned int)(index_offset_5 * 4))); + #pragma unroll + for (int acc_stage_11 = 0; acc_stage_11 < 2; acc_stage_11++) { + unsigned int copy_slot_2 = (k_stage_1 + (unsigned int)(4 + acc_stage_11 * 2 + ps_3)) % 8; + tma_gather4_gmem2smem_mc_cta2(smem_k_addr + copy_slot_2 * 16384 + (unsigned int)(group * 512), (&tmap_swa_v), (acc_stage_11 * 2 + (1 - ps_3)) * 128, ((raw_rows_5[0] >= 0) ? raw_rows_5[0] : 0), ((raw_rows_5[1] >= 0) ? raw_rows_5[1] : 0), ((raw_rows_5[2] >= 0) ? raw_rows_5[2] : 0), ((raw_rows_5[3] >= 0) ? raw_rows_5[3] : 0), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + tma_gather4_gmem2smem_mc_cta2(smem_k_addr + copy_slot_2 * 16384 + (unsigned int)(group * 512) + 8192, (&tmap_swa_v), (acc_stage_11 * 2 + (1 - ps_3)) * 128 + 64, ((raw_rows_5[0] >= 0) ? raw_rows_5[0] : 0), ((raw_rows_5[1] >= 0) ? raw_rows_5[1] : 0), ((raw_rows_5[2] >= 0) ? raw_rows_5[2] : 0), ((raw_rows_5[3] >= 0) ? raw_rows_5[3] : 0), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + } + } + } + v_stage_1 += 1; + if (v_stage_1 == 2) { v_stage_1 = 0; _phase_v_empty ^= 1; } + } else { + mbarrier_wait(v_empty_addr + (v_stage_1) * 8, _phase_v_empty); + if (load_warp_rank == 0) { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + } + } + } + if (cta_rank != ps_3) { + int fill_base_3 = (load_tile_cursor + last_tile) * 4; + #pragma unroll + for (int acc_stage_12 = 0; acc_stage_12 < 2; acc_stage_12++) { + int fill_3 = fill_base_3 + acc_stage_12 * 2 + ps_3; + mbarrier_wait(k_empty_addr + ((k_stage_1 + (unsigned int)(4 + acc_stage_12 * 2 + ps_3)) % 8) * 8, fill_3 / 8 & 1); + } + if (lane < 4) { + int index_offset_6 = last_tile * 128 + ps_3 * 64 + group * 4; + int raw_rows_6[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_6[0])), "=r"(*reinterpret_cast(&raw_rows_6[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_6[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_6[(0) + 3])) + : "r"(smem_indices_addr + (unsigned int)(index_offset_6 * 4))); + #pragma unroll + for (int acc_stage_13 = 0; acc_stage_13 < 2; acc_stage_13++) { + unsigned int copy_slot_3 = (k_stage_1 + (unsigned int)(4 + acc_stage_13 * 2 + ps_3)) % 8; + tma_gather4_gmem2smem_mc_cta2(smem_k_addr + copy_slot_3 * 16384 + (unsigned int)(group * 512), (&tmap_compressed_v), (acc_stage_13 * 2 + (1 - ps_3)) * 128, ((raw_rows_6[0] >= 0) ? raw_rows_6[0] : 0), ((raw_rows_6[1] >= 0) ? raw_rows_6[1] : 0), ((raw_rows_6[2] >= 0) ? raw_rows_6[2] : 0), ((raw_rows_6[3] >= 0) ? raw_rows_6[3] : 0), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + tma_gather4_gmem2smem_mc_cta2(smem_k_addr + copy_slot_3 * 16384 + (unsigned int)(group * 512) + 8192, (&tmap_compressed_v), (acc_stage_13 * 2 + (1 - ps_3)) * 128 + 64, ((raw_rows_6[0] >= 0) ? raw_rows_6[0] : 0), ((raw_rows_6[1] >= 0) ? raw_rows_6[1] : 0), ((raw_rows_6[2] >= 0) ? raw_rows_6[2] : 0), ((raw_rows_6[3] >= 0) ? raw_rows_6[3] : 0), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + } + } + } + v_stage_1 += 1; + if (v_stage_1 == 2) { v_stage_1 = 0; _phase_v_empty ^= 1; } + } + } + asm volatile("barrier.sync 2, 160;" ::: "memory"); + load_tile_cursor = load_tile_cursor + num_kv_tiles_3; + } + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + } + } + // ---- Role: index_warp ---- + if (warp == 13) { + { // index_warp_main + const int index_dummy = 0; + int valid_iter_1 = 0; + #pragma unroll 1 + for (unsigned int work_idx_4 = cluster_id; work_idx_4 < total_work_items; work_idx_4 += num_clusters) { + int split_idx_3 = work_idx_4 % (unsigned int)(((1) ? 5 : 1)); + int query_idx_3 = work_idx_4 / (unsigned int)(((1) ? 5 : 1)); + int owner_col_base = split_idx_3 * ((9 + ((1) ? 5 : 1) - 1) / ((1) ? 5 : 1)) * 128; + int* swa_row = swa_indices + (query_idx_3 * swa_index_stride); + int* compressed_row = compressed_indices + (query_idx_3 * compressed_index_stride); + int valid_slot_1 = valid_iter_1 & 1; + int valid_empty_phase = valid_iter_1 >> 1 & 1 ^ 1; + mbarrier_wait(index_valid_empty_addr + (valid_slot_1) * 8, valid_empty_phase); + #pragma unroll + for (int index_pass = 0; index_pass < (9 + ((1) ? 5 : 1) - 1) / ((1) ? 5 : 1); index_pass++) { + int index_offset_7 = index_pass * 128 + lane * 4; + int combined_col = owner_col_base + index_offset_7; + int* index_src = ((combined_col < 128) ? swa_row + combined_col : compressed_row + (combined_col - 128)); + int values[4]; + #pragma unroll + for (int i = 0; i < 4; i++) { + values[i] = -1; + } + if (combined_col + 4 <= sparse_topk) { + int _vec_load_0[4]; + { + const int4* _ivptr_0 = reinterpret_cast(index_src + 0); + int4 _ivld_0; + _ivld_0 = *_ivptr_0; + _vec_load_0[0 + 0] = _ivld_0.x; + _vec_load_0[0 + 1] = _ivld_0.y; + _vec_load_0[0 + 2] = _ivld_0.z; + _vec_load_0[0 + 3] = _ivld_0.w; + } + #pragma unroll + for (int i_1 = 0; i_1 < 4; i_1++) { + values[i_1] = _vec_load_0[i_1]; + } + } + if (combined_col < sparse_topk && combined_col + 4 > sparse_topk) { + #pragma unroll + for (int i_2 = 0; i_2 < 4; i_2++) { + if (combined_col + i_2 < sparse_topk) { + values[i_2] = index_src[i_2]; + } + } + } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(smem_indices_addr + (unsigned int)(index_offset_7 * 4)), "r"(values[0]), "r"(values[1]), "r"(values[2]), "r"(values[3]) : "memory"); + unsigned int valid_word = (((values[0] >= 0) ? 1 : 0) | ((values[1] >= 0) ? 1 : 0) << 1 | ((values[2] >= 0) ? 1 : 0) << 2 | ((values[3] >= 0) ? 1 : 0) << 3) << (lane & 7) * 4; + unsigned int _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, valid_word, 1); + valid_word = valid_word | _shfl_xor_0; + unsigned int _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, valid_word, 2); + valid_word = valid_word | _shfl_xor_1; + unsigned int _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, valid_word, 4); + valid_word = valid_word | _shfl_xor_2; + if ((lane & 7) == 0) { + smem_index_valid[valid_slot_1 * 8 + index_pass * 4 + (lane >> 3)] = valid_word; + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(index_valid_full_addr + (valid_slot_1) * 8); + asm volatile("barrier.sync 2, 160;" ::: "memory"); + asm volatile("barrier.sync 2, 160;" ::: "memory"); + valid_iter_1 = valid_iter_1 + 1; + } + } + } + // ---- Role: empty_warps ---- + if (warp >= 14 && warp <= 15) { + { // empty_warps_main + const int empty_dummy = 0; + #pragma unroll 1 + for (unsigned int work_idx_5 = cluster_id; work_idx_5 < total_work_items; work_idx_5 += num_clusters) { + asm volatile("tcgen05.fence::after_thread_sync;"); + } + } + } + + // Cleanup +} + +} // extern "C" diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_868bb2291a1464f14e05_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_868bb2291a1464f14e05_binding.cu new file mode 100644 index 00000000000..fda3c5b89ff --- /dev/null +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_868bb2291a1464f14e05_binding.cu @@ -0,0 +1,525 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +// Generated by the Cake source exporter — do not edit. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_868bb2291a1464f14e05'. +#include +#include +#include +#include + +#include "tvm_ffi_utils.h" + +#include +#include +#include +#include +#include +#include +#include + +struct CakeTensorMap; +extern "C" __global__ void kernel_cake_dsv4_868bb2291a1464f14e05(CakeTensorMap const* tmap_q, CakeTensorMap const* tmap_swa_kv, CakeTensorMap const* tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items); + + +namespace cake_host_shim_262e38a78b1a7fce { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t CakeDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +// Explicit immutable caller storage: initialize once, outside graph capture, +// exactly as the production TmaDeviceSlot path. The owner must retain the +// allocation and must never expose it as mutable scratch or rewrite its bytes. +template +static inline void PrepareImmutableCallerTmaWorkspace( + const tvm::ffi::Tensor& retained_workspace, + const CUtensorMap (&maps)[N], cudaStream_t stream) { + void* workspace = retained_workspace.data_ptr(); + CUcontext context = nullptr; + CUresult result = cuCtxGetCurrent(&context); + TVM_FFI_CHECK(result == CUDA_SUCCESS && context != nullptr, RuntimeError) + << "immutable TMA workspace requires an active CUDA context"; + CUcontext allocation_context = nullptr; + result = cuPointerGetAttribute( + &allocation_context, CU_POINTER_ATTRIBUTE_CONTEXT, + reinterpret_cast(workspace)); + TVM_FFI_CHECK(result == CUDA_SUCCESS && allocation_context == context, RuntimeError) + << "immutable TMA workspace must belong to the active CUDA context"; + unsigned long long buffer_id = 0; + result = cuPointerGetAttribute( + &buffer_id, CU_POINTER_ATTRIBUTE_BUFFER_ID, + reinterpret_cast(workspace)); + TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) + << "querying immutable TMA allocation identity failed"; + std::string key = std::to_string(reinterpret_cast(context)); + key += ":" + std::to_string(reinterpret_cast(workspace)); + key += ":" + std::to_string(buffer_id); + std::string descriptor_bytes(reinterpret_cast(maps), sizeof(maps)); + struct InitializedWorkspace { + tvm::ffi::Tensor owner; + std::string descriptor_bytes; + }; + static std::mutex mu; + // Keep the actual Tensor allocation alive: a caching allocator may reuse a + // freed suballocation without changing its enclosing CUDA buffer identity. + static auto* initialized = new std::unordered_map(); + std::lock_guard lock(mu); + auto it = initialized->find(key); + if (it != initialized->end()) { + TVM_FFI_CHECK(it->second.descriptor_bytes == descriptor_bytes, ValueError) + << "immutable TMA descriptors changed; provide new private descriptor storage"; + return; + } + CUstreamCaptureStatus capture_status = CU_STREAM_CAPTURE_STATUS_NONE; + result = cuStreamIsCapturing(reinterpret_cast(stream), &capture_status); + TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) + << "querying immutable TMA initialization capture state failed"; + TVM_FFI_CHECK(capture_status == CU_STREAM_CAPTURE_STATUS_NONE, RuntimeError) + << "immutable TMA storage must be initialized before CUDA Graph capture"; + // Synchronous host-to-device initialization completes before publishing the + // slot to any thread/stream, matching TmaDeviceSlot. The consumer retains its + // generated tensor-map acquire; the mutable upload's release fence is unchanged. + result = cuMemcpyHtoD(reinterpret_cast(workspace), maps, sizeof(maps)); + TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) + << "initializing immutable caller TMA descriptors failed"; + initialized->emplace( + std::move(key), InitializedWorkspace{retained_workspace, std::move(descriptor_bytes)}); +} + +static inline void* CallerTmaWorkspaceSlot( + const tvm::ffi::TensorView& workspace, size_t slot) { + return static_cast(workspace.data_ptr()) + slot * sizeof(CUtensorMap); +} + +// 4D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + int64_t d3 = t.size(t.ndim() - 3); + TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + TVM_FFI_CHECK(d1 % 128 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 128; + uint64_t global_dim[4] = {(uint64_t)(128), (uint64_t)(d2), (uint64_t)((d1 / 128)), (uint64_t)(d3)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (128, 128, 1, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 128; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + (uint64_t)((carrier_stride_1 * 8) / 8), + (uint64_t)((carrier_stride_2 * 8) / 8), + }; + uint32_t box_dim[4] = {128u, 128u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_swa_kv' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_swa_kv' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; + TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (128, 1) exceeds resolved global dims for 'tmap_swa_kv'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + }; + uint32_t box_dim[2] = {128u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_compressed_kv' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_compressed_kv(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_compressed_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_compressed_kv' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_compressed_kv' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_compressed_kv"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_compressed_kv' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved a non-positive global dim"; + TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (128, 1) exceeds resolved global dims for 'tmap_compressed_kv'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + }; + uint32_t box_dim[2] = {128u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_kv') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_compressed_kv, TensorView arg_O, TensorView arg_partial_lse, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_num_query_tokens, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, tvm::ffi::Tensor arg_tma_descriptor_workspace, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 1, 8, 1); + CheckContiguous(arg_tmap_q, "tmap_q"); + CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); + CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 1, 8, 1); + CheckCudaTensor(arg_tmap_compressed_kv, "tmap_compressed_kv"); + CheckDtype(arg_tmap_compressed_kv, "tmap_compressed_kv", 1, 8, 1); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_partial_lse, "partial_lse"); + CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); + CheckContiguous(arg_partial_lse, "partial_lse"); + CheckCudaTensor(arg_swa_indices, "swa_indices"); + CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); + CheckContiguous(arg_swa_indices, "swa_indices"); + CheckCudaTensor(arg_compressed_indices, "compressed_indices"); + CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); + CheckContiguous(arg_compressed_indices, "compressed_indices"); + CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); + CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); + CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); + CheckContiguous(arg_bmm1_scale, "bmm1_scale"); + CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); + CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); + CheckContiguous(arg_bmm2_scale, "bmm2_scale"); + TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) + << "scalar 'num_heads' value " << arg_num_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) + << "scalar 'swa_index_stride' value " << arg_swa_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) + << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) + << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_query_tokens >= -2147483648LL && arg_num_query_tokens <= 2147483647LL, ValueError) + << "scalar 'num_query_tokens' value " << arg_num_query_tokens + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) + << "scalar 'sparse_topk' value " << arg_sparse_topk + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) + << "scalar 'has_sinks' value " << arg_has_sinks + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); + CheckSameCudaDevice(arg_tmap_compressed_kv, arg_tmap_q, "tmap_compressed_kv", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_partial_lse, arg_tmap_q, "partial_lse", "tmap_q"); + CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); + CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); + CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); + CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); + CheckCudaTensor(arg_tma_descriptor_workspace, "tma_descriptor_workspace"); + CheckSameCudaDevice(arg_tma_descriptor_workspace, arg_tmap_q, "tma_descriptor_workspace", "tmap_q"); + CheckContiguous(arg_tma_descriptor_workspace, "tma_descriptor_workspace"); + CheckDtype(arg_tma_descriptor_workspace, "tma_descriptor_workspace", 1, 8, 1); + TVM_FFI_CHECK(arg_tma_descriptor_workspace.numel() >= 384, ValueError) + << "tma_descriptor_workspace requires at least 384 bytes"; + TVM_FFI_CHECK( + reinterpret_cast(arg_tma_descriptor_workspace.data_ptr()) % 128 == 0, + ValueError) + << "tma_descriptor_workspace must be 128-byte aligned"; + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + + cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); + CUtensorMap h_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CakeTensorMap const* p_tmap_q = static_cast(CallerTmaWorkspaceSlot(arg_tma_descriptor_workspace, 0u)); + CUtensorMap h_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); + CakeTensorMap const* p_tmap_swa_kv = static_cast(CallerTmaWorkspaceSlot(arg_tma_descriptor_workspace, 1u)); + CUtensorMap h_tmap_compressed_kv = EncodeTma_tmap_compressed_kv(arg_tmap_compressed_kv); + CakeTensorMap const* p_tmap_compressed_kv = static_cast(CallerTmaWorkspaceSlot(arg_tma_descriptor_workspace, 2u)); + __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); + float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); + int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); + int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); + int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); + float* p_sinks = static_cast(arg_sinks.data_ptr()); + float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); + float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); + int32_t v_num_heads = (int32_t)arg_num_heads; + int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; + int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; + int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; + int32_t v_num_query_tokens = (int32_t)arg_num_query_tokens; + int32_t v_sparse_topk = (int32_t)arg_sparse_topk; + int32_t v_has_sinks = (int32_t)arg_has_sinks; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_tmap_compressed_kv, &p_O, &p_partial_lse, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_num_query_tokens, &v_sparse_topk, &v_has_sinks, &v_total_work_items}; + + static std::mutex smem_mu; + static auto* smem_status_by_device = + new std::unordered_map(); + cudaError_t smem_status = cudaSuccess; + { + std::lock_guard lock(smem_mu); + auto it = smem_status_by_device->find(dev.device_id); + if (it == smem_status_by_device->end()) { + int current_device = -1; + smem_status = cudaGetDevice(¤t_device); + if (smem_status == cudaSuccess && current_device != dev.device_id) { + smem_status = cudaErrorInvalidDevice; + } + if (smem_status == cudaSuccess) { + smem_status = cudaFuncSetAttribute( + reinterpret_cast(kernel_cake_dsv4_868bb2291a1464f14e05), + cudaFuncAttributeMaxDynamicSharedMemorySize, + 205824); + } + smem_status_by_device->emplace(dev.device_id, smem_status); + } else { + smem_status = it->second; + } + } + TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) + << "cudaFuncSetAttribute for kernel_cake_dsv4_868bb2291a1464f14e05 failed: " + << cudaGetErrorString(smem_status); + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(512u, 1u, 1u); + const CUtensorMap workspace_maps[] = {h_tmap_q, h_tmap_swa_kv, h_tmap_compressed_kv}; + PrepareImmutableCallerTmaWorkspace(arg_tma_descriptor_workspace, workspace_maps, stream); + + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes = 205824u; + config.stream = stream; + cudaError_t launch_status = cudaLaunchKernelExC( + &config, reinterpret_cast(kernel_cake_dsv4_868bb2291a1464f14e05), kargs); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernelExC for kernel_cake_dsv4_868bb2291a1464f14e05 failed: " + << cudaGetErrorString(launch_status); +} + +} // namespace cake_host_shim_262e38a78b1a7fce + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_262e38a78b1a7fce::Run); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_868bb2291a1464f14e05_kernel.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_868bb2291a1464f14e05_kernel.cu new file mode 100644 index 00000000000..47307e2798a --- /dev/null +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_868bb2291a1464f14e05_kernel.cu @@ -0,0 +1,2920 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); +template +struct __align__(128) CakeTensorMapPack { CakeTensorMap maps[N]; }; + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +#define CAKE_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_SCRATCH_OFFSET 0 +#define NUM_K_PIPE_STAGES 1 +#define NUM_V_PIPE_STAGES 1 +#define NUM_KV_PIPE_STAGES 8 +#define NUM_INDEX_PIPE_STAGES 1 +#define SMEM_SMEM_Q_OFF 1024 +#define SMEM_SMEM_Q_STAGE_BYTES 16384 +#define SMEM_SMEM_Q_STRIDE 16384 +#define SMEM_SMEM_KV_OFF 66560 +#define SMEM_SMEM_KV_STAGE_BYTES 16384 +#define SMEM_SMEM_KV_STRIDE 16384 +#define SMEM_SMEM_V_OFF 66560 +#define SMEM_SMEM_V_STAGE_BYTES 16384 +#define SMEM_SMEM_V_STRIDE 16384 +#define SMEM_SMEM_Q_FULL_OFF 1024 +#define SMEM_SMEM_Q_FULL_STAGE_BYTES 32768 +#define SMEM_SMEM_Q_FULL_STRIDE 32768 +#define SMEM_SMEM_K_FULL_OFF 66560 +#define SMEM_SMEM_K_FULL_STAGE_BYTES 32768 +#define SMEM_SMEM_K_FULL_STRIDE 32768 +#define SMEM_SMEM_V_FULL_OFF 132096 +#define SMEM_SMEM_V_FULL_STAGE_BYTES 32768 +#define SMEM_SMEM_V_FULL_STRIDE 32768 +#define SMEM_SMEM_STATS_MAX_OFF 197632 +#define SMEM_SMEM_STATS_MAX_STAGE_BYTES 1024 +#define SMEM_SMEM_STATS_MAX_STRIDE 1024 +#define SMEM_SMEM_STATS_SUM_OFF 198656 +#define SMEM_SMEM_STATS_SUM_STAGE_BYTES 512 +#define SMEM_SMEM_STATS_SUM_STRIDE 512 +#define SMEM_SMEM_STATS_FINAL_MAX_OFF 199168 +#define SMEM_SMEM_STATS_FINAL_MAX_STAGE_BYTES 512 +#define SMEM_SMEM_STATS_FINAL_MAX_STRIDE 512 +#define SMEM_SMEM_SOFTMAX_WARP_PAIR_EXCHANGE_OFF 199680 +#define SMEM_SMEM_SOFTMAX_WARP_PAIR_EXCHANGE_STAGE_BYTES 1024 +#define SMEM_SMEM_SOFTMAX_WARP_PAIR_EXCHANGE_STRIDE 1024 +#define SMEM_SMEM_CORR_WARP_PAIR_EXCHANGE_OFF 200704 +#define SMEM_SMEM_CORR_WARP_PAIR_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_CORR_WARP_PAIR_EXCHANGE_STRIDE 512 +#define SMEM_SMEM_SPARSE_INDICES_OFF 201216 +#define SMEM_SMEM_SPARSE_INDICES_STAGE_BYTES 4608 +#define SMEM_SMEM_SPARSE_INDICES_STRIDE 4608 +#define SMEM_SMEM_P_FP8_OFF 1024 +#define SMEM_SMEM_P_FP8_STAGE_BYTES 8192 +#define SMEM_SMEM_P_FP8_STRIDE 8192 +#define SMEM_TOTAL 205824 +#define THREADS 512 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count) : "memory"); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f8f6f4( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x4(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x4.b32" + " {%0, %1, %2, %3}, [%4];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); +} + + +__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ float2 ex2_emulation_f32x2_value(float2 value) { + const float c0 = 1.0f, c1 = 0.695146143436431884765625f; + const float c2 = 0.227564394474029541015625f, c3 = 0.077119089663028717041015625f; + const float magic = 12582912.0f; + float x0 = max_noftz(value.x, -127.0f), x1 = max_noftz(value.y, -127.0f); + float2 xc2 = make_float2(x0, x1), magic2 = make_float2(magic, magic); + float2 xr2; + asm("add.rm.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xr2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&magic2)); + float2 c3_2 = make_float2(c3, c3), c2_2 = make_float2(c2, c2); + float2 c1_2 = make_float2(c1, c1), c0_2 = make_float2(c0, c0); + float2 xrb2, xfrac2; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xrb2) + : "l"(*(unsigned long long*)&xr2), "l"(*(unsigned long long*)&magic2)); + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xfrac2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&xrb2)); + float2 poly2; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&c3_2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c2_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c1_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c0_2)); + int x0r_i, x1r_i, p0_i, p1_i; + asm("mov.b64 {%0, %1}, %2;" : "=r"(x0r_i), "=r"(x1r_i) : "l"(*(unsigned long long*)&xr2)); + asm("mov.b64 {%0, %1}, %2;" : "=r"(p0_i), "=r"(p1_i) : "l"(*(unsigned long long*)&poly2)); + float r0, r1; + asm("mov.b32 %0, %1;" : "=f"(r0) : "r"((x0r_i << 23) + p0_i)); + asm("mov.b32 %0, %1;" : "=f"(r1) : "r"((x1r_i << 23) + p1_i)); + return make_float2(r0, r1); +} + +__device__ __forceinline__ void ex2_emulation_f32x2(float* x0_ptr, float* x1_ptr) { + float2 result = ex2_emulation_f32x2_value(make_float2(*x0_ptr, *x1_ptr)); + *x0_ptr = result.x; *x1_ptr = result.y; +} + +__device__ __forceinline__ void softmax_frag_exp2_cast( + float* sv, uint32_t* pv, int use_emu) +{ + #pragma unroll + for (int j = 0; j < 16; j++) { + if (use_emu && j >= 12) + ex2_emulation_f32x2(&sv[j*2], &sv[j*2+1]); + else { + sv[j*2] = approx_exp2(sv[j*2]); + sv[j*2+1] = approx_exp2(sv[j*2+1]); + } + } + #pragma unroll + for (int j = 0; j < 16; j++) { + __nv_bfloat162 bf = __float22bfloat162_rn({sv[j*2], sv[j*2+1]}); + pv[j] = reinterpret_cast(bf); + } +} + + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_2d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3}], [%4];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr) { + // Canonical .shared::cta form for non-multicast gather4, matching + // trtllm-gen / cuda_ptx and the PTX ISA qualifier order + // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, + // .shared::cluster is reserved for the multicast variant (ctaMask). + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem_cta2( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr) { + // Canonical .shared::cta form; see tma_gather4_gmem2smem above. + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem_mc( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr, unsigned short cta_mask) { + // Multicast variant: the PTX grammar ties the .shared::cluster + // destination to .multicast::cluster + ctaMask (cf. cuda_ptx / + // SM100_TMA_LOAD_MULTICAST_2D_GATHER4 in CUTLASS). + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" + ".mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem_mc_cta2( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr, unsigned short cta_mask) { + // Multicast + cta_group::2 variant; see tma_gather4_gmem2smem_mc. + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" + ".mbarrier::complete_tx::bytes.multicast::cluster.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_st_x8_u32(int addr, const uint32_t* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x8.b32" + " [%0], {%1,%2,%3,%4,%5,%6,%7,%8};" + :: "r"(addr), + "r"(src[0]), "r"(src[1]), "r"(src[2]), "r"(src[3]), + "r"(src[4]), "r"(src[5]), "r"(src[6]), "r"(src[7])); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(512, 1) void +kernel_cake_dsv4_868bb2291a1464f14e05(CakeTensorMap const* tmap_q, CakeTensorMap const* tmap_swa_kv, CakeTensorMap const* tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int mbar_base = smem; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define k_full_addr (mbar_base + 16) + #define k_empty_addr (mbar_base + 24) + #define v_full_addr (mbar_base + 32) + #define v_empty_addr (mbar_base + 40) + #define kv_full_addr (mbar_base + 48) + #define kv_empty_addr (mbar_base + 112) + #define index_full_addr (mbar_base + 176) + #define index_empty_addr (mbar_base + 184) + #define s_full_addr (mbar_base + 192) + #define p_full_addr (mbar_base + 208) + #define corr_done_addr (mbar_base + 224) + #define stats_addr (mbar_base + 240) + #define sum_ready_addr (mbar_base + 256) + #define o_done_addr (mbar_base + 264) + #define pv_done_addr (mbar_base + 272) + #define s_seeded_addr (mbar_base + 280) + #define q_pair_ready_addr (mbar_base + 288) + #define kv_pair_ready_addr (mbar_base + 296) + #define pv_pair_ready_addr (mbar_base + 360) + #define tmem_dealloc_addr (mbar_base + 376) + #define tmem_dealloc_peer_addr (mbar_base + 384) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + if (tid == 0) { + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_q)) : "memory"); + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_swa_kv)) : "memory"); + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_compressed_kv)) : "memory"); + } + __syncthreads(); + + + // Kernel setup ops + uint8_t* smem_q = reinterpret_cast(smem_raw + 1024); + const int smem_q_addr = smem + 1024; + uint8_t* smem_kv = reinterpret_cast(smem_raw + 66560); + const int smem_kv_addr = smem + 66560; + uint8_t* smem_v = reinterpret_cast(smem_raw + 66560); + const int smem_v_addr = smem + 66560; + uint8_t* smem_q_full = reinterpret_cast(smem_raw + 1024); + const int smem_q_full_addr = smem + 1024; + uint8_t* smem_k_full = reinterpret_cast(smem_raw + 66560); + const int smem_k_full_addr = smem + 66560; + uint8_t* smem_v_full = reinterpret_cast(smem_raw + 132096); + const int smem_v_full_addr = smem + 132096; + float* smem_stats_max = reinterpret_cast(smem_raw + 197632); + const int smem_stats_max_addr = smem + 197632; + float* smem_stats_sum = reinterpret_cast(smem_raw + 198656); + const int smem_stats_sum_addr = smem + 198656; + float* smem_stats_final_max = reinterpret_cast(smem_raw + 199168); + const int smem_stats_final_max_addr = smem + 199168; + float* smem_softmax_warp_pair_exchange = reinterpret_cast(smem_raw + 199680); + const int smem_softmax_warp_pair_exchange_addr = smem + 199680; + float* smem_corr_warp_pair_exchange = reinterpret_cast(smem_raw + 200704); + const int smem_corr_warp_pair_exchange_addr = smem + 200704; + int* smem_sparse_indices = reinterpret_cast(smem_raw + 201216); + const int smem_sparse_indices_addr = smem + 201216; + uint8_t* smem_p_fp8 = reinterpret_cast(smem_raw + 1024); + const int smem_p_fp8_addr = smem + 1024; + + // Mbarrier init (23 pipeline groups, 0 ordered-sequence groups, 49 barriers) + // Mbarriers at smem_raw[0..392) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // q_full: 1 barriers, init_count=1 + mbarrier_init(smem + 0, 1); + // q_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 8, 1); + // --- pipeline 'k_pipe' --- + // k_full: 1 barriers, init_count=1 + mbarrier_init(smem + 16, 1); + // k_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 24, 1); + // --- pipeline 'v_pipe' --- + // v_full: 1 barriers, init_count=1 + mbarrier_init(smem + 32, 1); + // v_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 40, 1); + // --- pipeline 'kv_pipe' --- + // kv_full: 8 barriers, init_count=1 + mbarrier_init(smem + 48, 1); + mbarrier_init(smem + 56, 1); + mbarrier_init(smem + 64, 1); + mbarrier_init(smem + 72, 1); + mbarrier_init(smem + 80, 1); + mbarrier_init(smem + 88, 1); + mbarrier_init(smem + 96, 1); + mbarrier_init(smem + 104, 1); + // kv_empty: 8 barriers, init_count=1 + mbarrier_init(smem + 112, 1); + mbarrier_init(smem + 120, 1); + mbarrier_init(smem + 128, 1); + mbarrier_init(smem + 136, 1); + mbarrier_init(smem + 144, 1); + mbarrier_init(smem + 152, 1); + mbarrier_init(smem + 160, 1); + mbarrier_init(smem + 168, 1); + // --- pipeline 'index_pipe' --- + // index_full: 1 barriers, init_count=1 + mbarrier_init(smem + 176, 1); + // index_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 184, 1); + // s_full: 2 barriers, init_count=1 + mbarrier_init(smem + 192, 1); + mbarrier_init(smem + 200, 1); + // p_full: 2 barriers, init_count=128 + mbarrier_init(smem + 208, 128); + mbarrier_init(smem + 216, 128); + // corr_done: 2 barriers, init_count=128 + mbarrier_init(smem + 224, 128); + mbarrier_init(smem + 232, 128); + // stats: 2 barriers, init_count=128 + mbarrier_init(smem + 240, 128); + mbarrier_init(smem + 248, 128); + // sum_ready: 1 barriers, init_count=128 + mbarrier_init(smem + 256, 128); + // o_done: 1 barriers, init_count=1 + mbarrier_init(smem + 264, 1); + // pv_done: 1 barriers, init_count=1 + mbarrier_init(smem + 272, 1); + // s_seeded: 1 barriers, init_count=128 + mbarrier_init(smem + 280, 128); + // q_pair_ready: 1 barriers, init_count=64 + mbarrier_init(smem + 288, 64); + // --- pipeline 'kv_pipe' --- + // kv_pair_ready: 8 barriers, init_count=64 + mbarrier_init(smem + 296, 64); + mbarrier_init(smem + 304, 64); + mbarrier_init(smem + 312, 64); + mbarrier_init(smem + 320, 64); + mbarrier_init(smem + 328, 64); + mbarrier_init(smem + 336, 64); + mbarrier_init(smem + 344, 64); + mbarrier_init(smem + 352, 64); + // pv_pair_ready: 2 barriers, init_count=64 + mbarrier_init(smem + 360, 64); + mbarrier_init(smem + 368, 64); + // tmem_dealloc: 1 barriers, init_count=416 + mbarrier_init(smem + 376, 416); + // tmem_dealloc_peer: 1 barriers, init_count=32 + mbarrier_init(smem + 384, 32); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncwarp(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 392); + if (warp == 0) { + int _tmem_hold = smem + 392; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_tmem_scratch = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 15) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 64;"); + } + + // ---- Role: index_warp ---- + if (warp == 13) { + { // index_warp_main + const int index_dummy = 0; + unsigned int _phase_index_empty = 1; + unsigned int _phase_q_full_0 = 0; + { + #pragma unroll 1 + for (unsigned int work_idx = bid; work_idx < total_work_items; work_idx += num_bids) { + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + } + } + } + } + // ---- Role: load_warp ---- + if (warp >= 9 && warp <= 12) { + { // load_warp_main + const int wg2_dummy = 0; + const int load_warp_rank = warp - ((0) ? 12 : 9); + int all_num_kv_tiles = (sparse_topk + 128 - 1) / 128; + unsigned int load_k_stage = 0; + unsigned int load_v_stage = 0; + unsigned int load_kv_stage = 0; + unsigned int load_k_index_stage = 0; + unsigned int load_v_index_stage = 0; + int k_cta_offset = 0; + unsigned int _phase_q_empty_0 = 1; + unsigned int _phase_index_full = 0; + unsigned int _phase_k_empty = 1; + unsigned int _phase_kv_empty = 1; + unsigned int _phase_v_empty = 1; + #pragma unroll 1 + for (unsigned int work_idx_1 = bid; work_idx_1 < total_work_items; work_idx_1 += num_bids) { + int split_idx = 0; + int query_idx = work_idx_1 >> 1; + int v_chunk = work_idx_1 & 1; + int tiles_per_split = all_num_kv_tiles; + int first_tile = split_idx * tiles_per_split; + int num_kv_tiles = tiles_per_split; + int sparse_extent = num_kv_tiles * 128; + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + if (load_warp_rank == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(q_full_addr, 65536); + #pragma unroll + for (int q_stage = 0; q_stage < 4; q_stage++) { + tma_4d_gmem2smem(smem_q_addr + (unsigned int)(q_stage * 16384), tmap_q, 0, 0, q_stage, query_idx, q_full_addr); + } + } + } + { + int num_index_passes = sparse_extent / 128; + #pragma unroll 1 + for (int index_pass = load_warp_rank; index_pass < num_index_passes; index_pass += 4) { + int index_offset = index_pass * 128 + lane * 4; + int tile_col_base = (first_tile + index_pass) * 128; + int* row_tile_ptr = ((tile_col_base < 128) ? swa_indices + (query_idx * swa_index_stride) : compressed_indices + (query_idx * compressed_index_stride + (tile_col_base - 128))); + int lane_col = tile_col_base + lane * 4; + int sparse_rows[4]; + #pragma unroll + for (int row_i = 0; row_i < 4; row_i++) { + sparse_rows[row_i] = -1; + } + if (lane_col + 4 <= sparse_topk) { + int _vec_load_0[4]; + { + const int4* _ivptr_0 = reinterpret_cast(row_tile_ptr + (lane * 4) + 0); + int4 _ivld_0; + _ivld_0 = *_ivptr_0; + _vec_load_0[0 + 0] = _ivld_0.x; + _vec_load_0[0 + 1] = _ivld_0.y; + _vec_load_0[0 + 2] = _ivld_0.z; + _vec_load_0[0 + 3] = _ivld_0.w; + } + #pragma unroll + for (int row_i_1 = 0; row_i_1 < 4; row_i_1++) { + sparse_rows[row_i_1] = _vec_load_0[row_i_1]; + } + } else { + #pragma unroll + for (int row_i_2 = 0; row_i_2 < 4; row_i_2++) { + if (lane_col + row_i_2 < sparse_topk) { + sparse_rows[row_i_2] = row_tile_ptr[lane * 4 + row_i_2]; + } + } + } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(smem_sparse_indices_addr + (unsigned int)(index_offset * 4)), "r"(sparse_rows[0]), "r"(sparse_rows[1]), "r"(sparse_rows[2]), "r"(sparse_rows[3]) : "memory"); + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("barrier.sync 8, 256;" ::: "memory"); + } + int k_index_stage_base = smem_sparse_indices_addr; + { + #pragma unroll + for (int qk_stage = 0; qk_stage < 4; qk_stage++) { + if (qk_stage == load_warp_rank) { + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (load_kv_stage) * 8, 16384); + } + int k_dst = smem_kv_addr + load_kv_stage * 16384; + int group = lane; + int k_group_limit = ((work_idx_1 == (unsigned int)bid) ? 16 : 32); + if (group < k_group_limit) { + int group_offset = k_cta_offset + group * 4; + int raw_rows[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows[0])), "=r"(*reinterpret_cast(&raw_rows[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows[(0) + 3])) + : "r"(k_index_stage_base + group_offset * 4)); + int row0 = ((raw_rows[0] >= 0) ? raw_rows[0] : 0); + int row1 = ((raw_rows[1] >= 0) ? raw_rows[1] : 0); + int row2 = ((raw_rows[2] >= 0) ? raw_rows[2] : 0); + int row3 = ((raw_rows[3] >= 0) ? raw_rows[3] : 0); + if (first_tile == 0) { + tma_gather4_gmem2smem(k_dst + group * 512, tmap_swa_kv, qk_stage * 128, row0, row1, row2, row3, kv_full_addr + (load_kv_stage) * 8); + } else { + tma_gather4_gmem2smem(k_dst + group * 512, tmap_compressed_kv, qk_stage * 128, row0, row1, row2, row3, kv_full_addr + (load_kv_stage) * 8); + } + } + } + load_kv_stage += 1; + if (load_kv_stage == 8) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + } + #pragma unroll 1 + for (int tile = 1; tile < num_kv_tiles; tile++) { + if ((0 & (int)((tile & 1) == 0)) != 0) { + mbarrier_wait(index_full_addr + (load_k_index_stage) * 8, _phase_index_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + } + k_index_stage_base = smem_sparse_indices_addr; + { + #pragma unroll + for (int qk_stage_1 = 0; qk_stage_1 < 4; qk_stage_1++) { + if (qk_stage_1 == load_warp_rank) { + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (load_kv_stage) * 8, 16384); + } + int k_dst_1 = smem_kv_addr + load_kv_stage * 16384; + int group_1 = lane; + if (group_1 < 32) { + int group_offset_1 = tile * 128 + k_cta_offset + group_1 * 4; + int raw_rows_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_1[0])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 3])) + : "r"(k_index_stage_base + group_offset_1 * 4)); + int row0_1 = ((raw_rows_1[0] >= 0) ? raw_rows_1[0] : 0); + int row1_1 = ((raw_rows_1[1] >= 0) ? raw_rows_1[1] : 0); + int row2_1 = ((raw_rows_1[2] >= 0) ? raw_rows_1[2] : 0); + int row3_1 = ((raw_rows_1[3] >= 0) ? raw_rows_1[3] : 0); + tma_gather4_gmem2smem(k_dst_1 + group_1 * 512, tmap_compressed_kv, qk_stage_1 * 128, row0_1, row1_1, row2_1, row3_1, kv_full_addr + (load_kv_stage) * 8); + } + } + load_kv_stage += 1; + if (load_kv_stage == 8) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + } + int prev_tile = tile - 1; + int global_prev_tile = first_tile + prev_tile; + int v_index_stage_base = smem_sparse_indices_addr; + { + #pragma unroll + for (int pv_stage = 0; pv_stage < 2; pv_stage++) { + if (pv_stage == load_warp_rank) { + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (load_kv_stage) * 8, 16384); + } + int v_dst = smem_kv_addr + load_kv_stage * 16384; + int v_col = v_chunk * 256 + pv_stage * 128; + int group_2 = lane; + int group_offset_2 = prev_tile * 128 + group_2 * 4; + int raw_rows_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_2[0])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 3])) + : "r"(v_index_stage_base + group_offset_2 * 4)); + int row0_2 = ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : 0); + int row1_2 = ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : 0); + int row2_2 = ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : 0); + int row3_2 = ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : 0); + if (global_prev_tile == 0) { + tma_gather4_gmem2smem(v_dst + group_2 * 512, tmap_swa_kv, v_col, row0_2, row1_2, row2_2, row3_2, kv_full_addr + (load_kv_stage) * 8); + } else { + tma_gather4_gmem2smem(v_dst + group_2 * 512, tmap_compressed_kv, v_col, row0_2, row1_2, row2_2, row3_2, kv_full_addr + (load_kv_stage) * 8); + } + } + load_kv_stage += 1; + if (load_kv_stage == 8) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + } + } + int last_tile = num_kv_tiles - 1; + int global_last_tile = first_tile + last_tile; + { + #pragma unroll + for (int pv_stage_1 = 0; pv_stage_1 < 2; pv_stage_1++) { + if (pv_stage_1 == load_warp_rank) { + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (load_kv_stage) * 8, 16384); + } + int v_dst_1 = smem_kv_addr + load_kv_stage * 16384; + int v_col_1 = v_chunk * 256 + pv_stage_1 * 128; + int group_3 = lane; + int group_offset_3 = last_tile * 128 + group_3 * 4; + int raw_rows_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_3[0])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 3])) + : "r"(smem_sparse_indices_addr + (unsigned int)(group_offset_3 * 4))); + int row0_3 = ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : 0); + int row1_3 = ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : 0); + int row2_3 = ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : 0); + int row3_3 = ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : 0); + if (global_last_tile == 0) { + tma_gather4_gmem2smem(v_dst_1 + group_3 * 512, tmap_swa_kv, v_col_1, row0_3, row1_3, row2_3, row3_3, kv_full_addr + (load_kv_stage) * 8); + } else { + tma_gather4_gmem2smem(v_dst_1 + group_3 * 512, tmap_compressed_kv, v_col_1, row0_3, row1_3, row2_3, row3_3, kv_full_addr + (load_kv_stage) * 8); + } + } + load_kv_stage += 1; + if (load_kv_stage == 8) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + } + } + mbarrier_arrive(tmem_dealloc_addr); + } + } + // ---- Role: softmax_wg ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 184;"); + { // softmax_wg_main + float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; + const int wg_dummy_inc = 0; + int all_num_kv_tiles_1 = (sparse_topk + 128 - 1) / 128; + const int tmem_row_base = ((0) ? warp % 2 * 32 : warp % 4 * 32); + const int tmem_score_row_base = ((0) ? (int)(warp % 4 * 32) : tmem_row_base); + const int n_half = ((0) ? (int)(warp % 4 / 2) : 0); + const int my_row = tmem_row_base + lane; + const int stats_row = n_half * 64 + my_row; + int softmax_tile_cursor = 0; + unsigned int softmax_index_stage = 0; + unsigned int _phase_q_full_0_1 = 0; + #pragma unroll 1 + for (unsigned int work_idx_2 = bid; work_idx_2 < total_work_items; work_idx_2 += num_bids) { + int split_idx_1 = 0; + int query_idx_1 = work_idx_2 >> 1; + int tiles_per_split_1 = all_num_kv_tiles_1; + int first_tile_1 = split_idx_1 * tiles_per_split_1; + int num_kv_tiles_1 = tiles_per_split_1; + int _max_0 = ((sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) : (0)); + int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); + int active_topk = _min_0; + { + float seed_zero[4]; + #pragma unroll + for (int seed_c4 = 0; seed_c4 < 4; seed_c4++) { + seed_zero[seed_c4] = 0.0f; + } + #pragma unroll + for (int seed_half = 0; seed_half < 2; seed_half++) { + #pragma unroll + for (int seed_c = 0; seed_c < 8; seed_c++) { + int seed_addr = taddr + (unsigned int)(((seed_half == 0) ? 96 : 224)) + (unsigned int)(seed_c * 4) + (unsigned int)(tmem_row_base << 16); + tmem_st_x4_f32(seed_addr, seed_zero); + } + } + { + #pragma unroll + for (int seed_half_1 = 0; seed_half_1 < 2; seed_half_1++) { + #pragma unroll + for (int seed_c_1 = 16; seed_c_1 < 24; seed_c_1++) { + int seed_addr_1 = taddr + (unsigned int)(((seed_half_1 == 0) ? 0 : 128)) + (unsigned int)(seed_c_1 * 4) + (unsigned int)(tmem_row_base << 16); + tmem_st_x4_f32(seed_addr_1, seed_zero); + } + } + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(s_seeded_addr); + } + { + asm volatile("barrier.sync 8, 256;" ::: "memory"); + if (work_idx_2 == (unsigned int)bid) { + const int helper_stage = warp % 4; + int helper_dst = smem_kv_addr + (unsigned int)(helper_stage * 16384); + int helper_group = lane; + if (helper_group >= 16) { + int helper_rows[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&helper_rows[0])), "=r"(*reinterpret_cast(&helper_rows[(0) + 1])), "=r"(*reinterpret_cast(&helper_rows[(0) + 2])), "=r"(*reinterpret_cast(&helper_rows[(0) + 3])) + : "r"(smem_sparse_indices_addr + (unsigned int)(helper_group * 16))); + int hrow0 = ((helper_rows[0] >= 0) ? helper_rows[0] : 0); + int hrow1 = ((helper_rows[1] >= 0) ? helper_rows[1] : 0); + int hrow2 = ((helper_rows[2] >= 0) ? helper_rows[2] : 0); + int hrow3 = ((helper_rows[3] >= 0) ? helper_rows[3] : 0); + if (first_tile_1 == 0) { + tma_gather4_gmem2smem(helper_dst + helper_group * 512, tmap_swa_kv, helper_stage * 128, hrow0, hrow1, hrow2, hrow3, kv_full_addr + (helper_stage) * 8); + } else { + tma_gather4_gmem2smem(helper_dst + helper_group * 512, tmap_compressed_kv, helper_stage * 128, hrow0, hrow1, hrow2, hrow3, kv_full_addr + (helper_stage) * 8); + } + } + } + } + { + mbarrier_wait(q_full_addr, _phase_q_full_0_1); + _phase_q_full_0_1 ^= 1; + } + float row_max_val = -CAKE_INF; + float row_sum_val = 0.0f; + int sink_head = ((0) ? my_row : my_row); + if (has_sinks != 0 && sink_head < num_heads && split_idx_1 == 0) { + row_max_val = sinks[sink_head] * 1.4426950408889634f / softmax_scale_log2; + row_sum_val = 1.0f; + } + #pragma unroll 1 + for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { + int pipeline_tile = softmax_tile_cursor + tile_1; + int phase = pipeline_tile & 1; + int s_wait_phase = pipeline_tile >> 1 & 1; + mbarrier_wait(s_full_addr + (phase) * 8, s_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int s_off = ((phase != 0) ? 128 : 0); + int s_base = taddr + (unsigned int)s_off + (unsigned int)(tmem_score_row_base << 16); + float new_max = row_max_val; + int valid_sparse_cols = ((active_topk < sparse_topk) ? active_topk : sparse_topk); + valid_sparse_cols = valid_sparse_cols - (first_tile_1 + tile_1) * 128 - n_half * 64; + if (valid_sparse_cols < 0) { + valid_sparse_cols = 0; + } + if (valid_sparse_cols > 128) { + valid_sparse_cols = 128; + } + int index_tile_addr = smem_sparse_indices_addr + (unsigned int)(tile_1 * 512); + int index_lane_addr = index_tile_addr + (n_half * 64 + lane) * 4; + int staged_index[4]; + unsigned int invalid_cols[4]; + unsigned int any_invalid = 0; + #pragma unroll + for (int w = 0; w < 4; w++) { + asm volatile("ld.shared.b32 %0, [%1];" : "=r"(*reinterpret_cast(&staged_index[w])) : "r"(index_lane_addr + w * 128)); + unsigned int _vote_0 = __ballot_sync(0xFFFFFFFF, staged_index[w] < 0); + invalid_cols[w] = _vote_0; + any_invalid = any_invalid | invalid_cols[w]; + } + float _tmem_load_0[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]) + : "r"(s_base)); + float _tmem_load_1[4]; + tmem_ld_x4(&_tmem_load_1[0], s_base); + float _tmem_load_2[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]) + : "r"(s_base + 32)); + float _tmem_load_3[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]) + : "r"(s_base + 64)); + float _tmem_load_4[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_4[0]), "=f"(_tmem_load_4[1]), "=f"(_tmem_load_4[2]), "=f"(_tmem_load_4[3]), "=f"(_tmem_load_4[4]), "=f"(_tmem_load_4[5]), "=f"(_tmem_load_4[6]), "=f"(_tmem_load_4[7]), "=f"(_tmem_load_4[8]), "=f"(_tmem_load_4[9]), "=f"(_tmem_load_4[10]), "=f"(_tmem_load_4[11]), "=f"(_tmem_load_4[12]), "=f"(_tmem_load_4[13]), "=f"(_tmem_load_4[14]), "=f"(_tmem_load_4[15]), "=f"(_tmem_load_4[16]), "=f"(_tmem_load_4[17]), "=f"(_tmem_load_4[18]), "=f"(_tmem_load_4[19]), "=f"(_tmem_load_4[20]), "=f"(_tmem_load_4[21]), "=f"(_tmem_load_4[22]), "=f"(_tmem_load_4[23]), "=f"(_tmem_load_4[24]), "=f"(_tmem_load_4[25]), "=f"(_tmem_load_4[26]), "=f"(_tmem_load_4[27]), "=f"(_tmem_load_4[28]), "=f"(_tmem_load_4[29]), "=f"(_tmem_load_4[30]), "=f"(_tmem_load_4[31]) + : "r"(s_base + 96)); + { + int frag_valid_0 = valid_sparse_cols; + uint32_t _slice_lo_mask_1; + { + int _lim_0 = frag_valid_0; + if (_lim_0 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) _tmem_load_0[0] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) _tmem_load_0[1] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) _tmem_load_0[2] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) _tmem_load_0[3] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) _tmem_load_0[4] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) _tmem_load_0[5] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) _tmem_load_0[6] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) _tmem_load_0[7] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) _tmem_load_0[8] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) _tmem_load_0[9] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) _tmem_load_0[10] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) _tmem_load_0[11] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) _tmem_load_0[12] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) _tmem_load_0[13] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) _tmem_load_0[14] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) _tmem_load_0[15] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) _tmem_load_0[16] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) _tmem_load_0[17] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) _tmem_load_0[18] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) _tmem_load_0[19] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) _tmem_load_0[20] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) _tmem_load_0[21] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) _tmem_load_0[22] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) _tmem_load_0[23] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) _tmem_load_0[24] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) _tmem_load_0[25] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) _tmem_load_0[26] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) _tmem_load_0[27] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) _tmem_load_0[28] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) _tmem_load_0[29] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) _tmem_load_0[30] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) _tmem_load_0[31] = -CAKE_INF; + if (any_invalid != 0) { + unsigned int frag_mask_word_0 = invalid_cols[0]; + _tmem_load_0[0] = (((frag_mask_word_0 & 1) != 0) ? -CAKE_INF : _tmem_load_0[0]); + _tmem_load_0[1] = (((frag_mask_word_0 >> 1 & 1) != 0) ? -CAKE_INF : _tmem_load_0[1]); + _tmem_load_0[2] = (((frag_mask_word_0 >> 2 & 1) != 0) ? -CAKE_INF : _tmem_load_0[2]); + _tmem_load_0[3] = (((frag_mask_word_0 >> 3 & 1) != 0) ? -CAKE_INF : _tmem_load_0[3]); + _tmem_load_0[4] = (((frag_mask_word_0 >> 4 & 1) != 0) ? -CAKE_INF : _tmem_load_0[4]); + _tmem_load_0[5] = (((frag_mask_word_0 >> 5 & 1) != 0) ? -CAKE_INF : _tmem_load_0[5]); + _tmem_load_0[6] = (((frag_mask_word_0 >> 6 & 1) != 0) ? -CAKE_INF : _tmem_load_0[6]); + _tmem_load_0[7] = (((frag_mask_word_0 >> 7 & 1) != 0) ? -CAKE_INF : _tmem_load_0[7]); + _tmem_load_0[8] = (((frag_mask_word_0 >> 8 & 1) != 0) ? -CAKE_INF : _tmem_load_0[8]); + _tmem_load_0[9] = (((frag_mask_word_0 >> 9 & 1) != 0) ? -CAKE_INF : _tmem_load_0[9]); + _tmem_load_0[10] = (((frag_mask_word_0 >> 10 & 1) != 0) ? -CAKE_INF : _tmem_load_0[10]); + _tmem_load_0[11] = (((frag_mask_word_0 >> 11 & 1) != 0) ? -CAKE_INF : _tmem_load_0[11]); + _tmem_load_0[12] = (((frag_mask_word_0 >> 12 & 1) != 0) ? -CAKE_INF : _tmem_load_0[12]); + _tmem_load_0[13] = (((frag_mask_word_0 >> 13 & 1) != 0) ? -CAKE_INF : _tmem_load_0[13]); + _tmem_load_0[14] = (((frag_mask_word_0 >> 14 & 1) != 0) ? -CAKE_INF : _tmem_load_0[14]); + _tmem_load_0[15] = (((frag_mask_word_0 >> 15 & 1) != 0) ? -CAKE_INF : _tmem_load_0[15]); + _tmem_load_0[16] = (((frag_mask_word_0 >> 16 & 1) != 0) ? -CAKE_INF : _tmem_load_0[16]); + _tmem_load_0[17] = (((frag_mask_word_0 >> 17 & 1) != 0) ? -CAKE_INF : _tmem_load_0[17]); + _tmem_load_0[18] = (((frag_mask_word_0 >> 18 & 1) != 0) ? -CAKE_INF : _tmem_load_0[18]); + _tmem_load_0[19] = (((frag_mask_word_0 >> 19 & 1) != 0) ? -CAKE_INF : _tmem_load_0[19]); + _tmem_load_0[20] = (((frag_mask_word_0 >> 20 & 1) != 0) ? -CAKE_INF : _tmem_load_0[20]); + _tmem_load_0[21] = (((frag_mask_word_0 >> 21 & 1) != 0) ? -CAKE_INF : _tmem_load_0[21]); + _tmem_load_0[22] = (((frag_mask_word_0 >> 22 & 1) != 0) ? -CAKE_INF : _tmem_load_0[22]); + _tmem_load_0[23] = (((frag_mask_word_0 >> 23 & 1) != 0) ? -CAKE_INF : _tmem_load_0[23]); + _tmem_load_0[24] = (((frag_mask_word_0 >> 24 & 1) != 0) ? -CAKE_INF : _tmem_load_0[24]); + _tmem_load_0[25] = (((frag_mask_word_0 >> 25 & 1) != 0) ? -CAKE_INF : _tmem_load_0[25]); + _tmem_load_0[26] = (((frag_mask_word_0 >> 26 & 1) != 0) ? -CAKE_INF : _tmem_load_0[26]); + _tmem_load_0[27] = (((frag_mask_word_0 >> 27 & 1) != 0) ? -CAKE_INF : _tmem_load_0[27]); + _tmem_load_0[28] = (((frag_mask_word_0 >> 28 & 1) != 0) ? -CAKE_INF : _tmem_load_0[28]); + _tmem_load_0[29] = (((frag_mask_word_0 >> 29 & 1) != 0) ? -CAKE_INF : _tmem_load_0[29]); + _tmem_load_0[30] = (((frag_mask_word_0 >> 30 & 1) != 0) ? -CAKE_INF : _tmem_load_0[30]); + _tmem_load_0[31] = (((frag_mask_word_0 >> 31 & 1) != 0) ? -CAKE_INF : _tmem_load_0[31]); + } + float2 _reg_reduce_max2_1 = {-CAKE_INF, -CAKE_INF}; + _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_0[0], _tmem_load_0[1])); + _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_0[2], _tmem_load_0[3])); + _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_0[4], _tmem_load_0[5])); + _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_0[6], _tmem_load_0[7])); + _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_0[8], _tmem_load_0[9])); + _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_0[10], _tmem_load_0[11])); + _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_0[12], _tmem_load_0[13])); + _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_0[14], _tmem_load_0[15])); + _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_0[16], _tmem_load_0[17])); + _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_0[18], _tmem_load_0[19])); + _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_0[20], _tmem_load_0[21])); + _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_0[22], _tmem_load_0[23])); + _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_0[24], _tmem_load_0[25])); + _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_0[26], _tmem_load_0[27])); + _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_0[28], _tmem_load_0[29])); + _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_0[30], _tmem_load_0[31])); + float _tmem_load_0_max = row_max_reduce(_reg_reduce_max2_1); + float _max_3 = max_noftz(new_max, _tmem_load_0_max); + new_max = _max_3; + int frag_valid_1 = valid_sparse_cols - 32; + uint32_t _slice_lo_mask_2; + { + int _lim_2 = frag_valid_1; + if (_lim_2 <= 0) { _slice_lo_mask_2 = 0u; } + else if (_lim_2 >= 32) { _slice_lo_mask_2 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_2) : "r"(_lim_2)); + } + } + if (!(_slice_lo_mask_2 & (1u << 0))) _tmem_load_2[0] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 1))) _tmem_load_2[1] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 2))) _tmem_load_2[2] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 3))) _tmem_load_2[3] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 4))) _tmem_load_2[4] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 5))) _tmem_load_2[5] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 6))) _tmem_load_2[6] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 7))) _tmem_load_2[7] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 8))) _tmem_load_2[8] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 9))) _tmem_load_2[9] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 10))) _tmem_load_2[10] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 11))) _tmem_load_2[11] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 12))) _tmem_load_2[12] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 13))) _tmem_load_2[13] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 14))) _tmem_load_2[14] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 15))) _tmem_load_2[15] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 16))) _tmem_load_2[16] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 17))) _tmem_load_2[17] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 18))) _tmem_load_2[18] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 19))) _tmem_load_2[19] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 20))) _tmem_load_2[20] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 21))) _tmem_load_2[21] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 22))) _tmem_load_2[22] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 23))) _tmem_load_2[23] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 24))) _tmem_load_2[24] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 25))) _tmem_load_2[25] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 26))) _tmem_load_2[26] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 27))) _tmem_load_2[27] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 28))) _tmem_load_2[28] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 29))) _tmem_load_2[29] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 30))) _tmem_load_2[30] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 31))) _tmem_load_2[31] = -CAKE_INF; + if (any_invalid != 0) { + unsigned int frag_mask_word_1 = invalid_cols[1]; + _tmem_load_2[0] = (((frag_mask_word_1 & 1) != 0) ? -CAKE_INF : _tmem_load_2[0]); + _tmem_load_2[1] = (((frag_mask_word_1 >> 1 & 1) != 0) ? -CAKE_INF : _tmem_load_2[1]); + _tmem_load_2[2] = (((frag_mask_word_1 >> 2 & 1) != 0) ? -CAKE_INF : _tmem_load_2[2]); + _tmem_load_2[3] = (((frag_mask_word_1 >> 3 & 1) != 0) ? -CAKE_INF : _tmem_load_2[3]); + _tmem_load_2[4] = (((frag_mask_word_1 >> 4 & 1) != 0) ? -CAKE_INF : _tmem_load_2[4]); + _tmem_load_2[5] = (((frag_mask_word_1 >> 5 & 1) != 0) ? -CAKE_INF : _tmem_load_2[5]); + _tmem_load_2[6] = (((frag_mask_word_1 >> 6 & 1) != 0) ? -CAKE_INF : _tmem_load_2[6]); + _tmem_load_2[7] = (((frag_mask_word_1 >> 7 & 1) != 0) ? -CAKE_INF : _tmem_load_2[7]); + _tmem_load_2[8] = (((frag_mask_word_1 >> 8 & 1) != 0) ? -CAKE_INF : _tmem_load_2[8]); + _tmem_load_2[9] = (((frag_mask_word_1 >> 9 & 1) != 0) ? -CAKE_INF : _tmem_load_2[9]); + _tmem_load_2[10] = (((frag_mask_word_1 >> 10 & 1) != 0) ? -CAKE_INF : _tmem_load_2[10]); + _tmem_load_2[11] = (((frag_mask_word_1 >> 11 & 1) != 0) ? -CAKE_INF : _tmem_load_2[11]); + _tmem_load_2[12] = (((frag_mask_word_1 >> 12 & 1) != 0) ? -CAKE_INF : _tmem_load_2[12]); + _tmem_load_2[13] = (((frag_mask_word_1 >> 13 & 1) != 0) ? -CAKE_INF : _tmem_load_2[13]); + _tmem_load_2[14] = (((frag_mask_word_1 >> 14 & 1) != 0) ? -CAKE_INF : _tmem_load_2[14]); + _tmem_load_2[15] = (((frag_mask_word_1 >> 15 & 1) != 0) ? -CAKE_INF : _tmem_load_2[15]); + _tmem_load_2[16] = (((frag_mask_word_1 >> 16 & 1) != 0) ? -CAKE_INF : _tmem_load_2[16]); + _tmem_load_2[17] = (((frag_mask_word_1 >> 17 & 1) != 0) ? -CAKE_INF : _tmem_load_2[17]); + _tmem_load_2[18] = (((frag_mask_word_1 >> 18 & 1) != 0) ? -CAKE_INF : _tmem_load_2[18]); + _tmem_load_2[19] = (((frag_mask_word_1 >> 19 & 1) != 0) ? -CAKE_INF : _tmem_load_2[19]); + _tmem_load_2[20] = (((frag_mask_word_1 >> 20 & 1) != 0) ? -CAKE_INF : _tmem_load_2[20]); + _tmem_load_2[21] = (((frag_mask_word_1 >> 21 & 1) != 0) ? -CAKE_INF : _tmem_load_2[21]); + _tmem_load_2[22] = (((frag_mask_word_1 >> 22 & 1) != 0) ? -CAKE_INF : _tmem_load_2[22]); + _tmem_load_2[23] = (((frag_mask_word_1 >> 23 & 1) != 0) ? -CAKE_INF : _tmem_load_2[23]); + _tmem_load_2[24] = (((frag_mask_word_1 >> 24 & 1) != 0) ? -CAKE_INF : _tmem_load_2[24]); + _tmem_load_2[25] = (((frag_mask_word_1 >> 25 & 1) != 0) ? -CAKE_INF : _tmem_load_2[25]); + _tmem_load_2[26] = (((frag_mask_word_1 >> 26 & 1) != 0) ? -CAKE_INF : _tmem_load_2[26]); + _tmem_load_2[27] = (((frag_mask_word_1 >> 27 & 1) != 0) ? -CAKE_INF : _tmem_load_2[27]); + _tmem_load_2[28] = (((frag_mask_word_1 >> 28 & 1) != 0) ? -CAKE_INF : _tmem_load_2[28]); + _tmem_load_2[29] = (((frag_mask_word_1 >> 29 & 1) != 0) ? -CAKE_INF : _tmem_load_2[29]); + _tmem_load_2[30] = (((frag_mask_word_1 >> 30 & 1) != 0) ? -CAKE_INF : _tmem_load_2[30]); + _tmem_load_2[31] = (((frag_mask_word_1 >> 31 & 1) != 0) ? -CAKE_INF : _tmem_load_2[31]); + } + float2 _reg_reduce_max2_3 = {-CAKE_INF, -CAKE_INF}; + _reg_reduce_max2_3.x = max_noftz(_reg_reduce_max2_3.x, max_noftz(_tmem_load_2[0], _tmem_load_2[1])); + _reg_reduce_max2_3.y = max_noftz(_reg_reduce_max2_3.y, max_noftz(_tmem_load_2[2], _tmem_load_2[3])); + _reg_reduce_max2_3.x = max_noftz(_reg_reduce_max2_3.x, max_noftz(_tmem_load_2[4], _tmem_load_2[5])); + _reg_reduce_max2_3.y = max_noftz(_reg_reduce_max2_3.y, max_noftz(_tmem_load_2[6], _tmem_load_2[7])); + _reg_reduce_max2_3.x = max_noftz(_reg_reduce_max2_3.x, max_noftz(_tmem_load_2[8], _tmem_load_2[9])); + _reg_reduce_max2_3.y = max_noftz(_reg_reduce_max2_3.y, max_noftz(_tmem_load_2[10], _tmem_load_2[11])); + _reg_reduce_max2_3.x = max_noftz(_reg_reduce_max2_3.x, max_noftz(_tmem_load_2[12], _tmem_load_2[13])); + _reg_reduce_max2_3.y = max_noftz(_reg_reduce_max2_3.y, max_noftz(_tmem_load_2[14], _tmem_load_2[15])); + _reg_reduce_max2_3.x = max_noftz(_reg_reduce_max2_3.x, max_noftz(_tmem_load_2[16], _tmem_load_2[17])); + _reg_reduce_max2_3.y = max_noftz(_reg_reduce_max2_3.y, max_noftz(_tmem_load_2[18], _tmem_load_2[19])); + _reg_reduce_max2_3.x = max_noftz(_reg_reduce_max2_3.x, max_noftz(_tmem_load_2[20], _tmem_load_2[21])); + _reg_reduce_max2_3.y = max_noftz(_reg_reduce_max2_3.y, max_noftz(_tmem_load_2[22], _tmem_load_2[23])); + _reg_reduce_max2_3.x = max_noftz(_reg_reduce_max2_3.x, max_noftz(_tmem_load_2[24], _tmem_load_2[25])); + _reg_reduce_max2_3.y = max_noftz(_reg_reduce_max2_3.y, max_noftz(_tmem_load_2[26], _tmem_load_2[27])); + _reg_reduce_max2_3.x = max_noftz(_reg_reduce_max2_3.x, max_noftz(_tmem_load_2[28], _tmem_load_2[29])); + _reg_reduce_max2_3.y = max_noftz(_reg_reduce_max2_3.y, max_noftz(_tmem_load_2[30], _tmem_load_2[31])); + float _tmem_load_2_max = row_max_reduce(_reg_reduce_max2_3); + float _max_4 = max_noftz(new_max, _tmem_load_2_max); + new_max = _max_4; + int frag_valid_2 = valid_sparse_cols - 64; + uint32_t _slice_lo_mask_3; + { + int _lim_4 = frag_valid_2; + if (_lim_4 <= 0) { _slice_lo_mask_3 = 0u; } + else if (_lim_4 >= 32) { _slice_lo_mask_3 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_3) : "r"(_lim_4)); + } + } + if (!(_slice_lo_mask_3 & (1u << 0))) _tmem_load_3[0] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 1))) _tmem_load_3[1] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 2))) _tmem_load_3[2] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 3))) _tmem_load_3[3] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 4))) _tmem_load_3[4] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 5))) _tmem_load_3[5] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 6))) _tmem_load_3[6] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 7))) _tmem_load_3[7] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 8))) _tmem_load_3[8] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 9))) _tmem_load_3[9] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 10))) _tmem_load_3[10] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 11))) _tmem_load_3[11] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 12))) _tmem_load_3[12] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 13))) _tmem_load_3[13] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 14))) _tmem_load_3[14] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 15))) _tmem_load_3[15] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 16))) _tmem_load_3[16] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 17))) _tmem_load_3[17] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 18))) _tmem_load_3[18] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 19))) _tmem_load_3[19] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 20))) _tmem_load_3[20] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 21))) _tmem_load_3[21] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 22))) _tmem_load_3[22] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 23))) _tmem_load_3[23] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 24))) _tmem_load_3[24] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 25))) _tmem_load_3[25] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 26))) _tmem_load_3[26] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 27))) _tmem_load_3[27] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 28))) _tmem_load_3[28] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 29))) _tmem_load_3[29] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 30))) _tmem_load_3[30] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 31))) _tmem_load_3[31] = -CAKE_INF; + if (any_invalid != 0) { + unsigned int frag_mask_word_2 = invalid_cols[2]; + _tmem_load_3[0] = (((frag_mask_word_2 & 1) != 0) ? -CAKE_INF : _tmem_load_3[0]); + _tmem_load_3[1] = (((frag_mask_word_2 >> 1 & 1) != 0) ? -CAKE_INF : _tmem_load_3[1]); + _tmem_load_3[2] = (((frag_mask_word_2 >> 2 & 1) != 0) ? -CAKE_INF : _tmem_load_3[2]); + _tmem_load_3[3] = (((frag_mask_word_2 >> 3 & 1) != 0) ? -CAKE_INF : _tmem_load_3[3]); + _tmem_load_3[4] = (((frag_mask_word_2 >> 4 & 1) != 0) ? -CAKE_INF : _tmem_load_3[4]); + _tmem_load_3[5] = (((frag_mask_word_2 >> 5 & 1) != 0) ? -CAKE_INF : _tmem_load_3[5]); + _tmem_load_3[6] = (((frag_mask_word_2 >> 6 & 1) != 0) ? -CAKE_INF : _tmem_load_3[6]); + _tmem_load_3[7] = (((frag_mask_word_2 >> 7 & 1) != 0) ? -CAKE_INF : _tmem_load_3[7]); + _tmem_load_3[8] = (((frag_mask_word_2 >> 8 & 1) != 0) ? -CAKE_INF : _tmem_load_3[8]); + _tmem_load_3[9] = (((frag_mask_word_2 >> 9 & 1) != 0) ? -CAKE_INF : _tmem_load_3[9]); + _tmem_load_3[10] = (((frag_mask_word_2 >> 10 & 1) != 0) ? -CAKE_INF : _tmem_load_3[10]); + _tmem_load_3[11] = (((frag_mask_word_2 >> 11 & 1) != 0) ? -CAKE_INF : _tmem_load_3[11]); + _tmem_load_3[12] = (((frag_mask_word_2 >> 12 & 1) != 0) ? -CAKE_INF : _tmem_load_3[12]); + _tmem_load_3[13] = (((frag_mask_word_2 >> 13 & 1) != 0) ? -CAKE_INF : _tmem_load_3[13]); + _tmem_load_3[14] = (((frag_mask_word_2 >> 14 & 1) != 0) ? -CAKE_INF : _tmem_load_3[14]); + _tmem_load_3[15] = (((frag_mask_word_2 >> 15 & 1) != 0) ? -CAKE_INF : _tmem_load_3[15]); + _tmem_load_3[16] = (((frag_mask_word_2 >> 16 & 1) != 0) ? -CAKE_INF : _tmem_load_3[16]); + _tmem_load_3[17] = (((frag_mask_word_2 >> 17 & 1) != 0) ? -CAKE_INF : _tmem_load_3[17]); + _tmem_load_3[18] = (((frag_mask_word_2 >> 18 & 1) != 0) ? -CAKE_INF : _tmem_load_3[18]); + _tmem_load_3[19] = (((frag_mask_word_2 >> 19 & 1) != 0) ? -CAKE_INF : _tmem_load_3[19]); + _tmem_load_3[20] = (((frag_mask_word_2 >> 20 & 1) != 0) ? -CAKE_INF : _tmem_load_3[20]); + _tmem_load_3[21] = (((frag_mask_word_2 >> 21 & 1) != 0) ? -CAKE_INF : _tmem_load_3[21]); + _tmem_load_3[22] = (((frag_mask_word_2 >> 22 & 1) != 0) ? -CAKE_INF : _tmem_load_3[22]); + _tmem_load_3[23] = (((frag_mask_word_2 >> 23 & 1) != 0) ? -CAKE_INF : _tmem_load_3[23]); + _tmem_load_3[24] = (((frag_mask_word_2 >> 24 & 1) != 0) ? -CAKE_INF : _tmem_load_3[24]); + _tmem_load_3[25] = (((frag_mask_word_2 >> 25 & 1) != 0) ? -CAKE_INF : _tmem_load_3[25]); + _tmem_load_3[26] = (((frag_mask_word_2 >> 26 & 1) != 0) ? -CAKE_INF : _tmem_load_3[26]); + _tmem_load_3[27] = (((frag_mask_word_2 >> 27 & 1) != 0) ? -CAKE_INF : _tmem_load_3[27]); + _tmem_load_3[28] = (((frag_mask_word_2 >> 28 & 1) != 0) ? -CAKE_INF : _tmem_load_3[28]); + _tmem_load_3[29] = (((frag_mask_word_2 >> 29 & 1) != 0) ? -CAKE_INF : _tmem_load_3[29]); + _tmem_load_3[30] = (((frag_mask_word_2 >> 30 & 1) != 0) ? -CAKE_INF : _tmem_load_3[30]); + _tmem_load_3[31] = (((frag_mask_word_2 >> 31 & 1) != 0) ? -CAKE_INF : _tmem_load_3[31]); + } + float2 _reg_reduce_max2_5 = {-CAKE_INF, -CAKE_INF}; + _reg_reduce_max2_5.x = max_noftz(_reg_reduce_max2_5.x, max_noftz(_tmem_load_3[0], _tmem_load_3[1])); + _reg_reduce_max2_5.y = max_noftz(_reg_reduce_max2_5.y, max_noftz(_tmem_load_3[2], _tmem_load_3[3])); + _reg_reduce_max2_5.x = max_noftz(_reg_reduce_max2_5.x, max_noftz(_tmem_load_3[4], _tmem_load_3[5])); + _reg_reduce_max2_5.y = max_noftz(_reg_reduce_max2_5.y, max_noftz(_tmem_load_3[6], _tmem_load_3[7])); + _reg_reduce_max2_5.x = max_noftz(_reg_reduce_max2_5.x, max_noftz(_tmem_load_3[8], _tmem_load_3[9])); + _reg_reduce_max2_5.y = max_noftz(_reg_reduce_max2_5.y, max_noftz(_tmem_load_3[10], _tmem_load_3[11])); + _reg_reduce_max2_5.x = max_noftz(_reg_reduce_max2_5.x, max_noftz(_tmem_load_3[12], _tmem_load_3[13])); + _reg_reduce_max2_5.y = max_noftz(_reg_reduce_max2_5.y, max_noftz(_tmem_load_3[14], _tmem_load_3[15])); + _reg_reduce_max2_5.x = max_noftz(_reg_reduce_max2_5.x, max_noftz(_tmem_load_3[16], _tmem_load_3[17])); + _reg_reduce_max2_5.y = max_noftz(_reg_reduce_max2_5.y, max_noftz(_tmem_load_3[18], _tmem_load_3[19])); + _reg_reduce_max2_5.x = max_noftz(_reg_reduce_max2_5.x, max_noftz(_tmem_load_3[20], _tmem_load_3[21])); + _reg_reduce_max2_5.y = max_noftz(_reg_reduce_max2_5.y, max_noftz(_tmem_load_3[22], _tmem_load_3[23])); + _reg_reduce_max2_5.x = max_noftz(_reg_reduce_max2_5.x, max_noftz(_tmem_load_3[24], _tmem_load_3[25])); + _reg_reduce_max2_5.y = max_noftz(_reg_reduce_max2_5.y, max_noftz(_tmem_load_3[26], _tmem_load_3[27])); + _reg_reduce_max2_5.x = max_noftz(_reg_reduce_max2_5.x, max_noftz(_tmem_load_3[28], _tmem_load_3[29])); + _reg_reduce_max2_5.y = max_noftz(_reg_reduce_max2_5.y, max_noftz(_tmem_load_3[30], _tmem_load_3[31])); + float _tmem_load_3_max = row_max_reduce(_reg_reduce_max2_5); + float _max_5 = max_noftz(new_max, _tmem_load_3_max); + new_max = _max_5; + int frag_valid_3 = valid_sparse_cols - 96; + uint32_t _slice_lo_mask_4; + { + int _lim_6 = frag_valid_3; + if (_lim_6 <= 0) { _slice_lo_mask_4 = 0u; } + else if (_lim_6 >= 32) { _slice_lo_mask_4 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_4) : "r"(_lim_6)); + } + } + if (!(_slice_lo_mask_4 & (1u << 0))) _tmem_load_4[0] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 1))) _tmem_load_4[1] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 2))) _tmem_load_4[2] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 3))) _tmem_load_4[3] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 4))) _tmem_load_4[4] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 5))) _tmem_load_4[5] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 6))) _tmem_load_4[6] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 7))) _tmem_load_4[7] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 8))) _tmem_load_4[8] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 9))) _tmem_load_4[9] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 10))) _tmem_load_4[10] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 11))) _tmem_load_4[11] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 12))) _tmem_load_4[12] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 13))) _tmem_load_4[13] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 14))) _tmem_load_4[14] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 15))) _tmem_load_4[15] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 16))) _tmem_load_4[16] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 17))) _tmem_load_4[17] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 18))) _tmem_load_4[18] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 19))) _tmem_load_4[19] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 20))) _tmem_load_4[20] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 21))) _tmem_load_4[21] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 22))) _tmem_load_4[22] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 23))) _tmem_load_4[23] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 24))) _tmem_load_4[24] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 25))) _tmem_load_4[25] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 26))) _tmem_load_4[26] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 27))) _tmem_load_4[27] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 28))) _tmem_load_4[28] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 29))) _tmem_load_4[29] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 30))) _tmem_load_4[30] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 31))) _tmem_load_4[31] = -CAKE_INF; + if (any_invalid != 0) { + unsigned int frag_mask_word_3 = invalid_cols[3]; + _tmem_load_4[0] = (((frag_mask_word_3 & 1) != 0) ? -CAKE_INF : _tmem_load_4[0]); + _tmem_load_4[1] = (((frag_mask_word_3 >> 1 & 1) != 0) ? -CAKE_INF : _tmem_load_4[1]); + _tmem_load_4[2] = (((frag_mask_word_3 >> 2 & 1) != 0) ? -CAKE_INF : _tmem_load_4[2]); + _tmem_load_4[3] = (((frag_mask_word_3 >> 3 & 1) != 0) ? -CAKE_INF : _tmem_load_4[3]); + _tmem_load_4[4] = (((frag_mask_word_3 >> 4 & 1) != 0) ? -CAKE_INF : _tmem_load_4[4]); + _tmem_load_4[5] = (((frag_mask_word_3 >> 5 & 1) != 0) ? -CAKE_INF : _tmem_load_4[5]); + _tmem_load_4[6] = (((frag_mask_word_3 >> 6 & 1) != 0) ? -CAKE_INF : _tmem_load_4[6]); + _tmem_load_4[7] = (((frag_mask_word_3 >> 7 & 1) != 0) ? -CAKE_INF : _tmem_load_4[7]); + _tmem_load_4[8] = (((frag_mask_word_3 >> 8 & 1) != 0) ? -CAKE_INF : _tmem_load_4[8]); + _tmem_load_4[9] = (((frag_mask_word_3 >> 9 & 1) != 0) ? -CAKE_INF : _tmem_load_4[9]); + _tmem_load_4[10] = (((frag_mask_word_3 >> 10 & 1) != 0) ? -CAKE_INF : _tmem_load_4[10]); + _tmem_load_4[11] = (((frag_mask_word_3 >> 11 & 1) != 0) ? -CAKE_INF : _tmem_load_4[11]); + _tmem_load_4[12] = (((frag_mask_word_3 >> 12 & 1) != 0) ? -CAKE_INF : _tmem_load_4[12]); + _tmem_load_4[13] = (((frag_mask_word_3 >> 13 & 1) != 0) ? -CAKE_INF : _tmem_load_4[13]); + _tmem_load_4[14] = (((frag_mask_word_3 >> 14 & 1) != 0) ? -CAKE_INF : _tmem_load_4[14]); + _tmem_load_4[15] = (((frag_mask_word_3 >> 15 & 1) != 0) ? -CAKE_INF : _tmem_load_4[15]); + _tmem_load_4[16] = (((frag_mask_word_3 >> 16 & 1) != 0) ? -CAKE_INF : _tmem_load_4[16]); + _tmem_load_4[17] = (((frag_mask_word_3 >> 17 & 1) != 0) ? -CAKE_INF : _tmem_load_4[17]); + _tmem_load_4[18] = (((frag_mask_word_3 >> 18 & 1) != 0) ? -CAKE_INF : _tmem_load_4[18]); + _tmem_load_4[19] = (((frag_mask_word_3 >> 19 & 1) != 0) ? -CAKE_INF : _tmem_load_4[19]); + _tmem_load_4[20] = (((frag_mask_word_3 >> 20 & 1) != 0) ? -CAKE_INF : _tmem_load_4[20]); + _tmem_load_4[21] = (((frag_mask_word_3 >> 21 & 1) != 0) ? -CAKE_INF : _tmem_load_4[21]); + _tmem_load_4[22] = (((frag_mask_word_3 >> 22 & 1) != 0) ? -CAKE_INF : _tmem_load_4[22]); + _tmem_load_4[23] = (((frag_mask_word_3 >> 23 & 1) != 0) ? -CAKE_INF : _tmem_load_4[23]); + _tmem_load_4[24] = (((frag_mask_word_3 >> 24 & 1) != 0) ? -CAKE_INF : _tmem_load_4[24]); + _tmem_load_4[25] = (((frag_mask_word_3 >> 25 & 1) != 0) ? -CAKE_INF : _tmem_load_4[25]); + _tmem_load_4[26] = (((frag_mask_word_3 >> 26 & 1) != 0) ? -CAKE_INF : _tmem_load_4[26]); + _tmem_load_4[27] = (((frag_mask_word_3 >> 27 & 1) != 0) ? -CAKE_INF : _tmem_load_4[27]); + _tmem_load_4[28] = (((frag_mask_word_3 >> 28 & 1) != 0) ? -CAKE_INF : _tmem_load_4[28]); + _tmem_load_4[29] = (((frag_mask_word_3 >> 29 & 1) != 0) ? -CAKE_INF : _tmem_load_4[29]); + _tmem_load_4[30] = (((frag_mask_word_3 >> 30 & 1) != 0) ? -CAKE_INF : _tmem_load_4[30]); + _tmem_load_4[31] = (((frag_mask_word_3 >> 31 & 1) != 0) ? -CAKE_INF : _tmem_load_4[31]); + } + float2 _reg_reduce_max2_7 = {-CAKE_INF, -CAKE_INF}; + _reg_reduce_max2_7.x = max_noftz(_reg_reduce_max2_7.x, max_noftz(_tmem_load_4[0], _tmem_load_4[1])); + _reg_reduce_max2_7.y = max_noftz(_reg_reduce_max2_7.y, max_noftz(_tmem_load_4[2], _tmem_load_4[3])); + _reg_reduce_max2_7.x = max_noftz(_reg_reduce_max2_7.x, max_noftz(_tmem_load_4[4], _tmem_load_4[5])); + _reg_reduce_max2_7.y = max_noftz(_reg_reduce_max2_7.y, max_noftz(_tmem_load_4[6], _tmem_load_4[7])); + _reg_reduce_max2_7.x = max_noftz(_reg_reduce_max2_7.x, max_noftz(_tmem_load_4[8], _tmem_load_4[9])); + _reg_reduce_max2_7.y = max_noftz(_reg_reduce_max2_7.y, max_noftz(_tmem_load_4[10], _tmem_load_4[11])); + _reg_reduce_max2_7.x = max_noftz(_reg_reduce_max2_7.x, max_noftz(_tmem_load_4[12], _tmem_load_4[13])); + _reg_reduce_max2_7.y = max_noftz(_reg_reduce_max2_7.y, max_noftz(_tmem_load_4[14], _tmem_load_4[15])); + _reg_reduce_max2_7.x = max_noftz(_reg_reduce_max2_7.x, max_noftz(_tmem_load_4[16], _tmem_load_4[17])); + _reg_reduce_max2_7.y = max_noftz(_reg_reduce_max2_7.y, max_noftz(_tmem_load_4[18], _tmem_load_4[19])); + _reg_reduce_max2_7.x = max_noftz(_reg_reduce_max2_7.x, max_noftz(_tmem_load_4[20], _tmem_load_4[21])); + _reg_reduce_max2_7.y = max_noftz(_reg_reduce_max2_7.y, max_noftz(_tmem_load_4[22], _tmem_load_4[23])); + _reg_reduce_max2_7.x = max_noftz(_reg_reduce_max2_7.x, max_noftz(_tmem_load_4[24], _tmem_load_4[25])); + _reg_reduce_max2_7.y = max_noftz(_reg_reduce_max2_7.y, max_noftz(_tmem_load_4[26], _tmem_load_4[27])); + _reg_reduce_max2_7.x = max_noftz(_reg_reduce_max2_7.x, max_noftz(_tmem_load_4[28], _tmem_load_4[29])); + _reg_reduce_max2_7.y = max_noftz(_reg_reduce_max2_7.y, max_noftz(_tmem_load_4[30], _tmem_load_4[31])); + float _tmem_load_4_max = row_max_reduce(_reg_reduce_max2_7); + float _max_6 = max_noftz(new_max, _tmem_load_4_max); + new_max = _max_6; + } + float _fma_0 = __fmaf_rn(row_max_val, softmax_scale_log2, (-new_max) * softmax_scale_log2); + float delta = _fma_0; + float _exp2_0 = approx_exp2(delta); + float exp_delta = _exp2_0; + float acc_scale = ((row_max_val > -CAKE_INF) ? exp_delta : 1.0f); + smem_stats_max[phase * 128 + stats_row] = acc_scale; + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(stats_addr + (phase) * 8); + row_max_val = new_max; + float safe_max = ((new_max == -CAKE_INF) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + float block_sum = 0.0f; + { + int p_off_all = ((phase != 0) ? 224 : 96); + int p_base_all = taddr + (unsigned int)p_off_all + (unsigned int)(tmem_row_base << 16); + const float2 _fma_b2_8 = {softmax_scale_log2, softmax_scale_log2}; + const float2 _fma_c2_9 = {-max_scaled, -max_scaled}; + float2 _fma_pair_10 = fma_f32x2(make_float2(_tmem_load_0[0], _tmem_load_0[1]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[0] = _fma_pair_10.x; + _tmem_load_0[1] = _fma_pair_10.y; + float2 _fma_pair_11 = fma_f32x2(make_float2(_tmem_load_0[2], _tmem_load_0[3]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[2] = _fma_pair_11.x; + _tmem_load_0[3] = _fma_pair_11.y; + float2 _fma_pair_12 = fma_f32x2(make_float2(_tmem_load_0[4], _tmem_load_0[5]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[4] = _fma_pair_12.x; + _tmem_load_0[5] = _fma_pair_12.y; + float2 _fma_pair_13 = fma_f32x2(make_float2(_tmem_load_0[6], _tmem_load_0[7]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[6] = _fma_pair_13.x; + _tmem_load_0[7] = _fma_pair_13.y; + float2 _fma_pair_14 = fma_f32x2(make_float2(_tmem_load_0[8], _tmem_load_0[9]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[8] = _fma_pair_14.x; + _tmem_load_0[9] = _fma_pair_14.y; + float2 _fma_pair_15 = fma_f32x2(make_float2(_tmem_load_0[10], _tmem_load_0[11]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[10] = _fma_pair_15.x; + _tmem_load_0[11] = _fma_pair_15.y; + float2 _fma_pair_16 = fma_f32x2(make_float2(_tmem_load_0[12], _tmem_load_0[13]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[12] = _fma_pair_16.x; + _tmem_load_0[13] = _fma_pair_16.y; + float2 _fma_pair_17 = fma_f32x2(make_float2(_tmem_load_0[14], _tmem_load_0[15]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[14] = _fma_pair_17.x; + _tmem_load_0[15] = _fma_pair_17.y; + float2 _fma_pair_18 = fma_f32x2(make_float2(_tmem_load_0[16], _tmem_load_0[17]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[16] = _fma_pair_18.x; + _tmem_load_0[17] = _fma_pair_18.y; + float2 _fma_pair_19 = fma_f32x2(make_float2(_tmem_load_0[18], _tmem_load_0[19]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[18] = _fma_pair_19.x; + _tmem_load_0[19] = _fma_pair_19.y; + float2 _fma_pair_20 = fma_f32x2(make_float2(_tmem_load_0[20], _tmem_load_0[21]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[20] = _fma_pair_20.x; + _tmem_load_0[21] = _fma_pair_20.y; + float2 _fma_pair_21 = fma_f32x2(make_float2(_tmem_load_0[22], _tmem_load_0[23]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[22] = _fma_pair_21.x; + _tmem_load_0[23] = _fma_pair_21.y; + float2 _fma_pair_22 = fma_f32x2(make_float2(_tmem_load_0[24], _tmem_load_0[25]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[24] = _fma_pair_22.x; + _tmem_load_0[25] = _fma_pair_22.y; + float2 _fma_pair_23 = fma_f32x2(make_float2(_tmem_load_0[26], _tmem_load_0[27]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[26] = _fma_pair_23.x; + _tmem_load_0[27] = _fma_pair_23.y; + float2 _fma_pair_24 = fma_f32x2(make_float2(_tmem_load_0[28], _tmem_load_0[29]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[28] = _fma_pair_24.x; + _tmem_load_0[29] = _fma_pair_24.y; + float2 _fma_pair_25 = fma_f32x2(make_float2(_tmem_load_0[30], _tmem_load_0[31]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[30] = _fma_pair_25.x; + _tmem_load_0[31] = _fma_pair_25.y; + #pragma unroll + for (int _le = 0; _le < 32; _le++) { + _tmem_load_0[_le] = approx_exp2(_tmem_load_0[_le]); + } + float2 _reg_reduce_sum2_26 = make_float2(0.0f, 0.0f); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[0], _tmem_load_0[1])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[2], _tmem_load_0[3])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[4], _tmem_load_0[5])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[6], _tmem_load_0[7])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[8], _tmem_load_0[9])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[10], _tmem_load_0[11])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[12], _tmem_load_0[13])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[14], _tmem_load_0[15])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[16], _tmem_load_0[17])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[18], _tmem_load_0[19])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[20], _tmem_load_0[21])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[22], _tmem_load_0[23])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[24], _tmem_load_0[25])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[26], _tmem_load_0[27])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[28], _tmem_load_0[29])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[30], _tmem_load_0[31])); + float _tmem_load_0_sum = _reg_reduce_sum2_26.x + _reg_reduce_sum2_26.y; + block_sum = block_sum + _tmem_load_0_sum; + { + uint32_t _pv_packed[8]; + #pragma unroll + for (int _j = 0; _j < 8; _j++) { + uint32_t _pk; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_pk) : "f"(_tmem_load_0[0 + _j * 4]), "f"(_tmem_load_0[0 + _j * 4 + 1]), + "f"(_tmem_load_0[0 + _j * 4 + 2]), "f"(_tmem_load_0[0 + _j * 4 + 3])); + _pv_packed[_j] = _pk; + } + tmem_st_x8_u32(p_base_all, _pv_packed); + } + const float2 _fma_b2_27 = {softmax_scale_log2, softmax_scale_log2}; + const float2 _fma_c2_28 = {-max_scaled, -max_scaled}; + float2 _fma_pair_29 = fma_f32x2(make_float2(_tmem_load_2[0], _tmem_load_2[1]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[0] = _fma_pair_29.x; + _tmem_load_2[1] = _fma_pair_29.y; + float2 _fma_pair_30 = fma_f32x2(make_float2(_tmem_load_2[2], _tmem_load_2[3]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[2] = _fma_pair_30.x; + _tmem_load_2[3] = _fma_pair_30.y; + float2 _fma_pair_31 = fma_f32x2(make_float2(_tmem_load_2[4], _tmem_load_2[5]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[4] = _fma_pair_31.x; + _tmem_load_2[5] = _fma_pair_31.y; + float2 _fma_pair_32 = fma_f32x2(make_float2(_tmem_load_2[6], _tmem_load_2[7]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[6] = _fma_pair_32.x; + _tmem_load_2[7] = _fma_pair_32.y; + float2 _fma_pair_33 = fma_f32x2(make_float2(_tmem_load_2[8], _tmem_load_2[9]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[8] = _fma_pair_33.x; + _tmem_load_2[9] = _fma_pair_33.y; + float2 _fma_pair_34 = fma_f32x2(make_float2(_tmem_load_2[10], _tmem_load_2[11]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[10] = _fma_pair_34.x; + _tmem_load_2[11] = _fma_pair_34.y; + float2 _fma_pair_35 = fma_f32x2(make_float2(_tmem_load_2[12], _tmem_load_2[13]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[12] = _fma_pair_35.x; + _tmem_load_2[13] = _fma_pair_35.y; + float2 _fma_pair_36 = fma_f32x2(make_float2(_tmem_load_2[14], _tmem_load_2[15]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[14] = _fma_pair_36.x; + _tmem_load_2[15] = _fma_pair_36.y; + float2 _fma_pair_37 = fma_f32x2(make_float2(_tmem_load_2[16], _tmem_load_2[17]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[16] = _fma_pair_37.x; + _tmem_load_2[17] = _fma_pair_37.y; + float2 _fma_pair_38 = fma_f32x2(make_float2(_tmem_load_2[18], _tmem_load_2[19]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[18] = _fma_pair_38.x; + _tmem_load_2[19] = _fma_pair_38.y; + float2 _fma_pair_39 = fma_f32x2(make_float2(_tmem_load_2[20], _tmem_load_2[21]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[20] = _fma_pair_39.x; + _tmem_load_2[21] = _fma_pair_39.y; + float2 _fma_pair_40 = fma_f32x2(make_float2(_tmem_load_2[22], _tmem_load_2[23]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[22] = _fma_pair_40.x; + _tmem_load_2[23] = _fma_pair_40.y; + float2 _fma_pair_41 = fma_f32x2(make_float2(_tmem_load_2[24], _tmem_load_2[25]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[24] = _fma_pair_41.x; + _tmem_load_2[25] = _fma_pair_41.y; + float2 _fma_pair_42 = fma_f32x2(make_float2(_tmem_load_2[26], _tmem_load_2[27]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[26] = _fma_pair_42.x; + _tmem_load_2[27] = _fma_pair_42.y; + float2 _fma_pair_43 = fma_f32x2(make_float2(_tmem_load_2[28], _tmem_load_2[29]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[28] = _fma_pair_43.x; + _tmem_load_2[29] = _fma_pair_43.y; + float2 _fma_pair_44 = fma_f32x2(make_float2(_tmem_load_2[30], _tmem_load_2[31]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[30] = _fma_pair_44.x; + _tmem_load_2[31] = _fma_pair_44.y; + #pragma unroll + for (int _le = 0; _le < 32; _le++) { + _tmem_load_2[_le] = approx_exp2(_tmem_load_2[_le]); + } + float2 _reg_reduce_sum2_45 = make_float2(0.0f, 0.0f); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[0], _tmem_load_2[1])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[2], _tmem_load_2[3])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[4], _tmem_load_2[5])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[6], _tmem_load_2[7])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[8], _tmem_load_2[9])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[10], _tmem_load_2[11])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[12], _tmem_load_2[13])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[14], _tmem_load_2[15])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[16], _tmem_load_2[17])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[18], _tmem_load_2[19])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[20], _tmem_load_2[21])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[22], _tmem_load_2[23])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[24], _tmem_load_2[25])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[26], _tmem_load_2[27])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[28], _tmem_load_2[29])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[30], _tmem_load_2[31])); + float _tmem_load_2_sum = _reg_reduce_sum2_45.x + _reg_reduce_sum2_45.y; + block_sum = block_sum + _tmem_load_2_sum; + { + uint32_t _pv_packed[8]; + #pragma unroll + for (int _j = 0; _j < 8; _j++) { + uint32_t _pk; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_pk) : "f"(_tmem_load_2[0 + _j * 4]), "f"(_tmem_load_2[0 + _j * 4 + 1]), + "f"(_tmem_load_2[0 + _j * 4 + 2]), "f"(_tmem_load_2[0 + _j * 4 + 3])); + _pv_packed[_j] = _pk; + } + tmem_st_x8_u32(p_base_all + 8, _pv_packed); + } + const float2 _fma_b2_46 = {softmax_scale_log2, softmax_scale_log2}; + const float2 _fma_c2_47 = {-max_scaled, -max_scaled}; + float2 _fma_pair_48 = fma_f32x2(make_float2(_tmem_load_3[0], _tmem_load_3[1]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[0] = _fma_pair_48.x; + _tmem_load_3[1] = _fma_pair_48.y; + float2 _fma_pair_49 = fma_f32x2(make_float2(_tmem_load_3[2], _tmem_load_3[3]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[2] = _fma_pair_49.x; + _tmem_load_3[3] = _fma_pair_49.y; + float2 _fma_pair_50 = fma_f32x2(make_float2(_tmem_load_3[4], _tmem_load_3[5]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[4] = _fma_pair_50.x; + _tmem_load_3[5] = _fma_pair_50.y; + float2 _fma_pair_51 = fma_f32x2(make_float2(_tmem_load_3[6], _tmem_load_3[7]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[6] = _fma_pair_51.x; + _tmem_load_3[7] = _fma_pair_51.y; + float2 _fma_pair_52 = fma_f32x2(make_float2(_tmem_load_3[8], _tmem_load_3[9]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[8] = _fma_pair_52.x; + _tmem_load_3[9] = _fma_pair_52.y; + float2 _fma_pair_53 = fma_f32x2(make_float2(_tmem_load_3[10], _tmem_load_3[11]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[10] = _fma_pair_53.x; + _tmem_load_3[11] = _fma_pair_53.y; + float2 _fma_pair_54 = fma_f32x2(make_float2(_tmem_load_3[12], _tmem_load_3[13]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[12] = _fma_pair_54.x; + _tmem_load_3[13] = _fma_pair_54.y; + float2 _fma_pair_55 = fma_f32x2(make_float2(_tmem_load_3[14], _tmem_load_3[15]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[14] = _fma_pair_55.x; + _tmem_load_3[15] = _fma_pair_55.y; + float2 _fma_pair_56 = fma_f32x2(make_float2(_tmem_load_3[16], _tmem_load_3[17]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[16] = _fma_pair_56.x; + _tmem_load_3[17] = _fma_pair_56.y; + float2 _fma_pair_57 = fma_f32x2(make_float2(_tmem_load_3[18], _tmem_load_3[19]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[18] = _fma_pair_57.x; + _tmem_load_3[19] = _fma_pair_57.y; + float2 _fma_pair_58 = fma_f32x2(make_float2(_tmem_load_3[20], _tmem_load_3[21]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[20] = _fma_pair_58.x; + _tmem_load_3[21] = _fma_pair_58.y; + float2 _fma_pair_59 = fma_f32x2(make_float2(_tmem_load_3[22], _tmem_load_3[23]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[22] = _fma_pair_59.x; + _tmem_load_3[23] = _fma_pair_59.y; + float2 _fma_pair_60 = fma_f32x2(make_float2(_tmem_load_3[24], _tmem_load_3[25]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[24] = _fma_pair_60.x; + _tmem_load_3[25] = _fma_pair_60.y; + float2 _fma_pair_61 = fma_f32x2(make_float2(_tmem_load_3[26], _tmem_load_3[27]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[26] = _fma_pair_61.x; + _tmem_load_3[27] = _fma_pair_61.y; + float2 _fma_pair_62 = fma_f32x2(make_float2(_tmem_load_3[28], _tmem_load_3[29]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[28] = _fma_pair_62.x; + _tmem_load_3[29] = _fma_pair_62.y; + float2 _fma_pair_63 = fma_f32x2(make_float2(_tmem_load_3[30], _tmem_load_3[31]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[30] = _fma_pair_63.x; + _tmem_load_3[31] = _fma_pair_63.y; + #pragma unroll + for (int _le = 0; _le < 32; _le++) { + _tmem_load_3[_le] = approx_exp2(_tmem_load_3[_le]); + } + float2 _reg_reduce_sum2_64 = make_float2(0.0f, 0.0f); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[0], _tmem_load_3[1])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[2], _tmem_load_3[3])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[4], _tmem_load_3[5])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[6], _tmem_load_3[7])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[8], _tmem_load_3[9])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[10], _tmem_load_3[11])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[12], _tmem_load_3[13])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[14], _tmem_load_3[15])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[16], _tmem_load_3[17])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[18], _tmem_load_3[19])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[20], _tmem_load_3[21])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[22], _tmem_load_3[23])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[24], _tmem_load_3[25])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[26], _tmem_load_3[27])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[28], _tmem_load_3[29])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[30], _tmem_load_3[31])); + float _tmem_load_3_sum = _reg_reduce_sum2_64.x + _reg_reduce_sum2_64.y; + block_sum = block_sum + _tmem_load_3_sum; + { + uint32_t _pv_packed[8]; + #pragma unroll + for (int _j = 0; _j < 8; _j++) { + uint32_t _pk; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_pk) : "f"(_tmem_load_3[0 + _j * 4]), "f"(_tmem_load_3[0 + _j * 4 + 1]), + "f"(_tmem_load_3[0 + _j * 4 + 2]), "f"(_tmem_load_3[0 + _j * 4 + 3])); + _pv_packed[_j] = _pk; + } + tmem_st_x8_u32(p_base_all + 16, _pv_packed); + } + const float2 _fma_b2_65 = {softmax_scale_log2, softmax_scale_log2}; + const float2 _fma_c2_66 = {-max_scaled, -max_scaled}; + float2 _fma_pair_67 = fma_f32x2(make_float2(_tmem_load_4[0], _tmem_load_4[1]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[0] = _fma_pair_67.x; + _tmem_load_4[1] = _fma_pair_67.y; + float2 _fma_pair_68 = fma_f32x2(make_float2(_tmem_load_4[2], _tmem_load_4[3]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[2] = _fma_pair_68.x; + _tmem_load_4[3] = _fma_pair_68.y; + float2 _fma_pair_69 = fma_f32x2(make_float2(_tmem_load_4[4], _tmem_load_4[5]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[4] = _fma_pair_69.x; + _tmem_load_4[5] = _fma_pair_69.y; + float2 _fma_pair_70 = fma_f32x2(make_float2(_tmem_load_4[6], _tmem_load_4[7]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[6] = _fma_pair_70.x; + _tmem_load_4[7] = _fma_pair_70.y; + float2 _fma_pair_71 = fma_f32x2(make_float2(_tmem_load_4[8], _tmem_load_4[9]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[8] = _fma_pair_71.x; + _tmem_load_4[9] = _fma_pair_71.y; + float2 _fma_pair_72 = fma_f32x2(make_float2(_tmem_load_4[10], _tmem_load_4[11]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[10] = _fma_pair_72.x; + _tmem_load_4[11] = _fma_pair_72.y; + float2 _fma_pair_73 = fma_f32x2(make_float2(_tmem_load_4[12], _tmem_load_4[13]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[12] = _fma_pair_73.x; + _tmem_load_4[13] = _fma_pair_73.y; + float2 _fma_pair_74 = fma_f32x2(make_float2(_tmem_load_4[14], _tmem_load_4[15]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[14] = _fma_pair_74.x; + _tmem_load_4[15] = _fma_pair_74.y; + float2 _fma_pair_75 = fma_f32x2(make_float2(_tmem_load_4[16], _tmem_load_4[17]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[16] = _fma_pair_75.x; + _tmem_load_4[17] = _fma_pair_75.y; + float2 _fma_pair_76 = fma_f32x2(make_float2(_tmem_load_4[18], _tmem_load_4[19]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[18] = _fma_pair_76.x; + _tmem_load_4[19] = _fma_pair_76.y; + float2 _fma_pair_77 = fma_f32x2(make_float2(_tmem_load_4[20], _tmem_load_4[21]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[20] = _fma_pair_77.x; + _tmem_load_4[21] = _fma_pair_77.y; + float2 _fma_pair_78 = fma_f32x2(make_float2(_tmem_load_4[22], _tmem_load_4[23]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[22] = _fma_pair_78.x; + _tmem_load_4[23] = _fma_pair_78.y; + float2 _fma_pair_79 = fma_f32x2(make_float2(_tmem_load_4[24], _tmem_load_4[25]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[24] = _fma_pair_79.x; + _tmem_load_4[25] = _fma_pair_79.y; + float2 _fma_pair_80 = fma_f32x2(make_float2(_tmem_load_4[26], _tmem_load_4[27]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[26] = _fma_pair_80.x; + _tmem_load_4[27] = _fma_pair_80.y; + float2 _fma_pair_81 = fma_f32x2(make_float2(_tmem_load_4[28], _tmem_load_4[29]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[28] = _fma_pair_81.x; + _tmem_load_4[29] = _fma_pair_81.y; + float2 _fma_pair_82 = fma_f32x2(make_float2(_tmem_load_4[30], _tmem_load_4[31]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[30] = _fma_pair_82.x; + _tmem_load_4[31] = _fma_pair_82.y; + #pragma unroll + for (int _le = 0; _le < 32; _le++) { + _tmem_load_4[_le] = approx_exp2(_tmem_load_4[_le]); + } + float2 _reg_reduce_sum2_83 = make_float2(0.0f, 0.0f); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[0], _tmem_load_4[1])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[2], _tmem_load_4[3])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[4], _tmem_load_4[5])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[6], _tmem_load_4[7])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[8], _tmem_load_4[9])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[10], _tmem_load_4[11])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[12], _tmem_load_4[13])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[14], _tmem_load_4[15])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[16], _tmem_load_4[17])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[18], _tmem_load_4[19])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[20], _tmem_load_4[21])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[22], _tmem_load_4[23])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[24], _tmem_load_4[25])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[26], _tmem_load_4[27])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[28], _tmem_load_4[29])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[30], _tmem_load_4[31])); + float _tmem_load_4_sum = _reg_reduce_sum2_83.x + _reg_reduce_sum2_83.y; + block_sum = block_sum + _tmem_load_4_sum; + { + uint32_t _pv_packed[8]; + #pragma unroll + for (int _j = 0; _j < 8; _j++) { + uint32_t _pk; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_pk) : "f"(_tmem_load_4[0 + _j * 4]), "f"(_tmem_load_4[0 + _j * 4 + 1]), + "f"(_tmem_load_4[0 + _j * 4 + 2]), "f"(_tmem_load_4[0 + _j * 4 + 3])); + _pv_packed[_j] = _pk; + } + tmem_st_x8_u32(p_base_all + 24, _pv_packed); + } + } + { + float _fma_1 = __fmaf_rn(row_sum_val, acc_scale, block_sum); + row_sum_val = _fma_1; + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(p_full_addr + (phase) * 8); + } + smem_stats_sum[stats_row] = row_sum_val; + smem_stats_final_max[stats_row] = row_max_val; + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(sum_ready_addr); + softmax_tile_cursor = softmax_tile_cursor + num_kv_tiles_1; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_dealloc_addr); + } + } + // ---- Role: correction_wg ---- + if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 184;"); + { // correction_wg_main + float softmax_scale_log2_1 = bmm1_scale[0] * 1.4426950408889634f; + float output_scale = bmm2_scale[0]; + const int wg_dummy_inc_1 = 0; + int all_num_kv_tiles_2 = (sparse_topk + 128 - 1) / 128; + const int tmem_row_base_1 = ((0) ? warp % 2 * 32 : warp % 4 * 32); + const int n_half_1 = ((0) ? (int)(warp % 4 / 2) : 0); + const int my_row_1 = tmem_row_base_1 + lane; + const int stats_row_1 = n_half_1 * 64 + my_row_1; + const int corr_row = tmem_row_base_1 << 16; + int correction_tile_cursor = 0; + unsigned int _phase_q_full_0_2 = 0; + unsigned int _phase_pv_done_0 = 0; + unsigned int _phase_o_done_0 = 0; + unsigned int _phase_sum_ready_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx_3 = bid; work_idx_3 < total_work_items; work_idx_3 += num_bids) { + int split_idx_2 = 0; + int query_idx_2 = work_idx_3 >> 1; + int v_chunk_1 = work_idx_3 & 1; + int tiles_per_split_2 = all_num_kv_tiles_2; + int first_tile_2 = split_idx_2 * tiles_per_split_2; + int num_kv_tiles_2 = tiles_per_split_2; + { + mbarrier_wait(q_full_addr, _phase_q_full_0_2); + _phase_q_full_0_2 ^= 1; + } + #pragma unroll 1 + for (int tile_2 = 0; tile_2 < num_kv_tiles_2; tile_2++) { + int pipeline_tile_1 = correction_tile_cursor + tile_2; + int phase_1 = pipeline_tile_1 & 1; + int stats_wait_phase = pipeline_tile_1 >> 1 & 1; + mbarrier_wait(stats_addr + (phase_1) * 8, stats_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + float acc_scale_1 = smem_stats_max[phase_1 * 128 + stats_row_1]; + if (tile_2 > 0) { + mbarrier_wait(pv_done_addr, _phase_pv_done_0); + _phase_pv_done_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + } + if (tile_2 > 0) { + int _vote_1 = __any_sync(0xFFFFFFFF, acc_scale_1 < 1.0f); + int any_rescale = _vote_1; + if (any_rescale != 0) { + #pragma unroll + for (int vs = 0; vs < 2; vs++) { + int o_base = taddr + 256 + (unsigned int)(vs * 128) + (unsigned int)corr_row; + #pragma unroll + for (int c = 0; c < 128; c += 16) { + float _tmem_load_5[16]; + tmem_ld_x16(&_tmem_load_5[0], o_base + c); + const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + mul_f32x2_inplace(&reinterpret_cast(_tmem_load_5)[_ls], _scale2_0); + tmem_st_x16_f32(o_base + c, _tmem_load_5); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + } + } + } + } + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(corr_done_addr + (phase_1) * 8); + } + mbarrier_wait(o_done_addr, _phase_o_done_0); + _phase_o_done_0 ^= 1; + mbarrier_wait(sum_ready_addr, _phase_sum_ready_0); + _phase_sum_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + float total_sum = smem_stats_sum[stats_row_1]; + float final_max = smem_stats_final_max[stats_row_1]; + float _rcp_0 = approx_rcp(total_sum); + float inv_sum = ((total_sum > 0.0f) ? _rcp_0 : 0.0f); + int head_idx = my_row_1; + int direct_o_offset = (query_idx_2 * num_heads + head_idx) * 512 + v_chunk_1 * 256; + int partial_o_offset = (query_idx_2 * num_heads + head_idx + split_idx_2) * 512 + ((0) ? v_chunk_1 : n_half_1) * 256; + int o_offset = direct_o_offset; + if ((0 & (int)(n_half_1 == 0) & (int)(head_idx < num_heads)) != 0) { + int stat_offset = query_idx_2 * num_heads + head_idx + split_idx_2; + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(total_sum)); + partial_lse[stat_offset] = ((total_sum > 0.0f) ? final_max * softmax_scale_log2_1 + _log2_0 : -CAKE_INF); + } + #pragma unroll + for (int vs_1 = 0; vs_1 < 2; vs_1++) { + int o_base_epi = taddr + 256 + (unsigned int)(vs_1 * 128) + (unsigned int)corr_row; + #pragma unroll + for (int c_1 = 0; c_1 < 128; c_1 += 32) { + float _tmem_load_6[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_6[0]), "=f"(_tmem_load_6[1]), "=f"(_tmem_load_6[2]), "=f"(_tmem_load_6[3]), "=f"(_tmem_load_6[4]), "=f"(_tmem_load_6[5]), "=f"(_tmem_load_6[6]), "=f"(_tmem_load_6[7]), "=f"(_tmem_load_6[8]), "=f"(_tmem_load_6[9]), "=f"(_tmem_load_6[10]), "=f"(_tmem_load_6[11]), "=f"(_tmem_load_6[12]), "=f"(_tmem_load_6[13]), "=f"(_tmem_load_6[14]), "=f"(_tmem_load_6[15]), "=f"(_tmem_load_6[16]), "=f"(_tmem_load_6[17]), "=f"(_tmem_load_6[18]), "=f"(_tmem_load_6[19]), "=f"(_tmem_load_6[20]), "=f"(_tmem_load_6[21]), "=f"(_tmem_load_6[22]), "=f"(_tmem_load_6[23]), "=f"(_tmem_load_6[24]), "=f"(_tmem_load_6[25]), "=f"(_tmem_load_6[26]), "=f"(_tmem_load_6[27]), "=f"(_tmem_load_6[28]), "=f"(_tmem_load_6[29]), "=f"(_tmem_load_6[30]), "=f"(_tmem_load_6[31]) + : "r"(o_base_epi + c_1)); + int gmem_base = o_offset + vs_1 * 128 + c_1; + if (head_idx < num_heads) { + #pragma unroll + for (int j = 0; j < 32; j += 16) { + { + const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; + #if __CUDA_ARCH__ >= 1000 + #pragma unroll + for (int _ps = 0; _ps < 8; _ps++) + mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_6[j])[_ps], _prescale2_1); + #else + #pragma unroll + for (int _ps = 0; _ps < 16; _ps++) + _tmem_load_6[j + _ps] *= inv_sum * output_scale; + #endif + { + __nv_bfloat162 _pk0 = __floats2bfloat162_rn(_tmem_load_6[j + 0], _tmem_load_6[j + 1]); + unsigned _pk_u0 = *reinterpret_cast(&_pk0); + __nv_bfloat162 _pk1 = __floats2bfloat162_rn(_tmem_load_6[j + 2], _tmem_load_6[j + 3]); + unsigned _pk_u1 = *reinterpret_cast(&_pk1); + __nv_bfloat162 _pk2 = __floats2bfloat162_rn(_tmem_load_6[j + 4], _tmem_load_6[j + 5]); + unsigned _pk_u2 = *reinterpret_cast(&_pk2); + __nv_bfloat162 _pk3 = __floats2bfloat162_rn(_tmem_load_6[j + 6], _tmem_load_6[j + 7]); + unsigned _pk_u3 = *reinterpret_cast(&_pk3); + __nv_bfloat162 _pk4 = __floats2bfloat162_rn(_tmem_load_6[j + 8], _tmem_load_6[j + 9]); + unsigned _pk_u4 = *reinterpret_cast(&_pk4); + __nv_bfloat162 _pk5 = __floats2bfloat162_rn(_tmem_load_6[j + 10], _tmem_load_6[j + 11]); + unsigned _pk_u5 = *reinterpret_cast(&_pk5); + __nv_bfloat162 _pk6 = __floats2bfloat162_rn(_tmem_load_6[j + 12], _tmem_load_6[j + 13]); + unsigned _pk_u6 = *reinterpret_cast(&_pk6); + __nv_bfloat162 _pk7 = __floats2bfloat162_rn(_tmem_load_6[j + 14], _tmem_load_6[j + 15]); + unsigned _pk_u7 = *reinterpret_cast(&_pk7); + asm volatile( + "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"((void*)(&((__nv_bfloat16*)(O + (gmem_base + j)))[0])), "r"(_pk_u0), "r"(_pk_u1), "r"(_pk_u2), "r"(_pk_u3), "r"(_pk_u4), "r"(_pk_u5), "r"(_pk_u6), "r"(_pk_u7) : "memory"); + } + } + } + } + } + } + correction_tile_cursor = correction_tile_cursor + num_kv_tiles_2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_dealloc_addr); + } + } + // ---- Role: mma_warp ---- + if (warp == 8) { + { // mma_warp_main + const int wg2_dummy_1 = 0; + int all_num_kv_tiles_3 = (sparse_topk + 128 - 1) / 128; + unsigned int mma_k_stage = 0; + unsigned int mma_v_stage = 0; + unsigned int mma_kv_stage = 0; + int mma_tile_cursor = 0; + unsigned int _phase_s_seeded_0 = 0; + unsigned int _phase_q_full_0_3 = 0; + unsigned int _phase_q_pair_ready_0 = 0; + unsigned int _phase_k_full = 0; + unsigned int _phase_kv_full = 0; + unsigned int _phase_kv_pair_ready = 0; + unsigned int _phase_v_full = 0; + #pragma unroll 1 + for (unsigned int work_idx_4 = bid; work_idx_4 < total_work_items; work_idx_4 += num_bids) { + int split_idx_3 = 0; + int tiles_per_split_3 = all_num_kv_tiles_3; + int first_tile_3 = split_idx_3 * tiles_per_split_3; + int num_kv_tiles_3 = tiles_per_split_3; + { + { + mbarrier_wait(s_seeded_addr, _phase_s_seeded_0); + _phase_s_seeded_0 ^= 1; + } + asm volatile("tcgen05.fence::after_thread_sync;"); + } + { + mbarrier_wait(q_full_addr, _phase_q_full_0_3); + _phase_q_full_0_3 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + } + int first_pv = 1; + #pragma unroll 1 + for (int tile_3 = 0; tile_3 < num_kv_tiles_3; tile_3++) { + int pipeline_tile_2 = mma_tile_cursor + tile_3; + int phase_2 = pipeline_tile_2 & 1; + int score_col = ((phase_2 != 0) ? 128 : 0); + { + #pragma unroll 1 + for (int n = 0; n < 4; n++) { + mbarrier_wait(kv_full_addr + (mma_kv_stage) * 8, _phase_kv_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + { + int _mma_a_lo_4 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (n) * 1024); + int _mma_b_lo_4 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 1024); + { + uint64_t _mma_ss_a_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_4); + uint64_t _mma_ss_b_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_4); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem_scratch + (score_col)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 136314896, ((n == 0) ? 0 : 1)); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem_scratch + (score_col)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 136314896, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem_scratch + (score_col)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 136314896, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem_scratch + (score_col)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 136314896, 1); + } + } + if (n == 3) { + elect_commit(s_full_addr + (phase_2) * 8); + } + elect_commit(kv_empty_addr + (mma_kv_stage) * 8); + } + mma_kv_stage += 1; + if (mma_kv_stage == 8) { mma_kv_stage = 0; _phase_kv_full ^= 1; _phase_kv_pair_ready ^= 1; } + } + } + if (tile_3 > 0) { + int prev_pipeline_tile = pipeline_tile_2 - 1; + int prev_phase = prev_pipeline_tile & 1; + int pv_wait_phase = prev_pipeline_tile >> 1 & 1; + mbarrier_wait(p_full_addr + (prev_phase) * 8, pv_wait_phase); + mbarrier_wait(corr_done_addr + (prev_phase) * 8, pv_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int p_col = ((prev_phase != 0) ? 224 : 96); + { + #pragma unroll + for (int vs_2 = 0; vs_2 < 2; vs_2++) { + mbarrier_wait(kv_full_addr + (mma_kv_stage) * 8, _phase_kv_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + int output_col = 256 + vs_2 * 128; + { + int _mma_b_lo_8 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (mma_kv_stage) * 1024); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 dhi, blo, ta, id;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 136380432;\n\t" + "mov.b32 ta, %2;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p0;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" + "}\n" + :: "r"((tmem_tmem_scratch + (output_col))), "r"(_mma_b_lo_8), "r"(tmem_tmem_scratch + p_col), "r"(((first_pv) ? 0 : 1))); + elect_commit(kv_empty_addr + (mma_kv_stage) * 8); + } + mma_kv_stage += 1; + if (mma_kv_stage == 8) { mma_kv_stage = 0; _phase_kv_full ^= 1; _phase_kv_pair_ready ^= 1; } + } + } + first_pv = 0; + elect_commit(pv_done_addr); + } + } + int last_pipeline_tile = mma_tile_cursor + num_kv_tiles_3 - 1; + int last_phase = last_pipeline_tile & 1; + int drain_wait_phase = last_pipeline_tile >> 1 & 1; + mbarrier_wait(p_full_addr + (last_phase) * 8, drain_wait_phase); + mbarrier_wait(corr_done_addr + (last_phase) * 8, drain_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int p_col_last = ((last_phase != 0) ? 224 : 96); + { + #pragma unroll + for (int vs_3 = 0; vs_3 < 2; vs_3++) { + mbarrier_wait(kv_full_addr + (mma_kv_stage) * 8, _phase_kv_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + int output_col_d = 256 + vs_3 * 128; + { + int _mma_b_lo_12 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (mma_kv_stage) * 1024); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 dhi, blo, ta, id;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 136380432;\n\t" + "mov.b32 ta, %2;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p0;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" + "}\n" + :: "r"((tmem_tmem_scratch + (output_col_d))), "r"(_mma_b_lo_12), "r"(tmem_tmem_scratch + p_col_last), "r"(((first_pv) ? 0 : 1))); + elect_commit(kv_empty_addr + (mma_kv_stage) * 8); + } + mma_kv_stage += 1; + if (mma_kv_stage == 8) { mma_kv_stage = 0; _phase_kv_full ^= 1; _phase_kv_pair_ready ^= 1; } + } + } + elect_commit(q_empty_addr); + elect_commit(o_done_addr); + mma_tile_cursor = mma_tile_cursor + num_kv_tiles_3; + } + mbarrier_arrive(tmem_dealloc_addr); + unsigned int _phase_tmem_dealloc_0 = 0; + mbarrier_wait(tmem_dealloc_addr, _phase_tmem_dealloc_0); + _phase_tmem_dealloc_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + { + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + } + } + // ---- Role: empty1 ---- + if (warp >= 14 && warp <= 15) { + { // empty1_main + const int wg2_dummy_2 = 0; + unsigned int _phase_q_full_0_4 = 0; + { + #pragma unroll 1 + for (unsigned int work_idx_5 = bid; work_idx_5 < total_work_items; work_idx_5 += num_bids) { + mbarrier_wait(q_full_addr, _phase_q_full_0_4); + _phase_q_full_0_4 ^= 1; + } + } + } + } + + // Cleanup +} + +} // extern "C" diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_983db4c490ceed3c455a_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_983db4c490ceed3c455a_binding.cu new file mode 100644 index 00000000000..b5c6878ec55 --- /dev/null +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_983db4c490ceed3c455a_binding.cu @@ -0,0 +1,461 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +// Generated by the Cake source exporter — do not edit. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_983db4c490ceed3c455a'. +#include +#include +#include +#include + +#include "tvm_ffi_utils.h" + +#include +#include +#include +#include +#include +#include +#include + +extern "C" __global__ void kernel_cake_dsv4_983db4c490ceed3c455a(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items); + + +namespace cake_host_shim_8e7f30cc5fb7222a { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t CakeDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +// 4D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + int64_t d3 = t.size(t.ndim() - 3); + TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + TVM_FFI_CHECK(d1 % 128 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 128; + uint64_t global_dim[4] = {(uint64_t)(128), (uint64_t)(d2), (uint64_t)((d1 / 128)), (uint64_t)(d3)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (128, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 128; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + (uint64_t)((carrier_stride_1 * 8) / 8), + (uint64_t)((carrier_stride_2 * 8) / 8), + }; + uint32_t box_dim[4] = {128u, 64u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_swa_kv' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_swa_kv' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; + TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (128, 1) exceeds resolved global dims for 'tmap_swa_kv'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + }; + uint32_t box_dim[2] = {128u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_compressed_kv' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_compressed_kv(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_compressed_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_compressed_kv' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_compressed_kv' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_compressed_kv"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_compressed_kv' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved a non-positive global dim"; + TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (128, 1) exceeds resolved global dims for 'tmap_compressed_kv'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + }; + uint32_t box_dim[2] = {128u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_kv') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_compressed_kv, TensorView arg_O, TensorView arg_partial_lse, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_num_query_tokens, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 1, 8, 1); + CheckContiguous(arg_tmap_q, "tmap_q"); + CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); + CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 1, 8, 1); + CheckCudaTensor(arg_tmap_compressed_kv, "tmap_compressed_kv"); + CheckDtype(arg_tmap_compressed_kv, "tmap_compressed_kv", 1, 8, 1); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_partial_lse, "partial_lse"); + CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); + CheckContiguous(arg_partial_lse, "partial_lse"); + CheckCudaTensor(arg_swa_indices, "swa_indices"); + CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); + CheckContiguous(arg_swa_indices, "swa_indices"); + CheckCudaTensor(arg_compressed_indices, "compressed_indices"); + CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); + CheckContiguous(arg_compressed_indices, "compressed_indices"); + CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); + CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); + CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); + CheckContiguous(arg_bmm1_scale, "bmm1_scale"); + CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); + CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); + CheckContiguous(arg_bmm2_scale, "bmm2_scale"); + TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) + << "scalar 'num_heads' value " << arg_num_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) + << "scalar 'swa_index_stride' value " << arg_swa_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) + << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) + << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_query_tokens >= -2147483648LL && arg_num_query_tokens <= 2147483647LL, ValueError) + << "scalar 'num_query_tokens' value " << arg_num_query_tokens + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) + << "scalar 'sparse_topk' value " << arg_sparse_topk + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) + << "scalar 'has_sinks' value " << arg_has_sinks + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); + CheckSameCudaDevice(arg_tmap_compressed_kv, arg_tmap_q, "tmap_compressed_kv", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_partial_lse, arg_tmap_q, "partial_lse", "tmap_q"); + CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); + CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); + CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); + CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") must be divisible by cluster dims (2, 1, 1)"; + + cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); + CUtensorMap p_tmap_compressed_kv = EncodeTma_tmap_compressed_kv(arg_tmap_compressed_kv); + __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); + float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); + int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); + int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); + int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); + float* p_sinks = static_cast(arg_sinks.data_ptr()); + float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); + float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); + int32_t v_num_heads = (int32_t)arg_num_heads; + int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; + int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; + int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; + int32_t v_num_query_tokens = (int32_t)arg_num_query_tokens; + int32_t v_sparse_topk = (int32_t)arg_sparse_topk; + int32_t v_has_sinks = (int32_t)arg_has_sinks; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_tmap_compressed_kv, &p_O, &p_partial_lse, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_num_query_tokens, &v_sparse_topk, &v_has_sinks, &v_total_work_items}; + + static std::mutex smem_mu; + static auto* smem_status_by_device = + new std::unordered_map(); + cudaError_t smem_status = cudaSuccess; + { + std::lock_guard lock(smem_mu); + auto it = smem_status_by_device->find(dev.device_id); + if (it == smem_status_by_device->end()) { + int current_device = -1; + smem_status = cudaGetDevice(¤t_device); + if (smem_status == cudaSuccess && current_device != dev.device_id) { + smem_status = cudaErrorInvalidDevice; + } + if (smem_status == cudaSuccess) { + smem_status = cudaFuncSetAttribute( + reinterpret_cast(kernel_cake_dsv4_983db4c490ceed3c455a), + cudaFuncAttributeMaxDynamicSharedMemorySize, + 223744); + } + smem_status_by_device->emplace(dev.device_id, smem_status); + } else { + smem_status = it->second; + } + } + TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) + << "cudaFuncSetAttribute for kernel_cake_dsv4_983db4c490ceed3c455a failed: " + << cudaGetErrorString(smem_status); + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(512u, 1u, 1u); + + // Extended direct-source launch for a linked CUDA kernel symbol. + cudaLaunchAttribute attrs[2]{}; + int n = 0; + attrs[n].id = cudaLaunchAttributeClusterDimension; + attrs[n].val.clusterDim.x = 2u; + attrs[n].val.clusterDim.y = 1u; + attrs[n].val.clusterDim.z = 1u; + ++n; + attrs[n].id = cudaLaunchAttributeClusterSchedulingPolicyPreference; + attrs[n].val.clusterSchedulingPolicyPreference = cudaClusterSchedulingPolicySpread; + ++n; + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes = 223744u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; + cudaError_t launch_status = cudaLaunchKernelExC( + &config, reinterpret_cast(kernel_cake_dsv4_983db4c490ceed3c455a), kargs); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernelExC for kernel_cake_dsv4_983db4c490ceed3c455a failed: " + << cudaGetErrorString(launch_status); +} + +} // namespace cake_host_shim_8e7f30cc5fb7222a + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_8e7f30cc5fb7222a::Run); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_983db4c490ceed3c455a_kernel.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_983db4c490ceed3c455a_kernel.cu new file mode 100644 index 00000000000..a01f2cc29ce --- /dev/null +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_983db4c490ceed3c455a_kernel.cu @@ -0,0 +1,2802 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +#define CAKE_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_SCRATCH_OFFSET 0 +#define NUM_K_PIPE_STAGES 3 +#define NUM_V_PIPE_STAGES 2 +#define NUM_KV_PIPE_STAGES 1 +#define NUM_INDEX_PIPE_STAGES 6 +#define SMEM_SMEM_Q_OFF 1024 +#define SMEM_SMEM_Q_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_STRIDE 8192 +#define SMEM_SMEM_KV_OFF 33792 +#define SMEM_SMEM_KV_STAGE_BYTES 16384 +#define SMEM_SMEM_KV_STRIDE 16384 +#define SMEM_SMEM_V_OFF 33792 +#define SMEM_SMEM_V_STAGE_BYTES 16384 +#define SMEM_SMEM_V_STRIDE 16384 +#define SMEM_SMEM_Q_FULL_OFF 1024 +#define SMEM_SMEM_Q_FULL_STAGE_BYTES 32768 +#define SMEM_SMEM_Q_FULL_STRIDE 32768 +#define SMEM_SMEM_K_FULL_OFF 33792 +#define SMEM_SMEM_K_FULL_STAGE_BYTES 32768 +#define SMEM_SMEM_K_FULL_STRIDE 32768 +#define SMEM_SMEM_V_FULL_OFF 132096 +#define SMEM_SMEM_V_FULL_STAGE_BYTES 32768 +#define SMEM_SMEM_V_FULL_STRIDE 32768 +#define SMEM_SMEM_STATS_MAX_OFF 197632 +#define SMEM_SMEM_STATS_MAX_STAGE_BYTES 1024 +#define SMEM_SMEM_STATS_MAX_STRIDE 1024 +#define SMEM_SMEM_STATS_SUM_OFF 198656 +#define SMEM_SMEM_STATS_SUM_STAGE_BYTES 512 +#define SMEM_SMEM_STATS_SUM_STRIDE 512 +#define SMEM_SMEM_STATS_FINAL_MAX_OFF 199168 +#define SMEM_SMEM_STATS_FINAL_MAX_STAGE_BYTES 512 +#define SMEM_SMEM_STATS_FINAL_MAX_STRIDE 512 +#define SMEM_SMEM_SOFTMAX_WARP_PAIR_EXCHANGE_OFF 199680 +#define SMEM_SMEM_SOFTMAX_WARP_PAIR_EXCHANGE_STAGE_BYTES 1024 +#define SMEM_SMEM_SOFTMAX_WARP_PAIR_EXCHANGE_STRIDE 1024 +#define SMEM_SMEM_CORR_WARP_PAIR_EXCHANGE_OFF 200704 +#define SMEM_SMEM_CORR_WARP_PAIR_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_CORR_WARP_PAIR_EXCHANGE_STRIDE 512 +#define SMEM_SMEM_SPARSE_INDICES_OFF 201216 +#define SMEM_SMEM_SPARSE_INDICES_STAGE_BYTES 1024 +#define SMEM_SMEM_SPARSE_INDICES_STRIDE 1024 +#define SMEM_SMEM_P_FP8_OFF 207360 +#define SMEM_SMEM_P_FP8_STAGE_BYTES 8192 +#define SMEM_SMEM_P_FP8_STRIDE 8192 +#define SMEM_TOTAL 223744 +#define THREADS 512 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count) : "memory"); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f8f6f4_cta2( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" + "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "tcgen05.mma.cta_group::2.kind::f8f6f4 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" + "}\n" + :: "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x4(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x4.b32" + " {%0, %1, %2, %3}, [%4];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); +} + + +__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ float2 ex2_emulation_f32x2_value(float2 value) { + const float c0 = 1.0f, c1 = 0.695146143436431884765625f; + const float c2 = 0.227564394474029541015625f, c3 = 0.077119089663028717041015625f; + const float magic = 12582912.0f; + float x0 = max_noftz(value.x, -127.0f), x1 = max_noftz(value.y, -127.0f); + float2 xc2 = make_float2(x0, x1), magic2 = make_float2(magic, magic); + float2 xr2; + asm("add.rm.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xr2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&magic2)); + float2 c3_2 = make_float2(c3, c3), c2_2 = make_float2(c2, c2); + float2 c1_2 = make_float2(c1, c1), c0_2 = make_float2(c0, c0); + float2 xrb2, xfrac2; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xrb2) + : "l"(*(unsigned long long*)&xr2), "l"(*(unsigned long long*)&magic2)); + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xfrac2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&xrb2)); + float2 poly2; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&c3_2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c2_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c1_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c0_2)); + int x0r_i, x1r_i, p0_i, p1_i; + asm("mov.b64 {%0, %1}, %2;" : "=r"(x0r_i), "=r"(x1r_i) : "l"(*(unsigned long long*)&xr2)); + asm("mov.b64 {%0, %1}, %2;" : "=r"(p0_i), "=r"(p1_i) : "l"(*(unsigned long long*)&poly2)); + float r0, r1; + asm("mov.b32 %0, %1;" : "=f"(r0) : "r"((x0r_i << 23) + p0_i)); + asm("mov.b32 %0, %1;" : "=f"(r1) : "r"((x1r_i << 23) + p1_i)); + return make_float2(r0, r1); +} + +__device__ __forceinline__ void ex2_emulation_f32x2(float* x0_ptr, float* x1_ptr) { + float2 result = ex2_emulation_f32x2_value(make_float2(*x0_ptr, *x1_ptr)); + *x0_ptr = result.x; *x1_ptr = result.y; +} + +__device__ __forceinline__ void softmax_frag_exp2_cast( + float* sv, uint32_t* pv, int use_emu) +{ + #pragma unroll + for (int j = 0; j < 16; j++) { + if (use_emu && j >= 12) + ex2_emulation_f32x2(&sv[j*2], &sv[j*2+1]); + else { + sv[j*2] = approx_exp2(sv[j*2]); + sv[j*2+1] = approx_exp2(sv[j*2+1]); + } + } + #pragma unroll + for (int j = 0; j < 16; j++) { + __nv_bfloat162 bf = __float22bfloat162_rn({sv[j*2], sv[j*2+1]}); + pv[j] = reinterpret_cast(bf); + } +} + + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_2d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3}], [%4];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr) { + // Canonical .shared::cta form for non-multicast gather4, matching + // trtllm-gen / cuda_ptx and the PTX ISA qualifier order + // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, + // .shared::cluster is reserved for the multicast variant (ctaMask). + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem_cta2( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr) { + // Canonical .shared::cta form; see tma_gather4_gmem2smem above. + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem_mc( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr, unsigned short cta_mask) { + // Multicast variant: the PTX grammar ties the .shared::cluster + // destination to .multicast::cluster + ctaMask (cf. cuda_ptx / + // SM100_TMA_LOAD_MULTICAST_2D_GATHER4 in CUTLASS). + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" + ".mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem_mc_cta2( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr, unsigned short cta_mask) { + // Multicast + cta_group::2 variant; see tma_gather4_gmem2smem_mc. + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" + ".mbarrier::complete_tx::bytes.multicast::cluster.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %1;\n\t" + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" + "}\n" + :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); +} + + +__device__ __forceinline__ void tmem_st_x8_u32(int addr, const uint32_t* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x8.b32" + " [%0], {%1,%2,%3,%4,%5,%6,%7,%8};" + :: "r"(addr), + "r"(src[0]), "r"(src[1]), "r"(src[2]), "r"(src[3]), + "r"(src[4]), "r"(src[5]), "r"(src[6]), "r"(src[7])); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(512, 1) __cluster_dims__(2,1,1) void +kernel_cake_dsv4_983db4c490ceed3c455a(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int mbar_base = smem; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define k_full_addr (mbar_base + 16) + #define k_empty_addr (mbar_base + 40) + #define v_full_addr (mbar_base + 64) + #define v_empty_addr (mbar_base + 80) + #define index_full_addr (mbar_base + 96) + #define index_empty_addr (mbar_base + 144) + #define s_full_addr (mbar_base + 192) + #define p_full_addr (mbar_base + 208) + #define corr_done_addr (mbar_base + 224) + #define stats_addr (mbar_base + 240) + #define sum_ready_addr (mbar_base + 256) + #define o_done_addr (mbar_base + 264) + #define pv_done_addr (mbar_base + 272) + #define s_seeded_addr (mbar_base + 280) + #define q_pair_ready_addr (mbar_base + 288) + #define kv_pair_ready_addr (mbar_base + 296) + #define pv_pair_ready_addr (mbar_base + 304) + #define tmem_dealloc_addr (mbar_base + 320) + #define tmem_dealloc_peer_addr (mbar_base + 328) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 2; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + + // Kernel setup ops + uint8_t* smem_q = reinterpret_cast(smem_raw + 1024); + const int smem_q_addr = smem + 1024; + uint8_t* smem_kv = reinterpret_cast(smem_raw + 33792); + const int smem_kv_addr = smem + 33792; + uint8_t* smem_v = reinterpret_cast(smem_raw + 33792); + const int smem_v_addr = smem + 33792; + uint8_t* smem_q_full = reinterpret_cast(smem_raw + 1024); + const int smem_q_full_addr = smem + 1024; + uint8_t* smem_k_full = reinterpret_cast(smem_raw + 33792); + const int smem_k_full_addr = smem + 33792; + uint8_t* smem_v_full = reinterpret_cast(smem_raw + 132096); + const int smem_v_full_addr = smem + 132096; + float* smem_stats_max = reinterpret_cast(smem_raw + 197632); + const int smem_stats_max_addr = smem + 197632; + float* smem_stats_sum = reinterpret_cast(smem_raw + 198656); + const int smem_stats_sum_addr = smem + 198656; + float* smem_stats_final_max = reinterpret_cast(smem_raw + 199168); + const int smem_stats_final_max_addr = smem + 199168; + float* smem_softmax_warp_pair_exchange = reinterpret_cast(smem_raw + 199680); + const int smem_softmax_warp_pair_exchange_addr = smem + 199680; + float* smem_corr_warp_pair_exchange = reinterpret_cast(smem_raw + 200704); + const int smem_corr_warp_pair_exchange_addr = smem + 200704; + int* smem_sparse_indices = reinterpret_cast(smem_raw + 201216); + const int smem_sparse_indices_addr = smem + 201216; + uint8_t* smem_p_fp8 = reinterpret_cast(smem_raw + 207360); + const int smem_p_fp8_addr = smem + 207360; + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // Mbarrier init (21 pipeline groups, 0 ordered-sequence groups, 42 barriers) + // Mbarriers at smem_raw[0..336) + + if (warp == 12) { + uint32_t leader = elect_sync(); + if (leader) { + // q_full: 1 barriers, init_count=2 + mbarrier_init(smem + 0, 2); + // q_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 8, 1); + // --- pipeline 'k_pipe' --- + // k_full: 3 barriers, init_count=2 + mbarrier_init(smem + 16, 2); + mbarrier_init(smem + 24, 2); + mbarrier_init(smem + 32, 2); + // k_empty: 3 barriers, init_count=1 + mbarrier_init(smem + 40, 1); + mbarrier_init(smem + 48, 1); + mbarrier_init(smem + 56, 1); + // --- pipeline 'v_pipe' --- + // v_full: 2 barriers, init_count=2 + mbarrier_init(smem + 64, 2); + mbarrier_init(smem + 72, 2); + // v_empty: 2 barriers, init_count=1 + mbarrier_init(smem + 80, 1); + mbarrier_init(smem + 88, 1); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + if (warp == 9) { + uint32_t leader = elect_sync(); + if (leader) { + // --- pipeline 'index_pipe' --- + // index_full: 6 barriers, init_count=32 + mbarrier_init(smem + 96, 32); + mbarrier_init(smem + 104, 32); + mbarrier_init(smem + 112, 32); + mbarrier_init(smem + 120, 32); + mbarrier_init(smem + 128, 32); + mbarrier_init(smem + 136, 32); + // index_empty: 6 barriers, init_count=256 + mbarrier_init(smem + 144, 256); + mbarrier_init(smem + 152, 256); + mbarrier_init(smem + 160, 256); + mbarrier_init(smem + 168, 256); + mbarrier_init(smem + 176, 256); + mbarrier_init(smem + 184, 256); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // s_full: stages (0,), init_count=1 + mbarrier_init(smem + 192, 1); + // o_done: 1 barriers, init_count=1 + mbarrier_init(smem + 264, 1); + // pv_done: 1 barriers, init_count=1 + mbarrier_init(smem + 272, 1); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + if (warp == 1) { + uint32_t leader = elect_sync(); + if (leader) { + // s_full: stages (1,), init_count=1 + mbarrier_init(smem + 200, 1); + // p_full: 2 barriers, init_count=128 + mbarrier_init(smem + 208, 128); + mbarrier_init(smem + 216, 128); + // corr_done: stages (0,), init_count=128 + mbarrier_init(smem + 224, 128); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + if (warp == 5) { + uint32_t leader = elect_sync(); + if (leader) { + // corr_done: stages (1,), init_count=128 + mbarrier_init(smem + 232, 128); + // stats: 2 barriers, init_count=128 + mbarrier_init(smem + 240, 128); + mbarrier_init(smem + 248, 128); + // sum_ready: 1 barriers, init_count=128 + mbarrier_init(smem + 256, 128); + // tmem_dealloc_peer: 1 barriers, init_count=32 + mbarrier_init(smem + 328, 32); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + if (warp == 2) { + uint32_t leader = elect_sync(); + if (leader) { + // s_seeded: 1 barriers, init_count=256 + mbarrier_init(smem + 280, 256); + // q_pair_ready: 1 barriers, init_count=64 + mbarrier_init(smem + 288, 64); + // kv_pair_ready: 1 barriers, init_count=64 + mbarrier_init(smem + 296, 64); + // pv_pair_ready: stages (0,), init_count=64 + mbarrier_init(smem + 304, 64); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + if (warp == 4) { + uint32_t leader = elect_sync(); + if (leader) { + // pv_pair_ready: stages (1,), init_count=64 + mbarrier_init(smem + 312, 64); + // tmem_dealloc: 1 barriers, init_count=416 + mbarrier_init(smem + 320, 416); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncwarp(); + + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 336); + if (warp == 0) { + int _tmem_hold = smem + 336; + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); + } + + __syncthreads(); + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_tmem_scratch = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Inc phase consumes the registers released above. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 136;"); + } + + // ---- Role: index_warp ---- + if (warp == 9) { + { // index_warp_main + const int index_dummy = 0; + unsigned int _phase_index_empty = 1; + unsigned int _phase_q_full_0 = 0; + { + int all_num_kv_tiles = (sparse_topk + 128 - 1) / 128; + unsigned int index_prod_stage = 0; + #pragma unroll + for (unsigned int static_work = 0; static_work < 1; static_work++) { + int split_idx = 0; + int query_idx = bid / 2; + int tiles_per_split = all_num_kv_tiles + 1 - 1; + int first_tile = split_idx * tiles_per_split; + int num_index_passes = (tiles_per_split + 1) / 2; + #pragma unroll 1 + for (int index_pass = 0; index_pass < num_index_passes; index_pass++) { + mbarrier_wait(index_empty_addr + (index_prod_stage) * 8, _phase_index_empty); + int index_stage_base = smem_sparse_indices_addr + index_prod_stage * 1024; + int stage_row_base = index_pass * 256; + #pragma unroll + for (int index_half = 0; index_half < 2; index_half++) { + int index_offset = index_half * 128 + lane * 4; + int tile_col_base = first_tile * 128 + stage_row_base + index_half * 128; + int* row_tile_ptr = ((tile_col_base < 128) ? swa_indices + (query_idx * swa_index_stride) : compressed_indices + (query_idx * compressed_index_stride + (tile_col_base - 128))); + int lane_col = tile_col_base + lane * 4; + if (lane_col + 4 <= sparse_topk) { + asm volatile("cp.async.cg.shared::cta.global.L2::128B [%0], [%1], 16;" + :: "r"(index_stage_base + index_offset * 4), "l"(row_tile_ptr + (lane * 4))); + } else { + int tail_rows[4]; + #pragma unroll + for (int row_i = 0; row_i < 4; row_i++) { + tail_rows[row_i] = -1; + } + #pragma unroll + for (int row_i_1 = 0; row_i_1 < 4; row_i_1++) { + if (lane_col + row_i_1 < sparse_topk) { + tail_rows[row_i_1] = row_tile_ptr[lane * 4 + row_i_1]; + } + } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(smem_sparse_indices_addr + ((unsigned int)index_stage_base - smem_sparse_indices_addr + (unsigned int)(index_offset * 4))), "r"(tail_rows[0]), "r"(tail_rows[1]), "r"(tail_rows[2]), "r"(tail_rows[3]) : "memory"); + } + } + asm volatile("cp.async.commit_group;"); + asm volatile("cp.async.wait_group 0;"); + mbarrier_arrive(index_full_addr + (index_prod_stage) * 8); + index_prod_stage += 1; + if (index_prod_stage == 6) { index_prod_stage = 0; _phase_index_empty ^= 1; } + } + } + } + } + } + // ---- Role: load_warp ---- + if (warp >= 12 && warp <= 15) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 72;"); + { // load_warp_main + const int wg2_dummy = 0; + const int load_warp_rank = warp - ((1) ? 12 : 9); + int all_num_kv_tiles_1 = (sparse_topk + 128 - 1) / 128; + unsigned int load_k_stage = 0; + unsigned int load_v_stage = 0; + unsigned int load_kv_stage = 0; + unsigned int load_k_index_stage = 0; + unsigned int load_v_index_stage = 0; + int k_cta_offset = bid % 2 * 64; + unsigned int _phase_q_empty_0 = 1; + unsigned int _phase_index_full = 0; + unsigned int _phase_k_empty = 1; + unsigned int _phase_k_empty_1 = 1; + unsigned int _phase_v_empty = 1; + #pragma unroll + for (unsigned int work_idx = 0; work_idx < 1; work_idx++) { + int split_idx_1 = 0; + int query_idx_1 = bid / 2 >> 1; + int v_chunk = bid / 2 & 1; + { + split_idx_1 = 0; + query_idx_1 = bid / 2; + v_chunk = bid % 2; + } + int tiles_per_split_1 = all_num_kv_tiles_1 + 1 - 1; + int first_tile_1 = split_idx_1 * tiles_per_split_1; + int num_kv_tiles = tiles_per_split_1; + int sparse_extent = num_kv_tiles * 128; + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + if (load_warp_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + #pragma unroll + for (int q_stage = 0; q_stage < 4; q_stage++) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global.mbarrier::complete_tx::bytes.multicast::cluster.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5}], [%6], %7;" + :: "r"(smem_q_full_addr + (unsigned int)(q_stage * 8192)), "l"((&tmap_q)), "r"(0), "r"(bid % 2 * 64), "r"(q_stage), "r"(query_idx_1), + "r"(((q_full_addr) & 0xFEFFFFFF)), "h"((uint16_t)(1 << bid % 2)) : "memory"); + } + } + } + { + mbarrier_wait(index_full_addr + (load_k_index_stage) * 8, _phase_index_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + } + int k_index_stage_base = smem_sparse_indices_addr + load_k_index_stage * 1024; + { + mbarrier_wait(k_empty_addr + (load_k_stage) * 8, _phase_k_empty); + if (load_warp_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + } + } + int k_dst = smem_k_full_addr + load_k_stage * 32768 + (unsigned int)(load_warp_rank * 8192); + int group = lane; + if (group < 16) { + int group_offset = k_cta_offset + group * 4; + int raw_rows[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows[0])), "=r"(*reinterpret_cast(&raw_rows[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows[(0) + 3])) + : "r"(k_index_stage_base + group_offset * 4)); + int row0 = ((raw_rows[0] >= 0) ? raw_rows[0] : 0); + int row1 = ((raw_rows[1] >= 0) ? raw_rows[1] : 0); + int row2 = ((raw_rows[2] >= 0) ? raw_rows[2] : 0); + int row3 = ((raw_rows[3] >= 0) ? raw_rows[3] : 0); + if (first_tile_1 == 0) { + tma_gather4_gmem2smem_mc_cta2(k_dst + group * 512, (&tmap_swa_kv), load_warp_rank * 128, row0, row1, row2, row3, ((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); + } else { + tma_gather4_gmem2smem_mc_cta2(k_dst + group * 512, (&tmap_compressed_kv), load_warp_rank * 128, row0, row1, row2, row3, ((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); + } + } + load_k_stage += 1; + if (load_k_stage == 3) { load_k_stage = 0; _phase_k_empty ^= 1; } + } + #pragma unroll 1 + for (int tile = 1; tile < num_kv_tiles; tile++) { + if ((1 & (int)((tile & 1) == 0)) != 0) { + mbarrier_wait(index_full_addr + (load_k_index_stage) * 8, _phase_index_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + } + k_index_stage_base = smem_sparse_indices_addr + load_k_index_stage * 1024; + { + mbarrier_wait(k_empty_addr + (load_k_stage) * 8, _phase_k_empty); + if (load_warp_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + } + } + int k_dst_1 = smem_k_full_addr + load_k_stage * 32768 + (unsigned int)(load_warp_rank * 8192); + int group_1 = lane; + if (group_1 < 16) { + int group_offset_1 = (tile & 1) * 128 + k_cta_offset + group_1 * 4; + int raw_rows_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_1[0])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 3])) + : "r"(k_index_stage_base + group_offset_1 * 4)); + int row0_1 = ((raw_rows_1[0] >= 0) ? raw_rows_1[0] : 0); + int row1_1 = ((raw_rows_1[1] >= 0) ? raw_rows_1[1] : 0); + int row2_1 = ((raw_rows_1[2] >= 0) ? raw_rows_1[2] : 0); + int row3_1 = ((raw_rows_1[3] >= 0) ? raw_rows_1[3] : 0); + tma_gather4_gmem2smem_mc_cta2(k_dst_1 + group_1 * 512, (&tmap_compressed_kv), load_warp_rank * 128, row0_1, row1_1, row2_1, row3_1, ((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); + } + load_k_stage += 1; + if (load_k_stage == 3) { load_k_stage = 0; _phase_k_empty ^= 1; } + if ((tile & 1) != 0) { + load_k_index_stage += 1; + if (load_k_index_stage == 6) { load_k_index_stage = 0; _phase_index_full ^= 1; } + } + } + int prev_tile = tile - 1; + int global_prev_tile = first_tile_1 + prev_tile; + int v_index_stage_base = smem_sparse_indices_addr + load_v_index_stage * 1024; + { + mbarrier_wait(v_empty_addr + (load_v_stage) * 8, _phase_v_empty); + if (load_warp_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + } + } + if (lane < 16) { + int v_dst = smem_v_full_addr + load_v_stage * 32768 + (unsigned int)(load_warp_rank / 2 * 16384); + int v_col = v_chunk * 256 + load_warp_rank / 2 * 128; + int group_2 = (load_warp_rank & 1) * 16 + lane; + int group_offset_2 = (prev_tile & 1) * 128 + group_2 * 4; + int raw_rows_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_2[0])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 3])) + : "r"(v_index_stage_base + group_offset_2 * 4)); + int row0_2 = ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : 0); + int row1_2 = ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : 0); + int row2_2 = ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : 0); + int row3_2 = ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : 0); + if (global_prev_tile == 0) { + tma_gather4_gmem2smem_mc_cta2(v_dst + group_2 * 512, (&tmap_swa_kv), v_col, row0_2, row1_2, row2_2, row3_2, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); + } else { + tma_gather4_gmem2smem_mc_cta2(v_dst + group_2 * 512, (&tmap_compressed_kv), v_col, row0_2, row1_2, row2_2, row3_2, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); + } + } + load_v_stage += 1; + if (load_v_stage == 2) { load_v_stage = 0; _phase_v_empty ^= 1; } + if ((prev_tile & 1) != 0) { + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(index_empty_addr + (load_v_index_stage) * 8); + load_v_index_stage = load_v_index_stage + 1; + if (load_v_index_stage == 6) { + load_v_index_stage = 0; + } + } + } + } + int last_tile = num_kv_tiles - 1; + int global_last_tile = first_tile_1 + last_tile; + { + mbarrier_wait(v_empty_addr + (load_v_stage) * 8, _phase_v_empty); + if (load_warp_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + } + } + if (lane < 16) { + int v_dst_1 = smem_v_full_addr + load_v_stage * 32768 + (unsigned int)(load_warp_rank / 2 * 16384); + int v_col_1 = v_chunk * 256 + load_warp_rank / 2 * 128; + int group_3 = (load_warp_rank & 1) * 16 + lane; + int group_offset_3 = (last_tile & 1) * 128 + group_3 * 4; + int raw_rows_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_3[0])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 3])) + : "r"(smem_sparse_indices_addr + load_v_index_stage * 1024 + (unsigned int)(group_offset_3 * 4))); + int row0_3 = ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : 0); + int row1_3 = ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : 0); + int row2_3 = ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : 0); + int row3_3 = ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : 0); + if (global_last_tile == 0) { + tma_gather4_gmem2smem_mc_cta2(v_dst_1 + group_3 * 512, (&tmap_swa_kv), v_col_1, row0_3, row1_3, row2_3, row3_3, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); + } else { + tma_gather4_gmem2smem_mc_cta2(v_dst_1 + group_3 * 512, (&tmap_compressed_kv), v_col_1, row0_3, row1_3, row2_3, row3_3, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); + } + } + load_v_stage += 1; + if (load_v_stage == 2) { load_v_stage = 0; _phase_v_empty ^= 1; } + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(index_empty_addr + (load_v_index_stage) * 8); + load_v_index_stage = load_v_index_stage + 1; + if (load_v_index_stage == 6) { + load_v_index_stage = 0; + } + } + } + mbarrier_arrive(tmem_dealloc_addr); + } + } + // ---- Role: softmax_wg ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 152;"); + { // softmax_wg_main + float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; + const int wg_dummy_inc = 0; + int all_num_kv_tiles_2 = (sparse_topk + 128 - 1) / 128; + const int tmem_row_base = ((1) ? warp % 2 * 32 : warp % 4 * 32); + const int tmem_score_row_base = ((1) ? (int)(warp % 4 * 32) : tmem_row_base); + const int n_half = ((1) ? (int)(warp % 4 / 2) : 0); + const int my_row = tmem_row_base + lane; + const int stats_row = n_half * 64 + my_row; + int softmax_tile_cursor = 0; + unsigned int softmax_index_stage = 0; + unsigned int _phase_q_full_0_1 = 0; + #pragma unroll + for (unsigned int work_idx_1 = 0; work_idx_1 < 1; work_idx_1++) { + int split_idx_2 = 0; + int query_idx_2 = bid / 2 >> 1; + { + split_idx_2 = 0; + query_idx_2 = bid / 2; + } + int tiles_per_split_2 = all_num_kv_tiles_2 + 1 - 1; + int first_tile_2 = split_idx_2 * tiles_per_split_2; + int num_kv_tiles_1 = tiles_per_split_2; + int _max_0 = ((sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) : (0)); + int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); + int active_topk = _min_0; + float row_max_val = -CAKE_INF; + float row_sum_val = 0.0f; + int sink_head = ((1) ? bid % 2 * 64 + my_row : my_row); + if (has_sinks != 0 && sink_head < num_heads && split_idx_2 == 0) { + row_max_val = sinks[sink_head] * 1.4426950408889634f / softmax_scale_log2; + row_sum_val = 1.0f; + } + #pragma unroll 1 + for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { + int pipeline_tile = softmax_tile_cursor + tile_1; + int phase = pipeline_tile & 1; + int s_wait_phase = pipeline_tile >> 1 & 1; + mbarrier_wait(s_full_addr + (phase) * 8, s_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int s_off = ((phase != 0) ? 128 : 0); + int s_base = taddr + (unsigned int)s_off + (unsigned int)(tmem_score_row_base << 16); + float new_max = row_max_val; + int valid_sparse_cols = ((active_topk < sparse_topk) ? active_topk : sparse_topk); + valid_sparse_cols = valid_sparse_cols - (first_tile_2 + tile_1) * 128 - n_half * 64; + if (valid_sparse_cols < 0) { + valid_sparse_cols = 0; + } + if (valid_sparse_cols > 64) { + valid_sparse_cols = 64; + } + int index_tile_addr = smem_sparse_indices_addr + softmax_index_stage * 1024 + (unsigned int)((tile_1 & 1) * 512); + int index_lane_addr = index_tile_addr + (n_half * 64 + lane) * 4; + int staged_index[2]; + unsigned int invalid_cols[2]; + unsigned int any_invalid = 0; + #pragma unroll + for (int w = 0; w < 2; w++) { + asm volatile("ld.shared.b32 %0, [%1];" : "=r"(*reinterpret_cast(&staged_index[w])) : "r"(index_lane_addr + w * 128)); + unsigned int _vote_0 = __ballot_sync(0xFFFFFFFF, staged_index[w] < 0); + invalid_cols[w] = _vote_0; + any_invalid = any_invalid | invalid_cols[w]; + } + if ((tile_1 & 1) != 0 || tile_1 + 1 == num_kv_tiles_1) { + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(index_empty_addr + (softmax_index_stage) * 8); + softmax_index_stage = softmax_index_stage + 1; + if (softmax_index_stage == 6) { + softmax_index_stage = 0; + } + } + float _tmem_load_0[4]; + tmem_ld_x4(&_tmem_load_0[0], s_base); + float _tmem_load_1[64]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]) + : "r"(s_base)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]) + : "r"(s_base + 32)); + { + uint32_t _slice_lo_mask_0; + { + int _lim_0 = valid_sparse_cols; + if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_0 & (1u << 0))) _tmem_load_1[0] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 1))) _tmem_load_1[1] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 2))) _tmem_load_1[2] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 3))) _tmem_load_1[3] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 4))) _tmem_load_1[4] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 5))) _tmem_load_1[5] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 6))) _tmem_load_1[6] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 7))) _tmem_load_1[7] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 8))) _tmem_load_1[8] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 9))) _tmem_load_1[9] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 10))) _tmem_load_1[10] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 11))) _tmem_load_1[11] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 12))) _tmem_load_1[12] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 13))) _tmem_load_1[13] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 14))) _tmem_load_1[14] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 15))) _tmem_load_1[15] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 16))) _tmem_load_1[16] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 17))) _tmem_load_1[17] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 18))) _tmem_load_1[18] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 19))) _tmem_load_1[19] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 20))) _tmem_load_1[20] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 21))) _tmem_load_1[21] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 22))) _tmem_load_1[22] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 23))) _tmem_load_1[23] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 24))) _tmem_load_1[24] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 25))) _tmem_load_1[25] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 26))) _tmem_load_1[26] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 27))) _tmem_load_1[27] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 28))) _tmem_load_1[28] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 29))) _tmem_load_1[29] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 30))) _tmem_load_1[30] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 31))) _tmem_load_1[31] = -CAKE_INF; + uint32_t _slice_lo_mask_1; + { + int _lim_1 = valid_sparse_cols - 32; + if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) _tmem_load_1[32] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) _tmem_load_1[33] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) _tmem_load_1[34] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) _tmem_load_1[35] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) _tmem_load_1[36] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) _tmem_load_1[37] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) _tmem_load_1[38] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) _tmem_load_1[39] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) _tmem_load_1[40] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) _tmem_load_1[41] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) _tmem_load_1[42] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) _tmem_load_1[43] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) _tmem_load_1[44] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) _tmem_load_1[45] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) _tmem_load_1[46] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) _tmem_load_1[47] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) _tmem_load_1[48] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) _tmem_load_1[49] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) _tmem_load_1[50] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) _tmem_load_1[51] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) _tmem_load_1[52] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) _tmem_load_1[53] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) _tmem_load_1[54] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) _tmem_load_1[55] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) _tmem_load_1[56] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) _tmem_load_1[57] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) _tmem_load_1[58] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) _tmem_load_1[59] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) _tmem_load_1[60] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) _tmem_load_1[61] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) _tmem_load_1[62] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) _tmem_load_1[63] = -CAKE_INF; + if (any_invalid != 0) { + #pragma unroll + for (int w_1 = 0; w_1 < 2; w_1++) { + unsigned int mask_word = invalid_cols[w_1]; + _tmem_load_1[w_1 * 32] = (((mask_word & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32]); + _tmem_load_1[w_1 * 32 + 1] = (((mask_word >> 1 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 1]); + _tmem_load_1[w_1 * 32 + 2] = (((mask_word >> 2 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 2]); + _tmem_load_1[w_1 * 32 + 3] = (((mask_word >> 3 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 3]); + _tmem_load_1[w_1 * 32 + 4] = (((mask_word >> 4 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 4]); + _tmem_load_1[w_1 * 32 + 5] = (((mask_word >> 5 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 5]); + _tmem_load_1[w_1 * 32 + 6] = (((mask_word >> 6 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 6]); + _tmem_load_1[w_1 * 32 + 7] = (((mask_word >> 7 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 7]); + _tmem_load_1[w_1 * 32 + 8] = (((mask_word >> 8 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 8]); + _tmem_load_1[w_1 * 32 + 9] = (((mask_word >> 9 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 9]); + _tmem_load_1[w_1 * 32 + 10] = (((mask_word >> 10 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 10]); + _tmem_load_1[w_1 * 32 + 11] = (((mask_word >> 11 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 11]); + _tmem_load_1[w_1 * 32 + 12] = (((mask_word >> 12 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 12]); + _tmem_load_1[w_1 * 32 + 13] = (((mask_word >> 13 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 13]); + _tmem_load_1[w_1 * 32 + 14] = (((mask_word >> 14 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 14]); + _tmem_load_1[w_1 * 32 + 15] = (((mask_word >> 15 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 15]); + _tmem_load_1[w_1 * 32 + 16] = (((mask_word >> 16 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 16]); + _tmem_load_1[w_1 * 32 + 17] = (((mask_word >> 17 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 17]); + _tmem_load_1[w_1 * 32 + 18] = (((mask_word >> 18 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 18]); + _tmem_load_1[w_1 * 32 + 19] = (((mask_word >> 19 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 19]); + _tmem_load_1[w_1 * 32 + 20] = (((mask_word >> 20 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 20]); + _tmem_load_1[w_1 * 32 + 21] = (((mask_word >> 21 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 21]); + _tmem_load_1[w_1 * 32 + 22] = (((mask_word >> 22 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 22]); + _tmem_load_1[w_1 * 32 + 23] = (((mask_word >> 23 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 23]); + _tmem_load_1[w_1 * 32 + 24] = (((mask_word >> 24 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 24]); + _tmem_load_1[w_1 * 32 + 25] = (((mask_word >> 25 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 25]); + _tmem_load_1[w_1 * 32 + 26] = (((mask_word >> 26 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 26]); + _tmem_load_1[w_1 * 32 + 27] = (((mask_word >> 27 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 27]); + _tmem_load_1[w_1 * 32 + 28] = (((mask_word >> 28 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 28]); + _tmem_load_1[w_1 * 32 + 29] = (((mask_word >> 29 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 29]); + _tmem_load_1[w_1 * 32 + 30] = (((mask_word >> 30 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 30]); + _tmem_load_1[w_1 * 32 + 31] = (((mask_word >> 31 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 31]); + } + } + float2 _reg_reduce_max2_2 = {-CAKE_INF, -CAKE_INF}; + row_max_x32_accum(&_tmem_load_1[0], _reg_reduce_max2_2); + row_max_x32_accum(&_tmem_load_1[32], _reg_reduce_max2_2); + float _tmem_load_1_max = row_max_reduce(_reg_reduce_max2_2); + float _max_1 = max_noftz(new_max, _tmem_load_1_max); + new_max = _max_1; + smem_softmax_warp_pair_exchange[phase * 128 + stats_row] = new_max; + asm volatile("barrier.sync %0, 64;" :: "r"(2 + warp % 2) : "memory"); + float _max_2 = max_noftz(new_max, smem_softmax_warp_pair_exchange[phase * 128 + (stats_row ^ 64)]); + new_max = _max_2; + } + float _fma_0 = __fmaf_rn(row_max_val, softmax_scale_log2, (-new_max) * softmax_scale_log2); + float delta = _fma_0; + float _exp2_0 = approx_exp2(delta); + float exp_delta = _exp2_0; + float acc_scale = ((row_max_val > -CAKE_INF) ? exp_delta : 1.0f); + smem_stats_max[phase * 128 + stats_row] = acc_scale; + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(stats_addr + (phase) * 8); + row_max_val = new_max; + float safe_max = ((new_max == -CAKE_INF) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + float block_sum = 0.0f; + { + { + const float2 _fma_b2_3 = {softmax_scale_log2, softmax_scale_log2}; + const float2 _fma_c2_4 = {-max_scaled, -max_scaled}; + #pragma unroll + for (int _lf = 0; _lf < 32; _lf++) + fma_f32x2_inplace(&reinterpret_cast(_tmem_load_1)[_lf], _fma_b2_3, _fma_c2_4); + #pragma unroll + for (int _le = 0; _le < 64; _le++) { + _tmem_load_1[_le] = approx_exp2(_tmem_load_1[_le]); + } + uint32_t _fp8_0[16]; + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[0]), "f"(_tmem_load_1[1]), + "f"(_tmem_load_1[2]), "f"(_tmem_load_1[3])); + _fp8_0[0] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[4]), "f"(_tmem_load_1[5]), + "f"(_tmem_load_1[6]), "f"(_tmem_load_1[7])); + _fp8_0[1] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[8]), "f"(_tmem_load_1[9]), + "f"(_tmem_load_1[10]), "f"(_tmem_load_1[11])); + _fp8_0[2] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[12]), "f"(_tmem_load_1[13]), + "f"(_tmem_load_1[14]), "f"(_tmem_load_1[15])); + _fp8_0[3] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[16]), "f"(_tmem_load_1[17]), + "f"(_tmem_load_1[18]), "f"(_tmem_load_1[19])); + _fp8_0[4] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[20]), "f"(_tmem_load_1[21]), + "f"(_tmem_load_1[22]), "f"(_tmem_load_1[23])); + _fp8_0[5] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[24]), "f"(_tmem_load_1[25]), + "f"(_tmem_load_1[26]), "f"(_tmem_load_1[27])); + _fp8_0[6] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[28]), "f"(_tmem_load_1[29]), + "f"(_tmem_load_1[30]), "f"(_tmem_load_1[31])); + _fp8_0[7] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[32]), "f"(_tmem_load_1[33]), + "f"(_tmem_load_1[34]), "f"(_tmem_load_1[35])); + _fp8_0[8] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[36]), "f"(_tmem_load_1[37]), + "f"(_tmem_load_1[38]), "f"(_tmem_load_1[39])); + _fp8_0[9] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[40]), "f"(_tmem_load_1[41]), + "f"(_tmem_load_1[42]), "f"(_tmem_load_1[43])); + _fp8_0[10] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[44]), "f"(_tmem_load_1[45]), + "f"(_tmem_load_1[46]), "f"(_tmem_load_1[47])); + _fp8_0[11] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[48]), "f"(_tmem_load_1[49]), + "f"(_tmem_load_1[50]), "f"(_tmem_load_1[51])); + _fp8_0[12] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[52]), "f"(_tmem_load_1[53]), + "f"(_tmem_load_1[54]), "f"(_tmem_load_1[55])); + _fp8_0[13] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[56]), "f"(_tmem_load_1[57]), + "f"(_tmem_load_1[58]), "f"(_tmem_load_1[59])); + _fp8_0[14] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[60]), "f"(_tmem_load_1[61]), + "f"(_tmem_load_1[62]), "f"(_tmem_load_1[63])); + _fp8_0[15] = _packed; + } + int p_stage_base = smem_p_fp8_addr + (unsigned int)(phase * 8192); + #pragma unroll + for (int p_vec = 0; p_vec < 4; p_vec++) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_stage_base + (my_row * 128 + (n_half * 4 + p_vec) * 16 ^ (my_row * 128 + (n_half * 4 + p_vec) * 16 >> 7 & 7) << 4))), "r"(_fp8_0[p_vec * 4]), "r"(_fp8_0[p_vec * 4 + 1]), "r"(_fp8_0[p_vec * 4 + 2]), "r"(_fp8_0[p_vec * 4 + 3]) : "memory"); + } + } + } + { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + { + asm volatile("" ::: "memory"); + } + mbarrier_arrive(p_full_addr + (phase) * 8); + { + { + asm volatile("" ::: "memory"); + const float2* _reg_reduce_src2_5 = reinterpret_cast(&_tmem_load_1[0]); + float2 _reg_reduce_sum2_5_0 = make_float2(0.0f, 0.0f); + float2 _reg_reduce_sum2_5_1 = make_float2(0.0f, 0.0f); + float2 _reg_reduce_sum2_5_2 = make_float2(0.0f, 0.0f); + float2 _reg_reduce_sum2_5_3 = make_float2(0.0f, 0.0f); + #pragma unroll + for (int _rr = 0; _rr < 8; _rr++) { + add_f32x2_inplace(&_reg_reduce_sum2_5_0, _reg_reduce_src2_5[_rr * 4 + 0]); + add_f32x2_inplace(&_reg_reduce_sum2_5_1, _reg_reduce_src2_5[_rr * 4 + 1]); + add_f32x2_inplace(&_reg_reduce_sum2_5_2, _reg_reduce_src2_5[_rr * 4 + 2]); + add_f32x2_inplace(&_reg_reduce_sum2_5_3, _reg_reduce_src2_5[_rr * 4 + 3]); + } + add_f32x2_inplace(&_reg_reduce_sum2_5_0, _reg_reduce_sum2_5_1); + add_f32x2_inplace(&_reg_reduce_sum2_5_2, _reg_reduce_sum2_5_3); + add_f32x2_inplace(&_reg_reduce_sum2_5_0, _reg_reduce_sum2_5_2); + float _tmem_load_1_sum = _reg_reduce_sum2_5_0.x + _reg_reduce_sum2_5_0.y; + block_sum = _tmem_load_1_sum; + float _fma_2 = __fmaf_rn(row_sum_val, acc_scale, block_sum); + row_sum_val = _fma_2; + } + } + } + smem_stats_sum[stats_row] = row_sum_val; + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(sum_ready_addr); + softmax_tile_cursor = softmax_tile_cursor + num_kv_tiles_1; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_dealloc_addr); + } + } + // ---- Role: correction_wg ---- + if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 144;"); + { // correction_wg_main + float softmax_scale_log2_1 = bmm1_scale[0] * 1.4426950408889634f; + float output_scale = bmm2_scale[0]; + const int wg_dummy_inc_1 = 0; + int all_num_kv_tiles_3 = (sparse_topk + 128 - 1) / 128; + const int tmem_row_base_1 = ((1) ? warp % 2 * 32 : warp % 4 * 32); + const int n_half_1 = ((1) ? (int)(warp % 4 / 2) : 0); + const int my_row_1 = tmem_row_base_1 + lane; + const int stats_row_1 = n_half_1 * 64 + my_row_1; + const int corr_row = tmem_row_base_1 << 16; + int correction_tile_cursor = 0; + unsigned int _phase_q_full_0_2 = 0; + unsigned int _phase_pv_done_0 = 0; + unsigned int _phase_o_done_0 = 0; + unsigned int _phase_sum_ready_0 = 0; + #pragma unroll + for (unsigned int work_idx_2 = 0; work_idx_2 < 1; work_idx_2++) { + int split_idx_3 = 0; + int query_idx_3 = bid / 2 >> 1; + int v_chunk_1 = bid / 2 & 1; + { + split_idx_3 = 0; + query_idx_3 = bid / 2; + v_chunk_1 = bid % 2; + } + int tiles_per_split_3 = all_num_kv_tiles_3 + 1 - 1; + int first_tile_3 = split_idx_3 * tiles_per_split_3; + int num_kv_tiles_2 = tiles_per_split_3; + #pragma unroll 1 + for (int tile_2 = 0; tile_2 < num_kv_tiles_2; tile_2++) { + int pipeline_tile_1 = correction_tile_cursor + tile_2; + int phase_1 = pipeline_tile_1 & 1; + int stats_wait_phase = pipeline_tile_1 >> 1 & 1; + mbarrier_wait(stats_addr + (phase_1) * 8, stats_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + float acc_scale_1 = smem_stats_max[phase_1 * 128 + stats_row_1]; + if (tile_2 > 0) { + mbarrier_wait(pv_done_addr, _phase_pv_done_0); + _phase_pv_done_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + } + if (tile_2 > 0) { + int _vote_1 = __any_sync(0xFFFFFFFF, acc_scale_1 < 1.0f); + int any_rescale = _vote_1; + if (any_rescale != 0) { + #pragma unroll + for (int vs = 0; vs < 2; vs++) { + int o_base = taddr + 256 + (unsigned int)(vs * 128) + (unsigned int)corr_row; + #pragma unroll + for (int c = 0; c < 128; c += 16) { + float _tmem_load_4[16]; + tmem_ld_x16(&_tmem_load_4[0], o_base + c); + const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + mul_f32x2_inplace(&reinterpret_cast(_tmem_load_4)[_ls], _scale2_0); + tmem_st_x16_f32(o_base + c, _tmem_load_4); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + } + } + } + } + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(corr_done_addr + (phase_1) * 8); + } + mbarrier_wait(o_done_addr, _phase_o_done_0); + _phase_o_done_0 ^= 1; + mbarrier_wait(sum_ready_addr, _phase_sum_ready_0); + _phase_sum_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + float total_sum = smem_stats_sum[stats_row_1]; + float final_max = smem_stats_final_max[stats_row_1]; + { + smem_corr_warp_pair_exchange[stats_row_1] = total_sum; + asm volatile("barrier.sync %0, 64;" :: "r"(4 + warp % 2) : "memory"); + total_sum = total_sum + smem_corr_warp_pair_exchange[stats_row_1 ^ 64]; + } + float _rcp_0 = approx_rcp(total_sum); + float inv_sum = ((total_sum > 0.0f) ? _rcp_0 : 0.0f); + int head_idx = my_row_1; + { + head_idx = bid % 2 * 64 + my_row_1; + } + int direct_o_offset = (query_idx_3 * num_heads + head_idx) * 512 + v_chunk_1 * 256; + int partial_o_offset = (query_idx_3 * num_heads + head_idx + split_idx_3) * 512 + ((0) ? v_chunk_1 : n_half_1) * 256; + int o_offset = partial_o_offset; + #pragma unroll + for (int vs_1 = 0; vs_1 < 2; vs_1++) { + int o_base_epi = taddr + 256 + (unsigned int)(vs_1 * 128) + (unsigned int)corr_row; + #pragma unroll + for (int c_1 = 0; c_1 < 128; c_1 += 64) { + float _tmem_load_5[64]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x64.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, [%64];" + : "=f"(_tmem_load_5[0]), "=f"(_tmem_load_5[1]), "=f"(_tmem_load_5[2]), "=f"(_tmem_load_5[3]), "=f"(_tmem_load_5[4]), "=f"(_tmem_load_5[5]), "=f"(_tmem_load_5[6]), "=f"(_tmem_load_5[7]), "=f"(_tmem_load_5[8]), "=f"(_tmem_load_5[9]), "=f"(_tmem_load_5[10]), "=f"(_tmem_load_5[11]), "=f"(_tmem_load_5[12]), "=f"(_tmem_load_5[13]), "=f"(_tmem_load_5[14]), "=f"(_tmem_load_5[15]), "=f"(_tmem_load_5[16]), "=f"(_tmem_load_5[17]), "=f"(_tmem_load_5[18]), "=f"(_tmem_load_5[19]), "=f"(_tmem_load_5[20]), "=f"(_tmem_load_5[21]), "=f"(_tmem_load_5[22]), "=f"(_tmem_load_5[23]), "=f"(_tmem_load_5[24]), "=f"(_tmem_load_5[25]), "=f"(_tmem_load_5[26]), "=f"(_tmem_load_5[27]), "=f"(_tmem_load_5[28]), "=f"(_tmem_load_5[29]), "=f"(_tmem_load_5[30]), "=f"(_tmem_load_5[31]), "=f"(_tmem_load_5[32]), "=f"(_tmem_load_5[33]), "=f"(_tmem_load_5[34]), "=f"(_tmem_load_5[35]), "=f"(_tmem_load_5[36]), "=f"(_tmem_load_5[37]), "=f"(_tmem_load_5[38]), "=f"(_tmem_load_5[39]), "=f"(_tmem_load_5[40]), "=f"(_tmem_load_5[41]), "=f"(_tmem_load_5[42]), "=f"(_tmem_load_5[43]), "=f"(_tmem_load_5[44]), "=f"(_tmem_load_5[45]), "=f"(_tmem_load_5[46]), "=f"(_tmem_load_5[47]), "=f"(_tmem_load_5[48]), "=f"(_tmem_load_5[49]), "=f"(_tmem_load_5[50]), "=f"(_tmem_load_5[51]), "=f"(_tmem_load_5[52]), "=f"(_tmem_load_5[53]), "=f"(_tmem_load_5[54]), "=f"(_tmem_load_5[55]), "=f"(_tmem_load_5[56]), "=f"(_tmem_load_5[57]), "=f"(_tmem_load_5[58]), "=f"(_tmem_load_5[59]), "=f"(_tmem_load_5[60]), "=f"(_tmem_load_5[61]), "=f"(_tmem_load_5[62]), "=f"(_tmem_load_5[63]) + : "r"(o_base_epi + c_1)); + int gmem_base = o_offset + vs_1 * 128 + c_1; + if (bid % 2 == 0) { + #pragma unroll + for (int j = 0; j < 64; j += 16) { + { + const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; + #if __CUDA_ARCH__ >= 1000 + #pragma unroll + for (int _ps = 0; _ps < 8; _ps++) + mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_5[j])[_ps], _prescale2_1); + #else + #pragma unroll + for (int _ps = 0; _ps < 16; _ps++) + _tmem_load_5[j + _ps] *= inv_sum * output_scale; + #endif + { + __nv_bfloat162 _pk0 = __floats2bfloat162_rn(_tmem_load_5[j + 0], _tmem_load_5[j + 1]); + unsigned _pk_u0 = *reinterpret_cast(&_pk0); + __nv_bfloat162 _pk1 = __floats2bfloat162_rn(_tmem_load_5[j + 2], _tmem_load_5[j + 3]); + unsigned _pk_u1 = *reinterpret_cast(&_pk1); + __nv_bfloat162 _pk2 = __floats2bfloat162_rn(_tmem_load_5[j + 4], _tmem_load_5[j + 5]); + unsigned _pk_u2 = *reinterpret_cast(&_pk2); + __nv_bfloat162 _pk3 = __floats2bfloat162_rn(_tmem_load_5[j + 6], _tmem_load_5[j + 7]); + unsigned _pk_u3 = *reinterpret_cast(&_pk3); + __nv_bfloat162 _pk4 = __floats2bfloat162_rn(_tmem_load_5[j + 8], _tmem_load_5[j + 9]); + unsigned _pk_u4 = *reinterpret_cast(&_pk4); + __nv_bfloat162 _pk5 = __floats2bfloat162_rn(_tmem_load_5[j + 10], _tmem_load_5[j + 11]); + unsigned _pk_u5 = *reinterpret_cast(&_pk5); + __nv_bfloat162 _pk6 = __floats2bfloat162_rn(_tmem_load_5[j + 12], _tmem_load_5[j + 13]); + unsigned _pk_u6 = *reinterpret_cast(&_pk6); + __nv_bfloat162 _pk7 = __floats2bfloat162_rn(_tmem_load_5[j + 14], _tmem_load_5[j + 15]); + unsigned _pk_u7 = *reinterpret_cast(&_pk7); + asm volatile( + "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"((void*)(&((__nv_bfloat16*)(O + (gmem_base + j)))[0])), "r"(_pk_u0), "r"(_pk_u1), "r"(_pk_u2), "r"(_pk_u3), "r"(_pk_u4), "r"(_pk_u5), "r"(_pk_u6), "r"(_pk_u7) : "memory"); + } + } + } + } + } + } + correction_tile_cursor = correction_tile_cursor + num_kv_tiles_2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_dealloc_addr); + } + } + // ---- Role: mma_warp ---- + if (warp == 8) { + { // mma_warp_main + const int wg2_dummy_1 = 0; + int all_num_kv_tiles_4 = (sparse_topk + 128 - 1) / 128; + unsigned int mma_k_stage = 0; + unsigned int mma_v_stage = 0; + unsigned int mma_kv_stage = 0; + int mma_tile_cursor = 0; + unsigned int _phase_s_seeded_0 = 0; + unsigned int _phase_q_full_0_3 = 0; + unsigned int _phase_q_pair_ready_0 = 0; + unsigned int _phase_k_full = 0; + unsigned int _phase_k_full_1 = 0; + unsigned int _phase_kv_pair_ready = 0; + unsigned int _phase_v_full = 0; + #pragma unroll + for (unsigned int work_idx_3 = 0; work_idx_3 < 1; work_idx_3++) { + int split_idx_4 = 0; + { + split_idx_4 = 0; + } + int tiles_per_split_4 = all_num_kv_tiles_4 + 1 - 1; + int first_tile_4 = split_idx_4 * tiles_per_split_4; + int num_kv_tiles_3 = tiles_per_split_4; + { + if (cta_rank == 0) { + mbarrier_wait(q_full_addr, _phase_q_full_0_3); + _phase_q_full_0_3 ^= 1; + } + } + int first_pv = 1; + #pragma unroll 1 + for (int tile_3 = 0; tile_3 < num_kv_tiles_3; tile_3++) { + int pipeline_tile_2 = mma_tile_cursor + tile_3; + int phase_2 = pipeline_tile_2 & 1; + int score_col = ((phase_2 != 0) ? 128 : 0); + { + if (cta_rank == 0) { + mbarrier_wait(k_full_addr + (mma_k_stage) * 8, _phase_k_full); + int _mma_a_lo_0 = (((smem_q_full_addr) >> 4) & 0x3FFF) + (0) * 2048; + int _mma_b_lo_0 = (((smem_k_full_addr) >> 4) & 0x3FFF) + (mma_k_stage) * 2048; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136314896;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (score_col))), "r"(0)); + int _mma_a_lo_1 = (((smem_q_full_addr + 8192) >> 4) & 0x3FFF) + (0) * 2048; + int _mma_b_lo_1 = (((smem_k_full_addr + 8192) >> 4) & 0x3FFF) + (mma_k_stage) * 2048; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136314896;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + int _mma_a_lo_2 = (((smem_q_full_addr + 16384) >> 4) & 0x3FFF) + (0) * 2048; + int _mma_b_lo_2 = (((smem_k_full_addr + 16384) >> 4) & 0x3FFF) + (mma_k_stage) * 2048; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136314896;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + int _mma_a_lo_3 = (((smem_q_full_addr + 24576) >> 4) & 0x3FFF) + (0) * 2048; + int _mma_b_lo_3 = (((smem_k_full_addr + 24576) >> 4) & 0x3FFF) + (mma_k_stage) * 2048; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136314896;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + if (cta_rank == 0) { + elect_commit_cg2_multicast(s_full_addr + (phase_2) * 8, (uint16_t)(3)); + } + if (cta_rank == 0) { + elect_commit_cg2_multicast(k_empty_addr + (mma_k_stage) * 8, (uint16_t)(3)); + } + } + mma_k_stage += 1; + if (mma_k_stage == 3) { mma_k_stage = 0; _phase_k_full ^= 1; } + } + if (tile_3 > 0) { + int prev_pipeline_tile = pipeline_tile_2 - 1; + int prev_phase = prev_pipeline_tile & 1; + int pv_wait_phase = prev_pipeline_tile >> 1 & 1; + mbarrier_wait(p_full_addr + (prev_phase) * 8, pv_wait_phase); + mbarrier_wait(corr_done_addr + (prev_phase) * 8, pv_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + { + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cluster.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(pv_pair_ready_addr + (unsigned int)(prev_phase * 8)), "r"(0) : "memory"); + if (cta_rank == 0) { + mbarrier_wait_cluster_hint(pv_pair_ready_addr + (prev_phase) * 8, pv_wait_phase, 10000000); + } + } + int p_col = ((prev_phase != 0) ? 224 : 96); + { + if (cta_rank == 0) { + mbarrier_wait(v_full_addr + (mma_v_stage) * 8, _phase_v_full); + int _mma_a_lo_6 = (((smem_p_fp8_addr) >> 4) & 0x3FFF) + (prev_phase) * 512; + int _mma_b_lo_6 = ((((smem_v_full_addr) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138477584;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_6), "r"(_mma_b_lo_6), "r"((tmem_tmem_scratch + (256))), "r"(((first_pv) ? 0 : 1))); + int _mma_a_lo_7 = (((smem_p_fp8_addr) >> 4) & 0x3FFF) + (prev_phase) * 512; + int _mma_b_lo_7 = ((((smem_v_full_addr + 16384) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138477584;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_7), "r"(_mma_b_lo_7), "r"((tmem_tmem_scratch + (384))), "r"(((first_pv) ? 0 : 1))); + if (cta_rank == 0) { + elect_commit_cg2_multicast(v_empty_addr + (mma_v_stage) * 8, (uint16_t)(3)); + } + } + mma_v_stage += 1; + if (mma_v_stage == 2) { mma_v_stage = 0; _phase_v_full ^= 1; } + } + first_pv = 0; + if (cta_rank == 0) { + elect_commit_cg2_multicast(pv_done_addr, (uint16_t)(3)); + } + } + } + int last_pipeline_tile = mma_tile_cursor + num_kv_tiles_3 - 1; + int last_phase = last_pipeline_tile & 1; + int drain_wait_phase = last_pipeline_tile >> 1 & 1; + mbarrier_wait(p_full_addr + (last_phase) * 8, drain_wait_phase); + mbarrier_wait(corr_done_addr + (last_phase) * 8, drain_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + { + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cluster.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(pv_pair_ready_addr + (unsigned int)(last_phase * 8)), "r"(0) : "memory"); + if (cta_rank == 0) { + mbarrier_wait_cluster_hint(pv_pair_ready_addr + (last_phase) * 8, drain_wait_phase, 10000000); + } + } + int p_col_last = ((last_phase != 0) ? 224 : 96); + { + if (cta_rank == 0) { + mbarrier_wait(v_full_addr + (mma_v_stage) * 8, _phase_v_full); + int _mma_a_lo_10 = (((smem_p_fp8_addr) >> 4) & 0x3FFF) + (last_phase) * 512; + int _mma_b_lo_10 = ((((smem_v_full_addr) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138477584;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_10), "r"(_mma_b_lo_10), "r"((tmem_tmem_scratch + (256))), "r"(((first_pv) ? 0 : 1))); + int _mma_a_lo_11 = (((smem_p_fp8_addr) >> 4) & 0x3FFF) + (last_phase) * 512; + int _mma_b_lo_11 = ((((smem_v_full_addr + 16384) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138477584;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_11), "r"(_mma_b_lo_11), "r"((tmem_tmem_scratch + (384))), "r"(((first_pv) ? 0 : 1))); + if (cta_rank == 0) { + elect_commit_cg2_multicast(v_empty_addr + (mma_v_stage) * 8, (uint16_t)(3)); + } + } + mma_v_stage += 1; + if (mma_v_stage == 2) { mma_v_stage = 0; _phase_v_full ^= 1; } + } + if (cta_rank == 0) { + elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + } + if (cta_rank == 0) { + elect_commit_cg2_multicast(o_done_addr, (uint16_t)(3)); + } + mma_tile_cursor = mma_tile_cursor + num_kv_tiles_3; + } + mbarrier_arrive(tmem_dealloc_addr); + unsigned int _phase_tmem_dealloc_0 = 0; + mbarrier_wait(tmem_dealloc_addr, _phase_tmem_dealloc_0); + _phase_tmem_dealloc_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + { + int tmem_dealloc_peer_rank = bid % 2 ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cluster.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_dealloc_peer_addr), "r"(tmem_dealloc_peer_rank) : "memory"); + mbarrier_wait(tmem_dealloc_peer_addr, 0); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + } + } + // ---- Role: empty1 ---- + if (warp >= 10 && warp <= 11) { + { // empty1_main + const int wg2_dummy_2 = 0; + unsigned int _phase_q_full_0_4 = 0; + } + } + + // Cleanup +} + +} // extern "C" diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_98db1f55c8c129e16005_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_98db1f55c8c129e16005_binding.cu new file mode 100644 index 00000000000..0f76ba73e68 --- /dev/null +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_98db1f55c8c129e16005_binding.cu @@ -0,0 +1,461 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +// Generated by the Cake source exporter — do not edit. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_98db1f55c8c129e16005'. +#include +#include +#include +#include + +#include "tvm_ffi_utils.h" + +#include +#include +#include +#include +#include +#include +#include + +extern "C" __global__ void kernel_cake_dsv4_98db1f55c8c129e16005(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items); + + +namespace cake_host_shim_9d011e971631673b { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t CakeDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +// 4D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + int64_t d3 = t.size(t.ndim() - 3); + TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + TVM_FFI_CHECK(d1 % 128 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 128; + uint64_t global_dim[4] = {(uint64_t)(128), (uint64_t)(d2), (uint64_t)((d1 / 128)), (uint64_t)(d3)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (128, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 128; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + (uint64_t)((carrier_stride_1 * 8) / 8), + (uint64_t)((carrier_stride_2 * 8) / 8), + }; + uint32_t box_dim[4] = {128u, 64u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_swa_kv' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_swa_kv' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; + TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (128, 1) exceeds resolved global dims for 'tmap_swa_kv'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + }; + uint32_t box_dim[2] = {128u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_compressed_kv' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_compressed_kv(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_compressed_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_compressed_kv' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_compressed_kv' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_compressed_kv"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_compressed_kv' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved a non-positive global dim"; + TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (128, 1) exceeds resolved global dims for 'tmap_compressed_kv'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + }; + uint32_t box_dim[2] = {128u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_kv') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_compressed_kv, TensorView arg_O, TensorView arg_partial_lse, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_num_query_tokens, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 1, 8, 1); + CheckContiguous(arg_tmap_q, "tmap_q"); + CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); + CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 1, 8, 1); + CheckCudaTensor(arg_tmap_compressed_kv, "tmap_compressed_kv"); + CheckDtype(arg_tmap_compressed_kv, "tmap_compressed_kv", 1, 8, 1); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_partial_lse, "partial_lse"); + CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); + CheckContiguous(arg_partial_lse, "partial_lse"); + CheckCudaTensor(arg_swa_indices, "swa_indices"); + CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); + CheckContiguous(arg_swa_indices, "swa_indices"); + CheckCudaTensor(arg_compressed_indices, "compressed_indices"); + CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); + CheckContiguous(arg_compressed_indices, "compressed_indices"); + CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); + CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); + CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); + CheckContiguous(arg_bmm1_scale, "bmm1_scale"); + CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); + CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); + CheckContiguous(arg_bmm2_scale, "bmm2_scale"); + TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) + << "scalar 'num_heads' value " << arg_num_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) + << "scalar 'swa_index_stride' value " << arg_swa_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) + << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) + << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_query_tokens >= -2147483648LL && arg_num_query_tokens <= 2147483647LL, ValueError) + << "scalar 'num_query_tokens' value " << arg_num_query_tokens + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) + << "scalar 'sparse_topk' value " << arg_sparse_topk + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) + << "scalar 'has_sinks' value " << arg_has_sinks + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); + CheckSameCudaDevice(arg_tmap_compressed_kv, arg_tmap_q, "tmap_compressed_kv", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_partial_lse, arg_tmap_q, "partial_lse", "tmap_q"); + CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); + CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); + CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); + CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") must be divisible by cluster dims (2, 1, 1)"; + + cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); + CUtensorMap p_tmap_compressed_kv = EncodeTma_tmap_compressed_kv(arg_tmap_compressed_kv); + __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); + float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); + int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); + int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); + int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); + float* p_sinks = static_cast(arg_sinks.data_ptr()); + float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); + float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); + int32_t v_num_heads = (int32_t)arg_num_heads; + int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; + int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; + int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; + int32_t v_num_query_tokens = (int32_t)arg_num_query_tokens; + int32_t v_sparse_topk = (int32_t)arg_sparse_topk; + int32_t v_has_sinks = (int32_t)arg_has_sinks; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_tmap_compressed_kv, &p_O, &p_partial_lse, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_num_query_tokens, &v_sparse_topk, &v_has_sinks, &v_total_work_items}; + + static std::mutex smem_mu; + static auto* smem_status_by_device = + new std::unordered_map(); + cudaError_t smem_status = cudaSuccess; + { + std::lock_guard lock(smem_mu); + auto it = smem_status_by_device->find(dev.device_id); + if (it == smem_status_by_device->end()) { + int current_device = -1; + smem_status = cudaGetDevice(¤t_device); + if (smem_status == cudaSuccess && current_device != dev.device_id) { + smem_status = cudaErrorInvalidDevice; + } + if (smem_status == cudaSuccess) { + smem_status = cudaFuncSetAttribute( + reinterpret_cast(kernel_cake_dsv4_98db1f55c8c129e16005), + cudaFuncAttributeMaxDynamicSharedMemorySize, + 190976); + } + smem_status_by_device->emplace(dev.device_id, smem_status); + } else { + smem_status = it->second; + } + } + TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) + << "cudaFuncSetAttribute for kernel_cake_dsv4_98db1f55c8c129e16005 failed: " + << cudaGetErrorString(smem_status); + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(512u, 1u, 1u); + + // Extended direct-source launch for a linked CUDA kernel symbol. + cudaLaunchAttribute attrs[2]{}; + int n = 0; + attrs[n].id = cudaLaunchAttributeClusterDimension; + attrs[n].val.clusterDim.x = 2u; + attrs[n].val.clusterDim.y = 1u; + attrs[n].val.clusterDim.z = 1u; + ++n; + attrs[n].id = cudaLaunchAttributeClusterSchedulingPolicyPreference; + attrs[n].val.clusterSchedulingPolicyPreference = cudaClusterSchedulingPolicySpread; + ++n; + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes = 190976u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; + cudaError_t launch_status = cudaLaunchKernelExC( + &config, reinterpret_cast(kernel_cake_dsv4_98db1f55c8c129e16005), kargs); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernelExC for kernel_cake_dsv4_98db1f55c8c129e16005 failed: " + << cudaGetErrorString(launch_status); +} + +} // namespace cake_host_shim_9d011e971631673b + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_9d011e971631673b::Run); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_98db1f55c8c129e16005_kernel.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_98db1f55c8c129e16005_kernel.cu new file mode 100644 index 00000000000..6364f2a635d --- /dev/null +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_98db1f55c8c129e16005_kernel.cu @@ -0,0 +1,2835 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +#define CAKE_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_SCRATCH_OFFSET 0 +#define NUM_K_PIPE_STAGES 2 +#define NUM_V_PIPE_STAGES 2 +#define NUM_KV_PIPE_STAGES 1 +#define NUM_INDEX_PIPE_STAGES 6 +#define SMEM_SMEM_Q_OFF 1024 +#define SMEM_SMEM_Q_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_STRIDE 8192 +#define SMEM_SMEM_KV_OFF 33792 +#define SMEM_SMEM_KV_STAGE_BYTES 16384 +#define SMEM_SMEM_KV_STRIDE 16384 +#define SMEM_SMEM_V_OFF 33792 +#define SMEM_SMEM_V_STAGE_BYTES 16384 +#define SMEM_SMEM_V_STRIDE 16384 +#define SMEM_SMEM_Q_FULL_OFF 1024 +#define SMEM_SMEM_Q_FULL_STAGE_BYTES 32768 +#define SMEM_SMEM_Q_FULL_STRIDE 32768 +#define SMEM_SMEM_K_FULL_OFF 33792 +#define SMEM_SMEM_K_FULL_STAGE_BYTES 65536 +#define SMEM_SMEM_K_FULL_STRIDE 65536 +#define SMEM_SMEM_V_FULL_OFF 33792 +#define SMEM_SMEM_V_FULL_STAGE_BYTES 65536 +#define SMEM_SMEM_V_FULL_STRIDE 65536 +#define SMEM_SMEM_STATS_MAX_OFF 164864 +#define SMEM_SMEM_STATS_MAX_STAGE_BYTES 1024 +#define SMEM_SMEM_STATS_MAX_STRIDE 1024 +#define SMEM_SMEM_STATS_SUM_OFF 165888 +#define SMEM_SMEM_STATS_SUM_STAGE_BYTES 512 +#define SMEM_SMEM_STATS_SUM_STRIDE 512 +#define SMEM_SMEM_STATS_FINAL_MAX_OFF 166400 +#define SMEM_SMEM_STATS_FINAL_MAX_STAGE_BYTES 512 +#define SMEM_SMEM_STATS_FINAL_MAX_STRIDE 512 +#define SMEM_SMEM_SOFTMAX_WARP_PAIR_EXCHANGE_OFF 166912 +#define SMEM_SMEM_SOFTMAX_WARP_PAIR_EXCHANGE_STAGE_BYTES 1024 +#define SMEM_SMEM_SOFTMAX_WARP_PAIR_EXCHANGE_STRIDE 1024 +#define SMEM_SMEM_CORR_WARP_PAIR_EXCHANGE_OFF 167936 +#define SMEM_SMEM_CORR_WARP_PAIR_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_CORR_WARP_PAIR_EXCHANGE_STRIDE 512 +#define SMEM_SMEM_SPARSE_INDICES_OFF 168448 +#define SMEM_SMEM_SPARSE_INDICES_STAGE_BYTES 1024 +#define SMEM_SMEM_SPARSE_INDICES_STRIDE 1024 +#define SMEM_SMEM_P_FP8_OFF 174592 +#define SMEM_SMEM_P_FP8_STAGE_BYTES 8192 +#define SMEM_SMEM_P_FP8_STRIDE 8192 +#define SMEM_SMEM_MERGE_RECV_OFF 33792 +#define SMEM_SMEM_MERGE_RECV_STAGE_BYTES 65536 +#define SMEM_SMEM_MERGE_RECV_STRIDE 65536 +#define SMEM_SMEM_MERGE_SEND_OFF 99328 +#define SMEM_SMEM_MERGE_SEND_STAGE_BYTES 65536 +#define SMEM_SMEM_MERGE_SEND_STRIDE 65536 +#define SMEM_SMEM_MERGE_STATS_RECV_OFF 174592 +#define SMEM_SMEM_MERGE_STATS_RECV_STAGE_BYTES 512 +#define SMEM_SMEM_MERGE_STATS_RECV_STRIDE 512 +#define SMEM_SMEM_MERGE_STATS_SEND_OFF 182784 +#define SMEM_SMEM_MERGE_STATS_SEND_STAGE_BYTES 512 +#define SMEM_SMEM_MERGE_STATS_SEND_STRIDE 512 +#define SMEM_TOTAL 190976 +#define THREADS 512 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count) : "memory"); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f8f6f4( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x4(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x4.b32" + " {%0, %1, %2, %3}, [%4];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); +} + + +__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ float2 ex2_emulation_f32x2_value(float2 value) { + const float c0 = 1.0f, c1 = 0.695146143436431884765625f; + const float c2 = 0.227564394474029541015625f, c3 = 0.077119089663028717041015625f; + const float magic = 12582912.0f; + float x0 = max_noftz(value.x, -127.0f), x1 = max_noftz(value.y, -127.0f); + float2 xc2 = make_float2(x0, x1), magic2 = make_float2(magic, magic); + float2 xr2; + asm("add.rm.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xr2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&magic2)); + float2 c3_2 = make_float2(c3, c3), c2_2 = make_float2(c2, c2); + float2 c1_2 = make_float2(c1, c1), c0_2 = make_float2(c0, c0); + float2 xrb2, xfrac2; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xrb2) + : "l"(*(unsigned long long*)&xr2), "l"(*(unsigned long long*)&magic2)); + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xfrac2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&xrb2)); + float2 poly2; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&c3_2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c2_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c1_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c0_2)); + int x0r_i, x1r_i, p0_i, p1_i; + asm("mov.b64 {%0, %1}, %2;" : "=r"(x0r_i), "=r"(x1r_i) : "l"(*(unsigned long long*)&xr2)); + asm("mov.b64 {%0, %1}, %2;" : "=r"(p0_i), "=r"(p1_i) : "l"(*(unsigned long long*)&poly2)); + float r0, r1; + asm("mov.b32 %0, %1;" : "=f"(r0) : "r"((x0r_i << 23) + p0_i)); + asm("mov.b32 %0, %1;" : "=f"(r1) : "r"((x1r_i << 23) + p1_i)); + return make_float2(r0, r1); +} + +__device__ __forceinline__ void ex2_emulation_f32x2(float* x0_ptr, float* x1_ptr) { + float2 result = ex2_emulation_f32x2_value(make_float2(*x0_ptr, *x1_ptr)); + *x0_ptr = result.x; *x1_ptr = result.y; +} + +__device__ __forceinline__ void softmax_frag_exp2_cast( + float* sv, uint32_t* pv, int use_emu) +{ + #pragma unroll + for (int j = 0; j < 16; j++) { + if (use_emu && j >= 12) + ex2_emulation_f32x2(&sv[j*2], &sv[j*2+1]); + else { + sv[j*2] = approx_exp2(sv[j*2]); + sv[j*2+1] = approx_exp2(sv[j*2+1]); + } + } + #pragma unroll + for (int j = 0; j < 16; j++) { + __nv_bfloat162 bf = __float22bfloat162_rn({sv[j*2], sv[j*2+1]}); + pv[j] = reinterpret_cast(bf); + } +} + + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_2d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3}], [%4];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr) { + // Canonical .shared::cta form for non-multicast gather4, matching + // trtllm-gen / cuda_ptx and the PTX ISA qualifier order + // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, + // .shared::cluster is reserved for the multicast variant (ctaMask). + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_st_x8_u32(int addr, const uint32_t* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x8.b32" + " [%0], {%1,%2,%3,%4,%5,%6,%7,%8};" + :: "r"(addr), + "r"(src[0]), "r"(src[1]), "r"(src[2]), "r"(src[3]), + "r"(src[4]), "r"(src[5]), "r"(src[6]), "r"(src[7])); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(512, 1) __cluster_dims__(2,1,1) void +kernel_cake_dsv4_98db1f55c8c129e16005(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int mbar_base = smem; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define k_full_addr (mbar_base + 16) + #define k_empty_addr (mbar_base + 32) + #define v_full_addr (mbar_base + 48) + #define v_empty_addr (mbar_base + 64) + #define index_full_addr (mbar_base + 80) + #define index_empty_addr (mbar_base + 128) + #define s_full_addr (mbar_base + 176) + #define p_full_addr (mbar_base + 192) + #define corr_done_addr (mbar_base + 208) + #define stats_addr (mbar_base + 224) + #define sum_ready_addr (mbar_base + 240) + #define o_done_addr (mbar_base + 248) + #define pv_done_addr (mbar_base + 256) + #define s_seeded_addr (mbar_base + 264) + #define q_pair_ready_addr (mbar_base + 272) + #define kv_pair_ready_addr (mbar_base + 280) + #define pv_pair_ready_addr (mbar_base + 288) + #define tmem_dealloc_addr (mbar_base + 304) + #define tmem_dealloc_peer_addr (mbar_base + 312) + #define merge_ready_addr (mbar_base + 320) + #define merge_full_addr (mbar_base + 328) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 2; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + + // Kernel setup ops + uint8_t* smem_q = reinterpret_cast(smem_raw + 1024); + const int smem_q_addr = smem + 1024; + uint8_t* smem_kv = reinterpret_cast(smem_raw + 33792); + const int smem_kv_addr = smem + 33792; + uint8_t* smem_v = reinterpret_cast(smem_raw + 33792); + const int smem_v_addr = smem + 33792; + uint8_t* smem_q_full = reinterpret_cast(smem_raw + 1024); + const int smem_q_full_addr = smem + 1024; + uint8_t* smem_k_full = reinterpret_cast(smem_raw + 33792); + const int smem_k_full_addr = smem + 33792; + uint8_t* smem_v_full = reinterpret_cast(smem_raw + 33792); + const int smem_v_full_addr = smem + 33792; + float* smem_stats_max = reinterpret_cast(smem_raw + 164864); + const int smem_stats_max_addr = smem + 164864; + float* smem_stats_sum = reinterpret_cast(smem_raw + 165888); + const int smem_stats_sum_addr = smem + 165888; + float* smem_stats_final_max = reinterpret_cast(smem_raw + 166400); + const int smem_stats_final_max_addr = smem + 166400; + float* smem_softmax_warp_pair_exchange = reinterpret_cast(smem_raw + 166912); + const int smem_softmax_warp_pair_exchange_addr = smem + 166912; + float* smem_corr_warp_pair_exchange = reinterpret_cast(smem_raw + 167936); + const int smem_corr_warp_pair_exchange_addr = smem + 167936; + int* smem_sparse_indices = reinterpret_cast(smem_raw + 168448); + const int smem_sparse_indices_addr = smem + 168448; + uint8_t* smem_p_fp8 = reinterpret_cast(smem_raw + 174592); + const int smem_p_fp8_addr = smem + 174592; + float* smem_merge_recv = reinterpret_cast(smem_raw + 33792); + const int smem_merge_recv_addr = smem + 33792; + float* smem_merge_send = reinterpret_cast(smem_raw + 99328); + const int smem_merge_send_addr = smem + 99328; + float* smem_merge_stats_recv = reinterpret_cast(smem_raw + 174592); + const int smem_merge_stats_recv_addr = smem + 174592; + float* smem_merge_stats_send = reinterpret_cast(smem_raw + 182784); + const int smem_merge_stats_send_addr = smem + 182784; + { + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + } + + // Mbarrier init (23 pipeline groups, 0 ordered-sequence groups, 42 barriers) + // Mbarriers at smem_raw[0..336) + + if (warp == 12) { + uint32_t leader = elect_sync(); + if (leader) { + // q_full: 1 barriers, init_count=1 + mbarrier_init(smem + 0, 1); + // q_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 8, 1); + // --- pipeline 'k_pipe' --- + // k_full: 2 barriers, init_count=1 + mbarrier_init(smem + 16, 1); + mbarrier_init(smem + 24, 1); + // k_empty: 2 barriers, init_count=1 + mbarrier_init(smem + 32, 1); + mbarrier_init(smem + 40, 1); + // --- pipeline 'v_pipe' --- + // v_full: 2 barriers, init_count=2 + mbarrier_init(smem + 48, 2); + mbarrier_init(smem + 56, 2); + // v_empty: 2 barriers, init_count=1 + mbarrier_init(smem + 64, 1); + mbarrier_init(smem + 72, 1); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + if (warp == 9) { + uint32_t leader = elect_sync(); + if (leader) { + // --- pipeline 'index_pipe' --- + // index_full: 6 barriers, init_count=32 + mbarrier_init(smem + 80, 32); + mbarrier_init(smem + 88, 32); + mbarrier_init(smem + 96, 32); + mbarrier_init(smem + 104, 32); + mbarrier_init(smem + 112, 32); + mbarrier_init(smem + 120, 32); + // index_empty: 6 barriers, init_count=192 + mbarrier_init(smem + 128, 192); + mbarrier_init(smem + 136, 192); + mbarrier_init(smem + 144, 192); + mbarrier_init(smem + 152, 192); + mbarrier_init(smem + 160, 192); + mbarrier_init(smem + 168, 192); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // s_full: stages (0,), init_count=1 + mbarrier_init(smem + 176, 1); + // p_full: stages (0,), init_count=64 + mbarrier_init(smem + 192, 64); + // o_done: 1 barriers, init_count=1 + mbarrier_init(smem + 248, 1); + // pv_done: 1 barriers, init_count=1 + mbarrier_init(smem + 256, 1); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + if (warp == 1) { + uint32_t leader = elect_sync(); + if (leader) { + // s_full: stages (1,), init_count=1 + mbarrier_init(smem + 184, 1); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + if (warp == 2) { + uint32_t leader = elect_sync(); + if (leader) { + // p_full: stages (1,), init_count=64 + mbarrier_init(smem + 200, 64); + // s_seeded: 1 barriers, init_count=256 + mbarrier_init(smem + 264, 256); + // q_pair_ready: 1 barriers, init_count=64 + mbarrier_init(smem + 272, 64); + // --- pipeline 'kv_pipe' --- + // kv_pair_ready: 1 barriers, init_count=64 + mbarrier_init(smem + 280, 64); + // pv_pair_ready: stages (0,), init_count=64 + mbarrier_init(smem + 288, 64); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + if (warp == 4) { + uint32_t leader = elect_sync(); + if (leader) { + // corr_done: stages (0,), init_count=64 + mbarrier_init(smem + 208, 64); + // stats: stages (0,), init_count=64 + mbarrier_init(smem + 224, 64); + // sum_ready: 1 barriers, init_count=64 + mbarrier_init(smem + 240, 64); + // pv_pair_ready: stages (1,), init_count=64 + mbarrier_init(smem + 296, 64); + // tmem_dealloc: 1 barriers, init_count=288 + mbarrier_init(smem + 304, 288); + // merge_ready: 1 barriers, init_count=64 + mbarrier_init(smem + 320, 64); + // merge_full: 1 barriers, init_count=1 + mbarrier_init(smem + 328, 1); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + if (warp == 6) { + uint32_t leader = elect_sync(); + if (leader) { + // corr_done: stages (1,), init_count=64 + mbarrier_init(smem + 216, 64); + // stats: stages (1,), init_count=64 + mbarrier_init(smem + 232, 64); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + if (warp == 5) { + uint32_t leader = elect_sync(); + if (leader) { + // tmem_dealloc_peer: 1 barriers, init_count=32 + mbarrier_init(smem + 312, 32); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncwarp(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 336); + if (warp == 0) { + int _tmem_hold = smem + 336; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } + + __syncthreads(); + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_tmem_scratch = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Inc phase consumes the registers released above. + if (warp >= 0 && warp <= 11) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 136;"); + } + + // ---- Role: index_warp ---- + if (warp == 9) { + { // index_warp_main + const int index_dummy = 0; + unsigned int _phase_index_empty = 1; + unsigned int _phase_q_full_0 = 0; + { + int all_num_kv_tiles = (sparse_topk + 128 - 1) / 128; + unsigned int index_prod_stage = 0; + #pragma unroll + for (unsigned int static_work = 0; static_work < 1; static_work++) { + int split_idx = 0; + int query_idx = bid / 2; + int tiles_per_split = all_num_kv_tiles + 1 - 1; + int first_tile = split_idx * tiles_per_split; + int num_index_passes = (tiles_per_split + 1) / 2; + #pragma unroll 1 + for (int index_pass = 0; index_pass < num_index_passes; index_pass++) { + mbarrier_wait(index_empty_addr + (index_prod_stage) * 8, _phase_index_empty); + int index_stage_base = smem_sparse_indices_addr + index_prod_stage * 1024; + int stage_row_base = index_pass * 256; + #pragma unroll + for (int index_half = 0; index_half < 2; index_half++) { + int index_offset = index_half * 128 + lane * 4; + int tile_col_base = first_tile * 128 + stage_row_base + index_half * 128; + int* row_tile_ptr = ((tile_col_base < 128) ? swa_indices + (query_idx * swa_index_stride) : compressed_indices + (query_idx * compressed_index_stride + (tile_col_base - 128))); + int lane_col = tile_col_base + lane * 4; + if (lane_col + 4 <= sparse_topk) { + asm volatile("cp.async.cg.shared::cta.global.L2::128B [%0], [%1], 16;" + :: "r"(index_stage_base + index_offset * 4), "l"(row_tile_ptr + (lane * 4))); + } else { + int tail_rows[4]; + #pragma unroll + for (int row_i = 0; row_i < 4; row_i++) { + tail_rows[row_i] = -1; + } + #pragma unroll + for (int row_i_1 = 0; row_i_1 < 4; row_i_1++) { + if (lane_col + row_i_1 < sparse_topk) { + tail_rows[row_i_1] = row_tile_ptr[lane * 4 + row_i_1]; + } + } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(smem_sparse_indices_addr + ((unsigned int)index_stage_base - smem_sparse_indices_addr + (unsigned int)(index_offset * 4))), "r"(tail_rows[0]), "r"(tail_rows[1]), "r"(tail_rows[2]), "r"(tail_rows[3]) : "memory"); + } + } + asm volatile("cp.async.commit_group;"); + asm volatile("cp.async.wait_group 0;"); + mbarrier_arrive(index_full_addr + (index_prod_stage) * 8); + index_prod_stage += 1; + if (index_prod_stage == 6) { index_prod_stage = 0; _phase_index_empty ^= 1; } + } + } + } + } + } + // ---- Role: load_warp ---- + if (warp >= 12 && warp <= 15) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 96;"); + { // load_warp_main + const int wg2_dummy = 0; + const int load_warp_rank = warp - ((1) ? 12 : 9); + int all_num_kv_tiles_1 = (sparse_topk + 128 - 1) / 128; + unsigned int load_k_stage = 0; + unsigned int load_v_stage = 0; + unsigned int load_kv_stage = 0; + unsigned int load_k_index_stage = 0; + unsigned int load_v_index_stage = 0; + int k_cta_offset = bid % 2 * 64; + unsigned int _phase_q_empty_0 = 1; + unsigned int _phase_index_full = 0; + unsigned int _phase_k_empty = 1; + #pragma unroll + for (unsigned int work_idx = 0; work_idx < 1; work_idx++) { + int split_idx_1 = 0; + int query_idx_1 = bid / 2 >> 1; + int v_chunk = bid / 2 & 1; + { + split_idx_1 = 0; + query_idx_1 = bid / 2; + v_chunk = bid % 2; + } + int tiles_per_split_1 = all_num_kv_tiles_1 + 1 - 1; + int first_tile_1 = split_idx_1 * tiles_per_split_1; + int num_kv_tiles = tiles_per_split_1; + int sparse_extent = num_kv_tiles * 128; + int wt_num_local_tiles = (num_kv_tiles - bid % 2 + 1) / 2; + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + if (load_warp_rank == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(q_full_addr, 32768); + #pragma unroll + for (int q_stage = 0; q_stage < 4; q_stage++) { + tma_4d_gmem2smem(smem_q_full_addr + (unsigned int)(q_stage * 8192), (&tmap_q), 0, 0, q_stage, query_idx_1, q_full_addr); + } + } + } + #pragma unroll 1 + for (int wt_tile = 0; wt_tile < wt_num_local_tiles; wt_tile++) { + int wt_global_tile = first_tile_1 + wt_tile * 2 + bid % 2; + mbarrier_wait(index_full_addr + (load_k_index_stage) * 8, _phase_index_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + int wt_index_base = smem_sparse_indices_addr + load_k_index_stage * 1024; + mbarrier_wait(k_empty_addr + (load_k_stage) * 8, _phase_k_empty); + if (load_warp_rank == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(k_full_addr + (load_k_stage) * 8, 65536); + } + } + int wt_dst = smem_k_full_addr + load_k_stage * 65536 + (unsigned int)(load_warp_rank * 16384); + int wt_group = lane; + int wt_group_offset = bid % 2 * 128 + wt_group * 4; + int wt_rows[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&wt_rows[0])), "=r"(*reinterpret_cast(&wt_rows[(0) + 1])), "=r"(*reinterpret_cast(&wt_rows[(0) + 2])), "=r"(*reinterpret_cast(&wt_rows[(0) + 3])) + : "r"(wt_index_base + wt_group_offset * 4)); + int wt_row0 = ((wt_rows[0] >= 0) ? wt_rows[0] : 0); + int wt_row1 = ((wt_rows[1] >= 0) ? wt_rows[1] : 0); + int wt_row2 = ((wt_rows[2] >= 0) ? wt_rows[2] : 0); + int wt_row3 = ((wt_rows[3] >= 0) ? wt_rows[3] : 0); + if (wt_global_tile == 0) { + tma_gather4_gmem2smem(wt_dst + wt_group * 512, (&tmap_swa_kv), load_warp_rank * 128, wt_row0, wt_row1, wt_row2, wt_row3, k_full_addr + (load_k_stage) * 8); + } else { + tma_gather4_gmem2smem(wt_dst + wt_group * 512, (&tmap_compressed_kv), load_warp_rank * 128, wt_row0, wt_row1, wt_row2, wt_row3, k_full_addr + (load_k_stage) * 8); + } + load_k_stage += 1; + if (load_k_stage == 2) { load_k_stage = 0; _phase_k_empty ^= 1; } + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(index_empty_addr + (load_k_index_stage) * 8); + load_k_index_stage += 1; + if (load_k_index_stage == 6) { load_k_index_stage = 0; _phase_index_full ^= 1; } + } + } + mbarrier_arrive(tmem_dealloc_addr); + } + } + // ---- Role: softmax_wg ---- + if (warp == 0 || warp == 2) { + { // softmax_wg_main + float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; + const int wg_dummy_inc = 0; + int all_num_kv_tiles_2 = (sparse_topk + 128 - 1) / 128; + const int tmem_row_base = ((1) ? warp % 2 * 32 : warp % 4 * 32); + const int tmem_score_row_base = ((1) ? (int)(warp % 4 * 32) : tmem_row_base); + const int n_half = ((1) ? (int)(warp % 4 / 2) : 0); + const int my_row = tmem_row_base + lane; + const int stats_row = n_half * 64 + my_row; + int softmax_tile_cursor = 0; + unsigned int softmax_index_stage = 0; + unsigned int _phase_q_full_0_1 = 0; + #pragma unroll + for (unsigned int work_idx_1 = 0; work_idx_1 < 1; work_idx_1++) { + int split_idx_2 = 0; + int query_idx_2 = bid / 2 >> 1; + { + split_idx_2 = 0; + query_idx_2 = bid / 2; + } + int tiles_per_split_2 = all_num_kv_tiles_2 + 1 - 1; + int first_tile_2 = split_idx_2 * tiles_per_split_2; + int num_kv_tiles_1 = tiles_per_split_2; + int _max_0 = ((sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) : (0)); + int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); + int active_topk = _min_0; + float row_max_val = -CAKE_INF; + float row_sum_val = 0.0f; + int sink_head = ((1) ? my_row : my_row); + if (has_sinks != 0 && sink_head < num_heads && split_idx_2 + bid % 2 == 0) { + row_max_val = sinks[sink_head] * 1.4426950408889634f / softmax_scale_log2; + row_sum_val = 1.0f; + } + #pragma unroll 1 + for (int tile = 0; tile < (num_kv_tiles_1 - bid % 2 + 1) / 2; tile++) { + int pipeline_tile = softmax_tile_cursor + tile; + int phase = pipeline_tile & 1; + int s_wait_phase = pipeline_tile >> 1 & 1; + mbarrier_wait(s_full_addr + (phase) * 8, s_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int s_off = ((phase != 0) ? 128 : 0); + int s_base = taddr + (unsigned int)s_off + (unsigned int)(tmem_score_row_base << 16); + float new_max = row_max_val; + int valid_sparse_cols = ((active_topk < sparse_topk) ? active_topk : sparse_topk); + valid_sparse_cols = valid_sparse_cols - (first_tile_2 + tile * 2 + bid % 2) * 128 - n_half * 64; + if (valid_sparse_cols < 0) { + valid_sparse_cols = 0; + } + if (valid_sparse_cols > 64) { + valid_sparse_cols = 64; + } + int index_tile_addr = smem_sparse_indices_addr + softmax_index_stage * 1024 + (unsigned int)(bid % 2 * 512); + int index_lane_addr = index_tile_addr + (n_half * 64 + lane) * 4; + int staged_index[2]; + unsigned int invalid_cols[2]; + unsigned int any_invalid = 0; + #pragma unroll + for (int w = 0; w < 2; w++) { + asm volatile("ld.shared.b32 %0, [%1];" : "=r"(*reinterpret_cast(&staged_index[w])) : "r"(index_lane_addr + w * 128)); + unsigned int _vote_0 = __ballot_sync(0xFFFFFFFF, staged_index[w] < 0); + invalid_cols[w] = _vote_0; + any_invalid = any_invalid | invalid_cols[w]; + } + if (tile >= 0) { + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(index_empty_addr + (softmax_index_stage) * 8); + softmax_index_stage = softmax_index_stage + 1; + if (softmax_index_stage == 6) { + softmax_index_stage = 0; + } + } + float _tmem_load_0[4]; + tmem_ld_x4(&_tmem_load_0[0], s_base); + float sv_split[64]; + float half_max = -CAKE_INF; + if (valid_sparse_cols == 64 && any_invalid == 0) { + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 1000) + #error "TmemLoadRed requires tcgen05.ld.red support (sm_103/sm_101-sm_110 family), not sm_100" + #endif + asm volatile( + "tcgen05.ld.red.sync.aligned.32x32b.x64.max.f32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, %64, [%65];" + : "=f"(sv_split[0]), "=f"(sv_split[1]), "=f"(sv_split[2]), "=f"(sv_split[3]), "=f"(sv_split[4]), "=f"(sv_split[5]), "=f"(sv_split[6]), "=f"(sv_split[7]), "=f"(sv_split[8]), "=f"(sv_split[9]), "=f"(sv_split[10]), "=f"(sv_split[11]), "=f"(sv_split[12]), "=f"(sv_split[13]), "=f"(sv_split[14]), "=f"(sv_split[15]), "=f"(sv_split[16]), "=f"(sv_split[17]), "=f"(sv_split[18]), "=f"(sv_split[19]), "=f"(sv_split[20]), "=f"(sv_split[21]), "=f"(sv_split[22]), "=f"(sv_split[23]), "=f"(sv_split[24]), "=f"(sv_split[25]), "=f"(sv_split[26]), "=f"(sv_split[27]), "=f"(sv_split[28]), "=f"(sv_split[29]), "=f"(sv_split[30]), "=f"(sv_split[31]), "=f"(sv_split[32]), "=f"(sv_split[33]), "=f"(sv_split[34]), "=f"(sv_split[35]), "=f"(sv_split[36]), "=f"(sv_split[37]), "=f"(sv_split[38]), "=f"(sv_split[39]), "=f"(sv_split[40]), "=f"(sv_split[41]), "=f"(sv_split[42]), "=f"(sv_split[43]), "=f"(sv_split[44]), "=f"(sv_split[45]), "=f"(sv_split[46]), "=f"(sv_split[47]), "=f"(sv_split[48]), "=f"(sv_split[49]), "=f"(sv_split[50]), "=f"(sv_split[51]), "=f"(sv_split[52]), "=f"(sv_split[53]), "=f"(sv_split[54]), "=f"(sv_split[55]), "=f"(sv_split[56]), "=f"(sv_split[57]), "=f"(sv_split[58]), "=f"(sv_split[59]), "=f"(sv_split[60]), "=f"(sv_split[61]), "=f"(sv_split[62]), "=f"(sv_split[63]), "=f"(half_max) + : "r"(s_base)); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + } else { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv_split[0]), "=f"(sv_split[1]), "=f"(sv_split[2]), "=f"(sv_split[3]), "=f"(sv_split[4]), "=f"(sv_split[5]), "=f"(sv_split[6]), "=f"(sv_split[7]), "=f"(sv_split[8]), "=f"(sv_split[9]), "=f"(sv_split[10]), "=f"(sv_split[11]), "=f"(sv_split[12]), "=f"(sv_split[13]), "=f"(sv_split[14]), "=f"(sv_split[15]), "=f"(sv_split[16]), "=f"(sv_split[17]), "=f"(sv_split[18]), "=f"(sv_split[19]), "=f"(sv_split[20]), "=f"(sv_split[21]), "=f"(sv_split[22]), "=f"(sv_split[23]), "=f"(sv_split[24]), "=f"(sv_split[25]), "=f"(sv_split[26]), "=f"(sv_split[27]), "=f"(sv_split[28]), "=f"(sv_split[29]), "=f"(sv_split[30]), "=f"(sv_split[31]) + : "r"(s_base)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv_split[32]), "=f"(sv_split[33]), "=f"(sv_split[34]), "=f"(sv_split[35]), "=f"(sv_split[36]), "=f"(sv_split[37]), "=f"(sv_split[38]), "=f"(sv_split[39]), "=f"(sv_split[40]), "=f"(sv_split[41]), "=f"(sv_split[42]), "=f"(sv_split[43]), "=f"(sv_split[44]), "=f"(sv_split[45]), "=f"(sv_split[46]), "=f"(sv_split[47]), "=f"(sv_split[48]), "=f"(sv_split[49]), "=f"(sv_split[50]), "=f"(sv_split[51]), "=f"(sv_split[52]), "=f"(sv_split[53]), "=f"(sv_split[54]), "=f"(sv_split[55]), "=f"(sv_split[56]), "=f"(sv_split[57]), "=f"(sv_split[58]), "=f"(sv_split[59]), "=f"(sv_split[60]), "=f"(sv_split[61]), "=f"(sv_split[62]), "=f"(sv_split[63]) + : "r"(s_base + 32)); + uint32_t _slice_lo_mask_0; + { + int _lim_0 = valid_sparse_cols; + if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_0 & (1u << 0))) sv_split[0] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 1))) sv_split[1] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 2))) sv_split[2] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 3))) sv_split[3] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 4))) sv_split[4] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 5))) sv_split[5] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 6))) sv_split[6] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 7))) sv_split[7] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 8))) sv_split[8] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 9))) sv_split[9] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 10))) sv_split[10] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 11))) sv_split[11] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 12))) sv_split[12] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 13))) sv_split[13] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 14))) sv_split[14] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 15))) sv_split[15] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 16))) sv_split[16] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 17))) sv_split[17] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 18))) sv_split[18] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 19))) sv_split[19] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 20))) sv_split[20] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 21))) sv_split[21] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 22))) sv_split[22] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 23))) sv_split[23] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 24))) sv_split[24] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 25))) sv_split[25] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 26))) sv_split[26] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 27))) sv_split[27] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 28))) sv_split[28] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 29))) sv_split[29] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 30))) sv_split[30] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 31))) sv_split[31] = -CAKE_INF; + uint32_t _slice_lo_mask_1; + { + int _lim_1 = valid_sparse_cols - 32; + if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) sv_split[32] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) sv_split[33] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) sv_split[34] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) sv_split[35] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) sv_split[36] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) sv_split[37] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) sv_split[38] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) sv_split[39] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) sv_split[40] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) sv_split[41] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) sv_split[42] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) sv_split[43] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) sv_split[44] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) sv_split[45] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) sv_split[46] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) sv_split[47] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) sv_split[48] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) sv_split[49] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) sv_split[50] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) sv_split[51] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) sv_split[52] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) sv_split[53] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) sv_split[54] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) sv_split[55] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) sv_split[56] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) sv_split[57] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) sv_split[58] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) sv_split[59] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) sv_split[60] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) sv_split[61] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) sv_split[62] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) sv_split[63] = -CAKE_INF; + if (any_invalid != 0) { + #pragma unroll + for (int w_1 = 0; w_1 < 2; w_1++) { + unsigned int mask_word = invalid_cols[w_1]; + sv_split[w_1 * 32] = (((mask_word & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32]); + sv_split[w_1 * 32 + 1] = (((mask_word >> 1 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 1]); + sv_split[w_1 * 32 + 2] = (((mask_word >> 2 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 2]); + sv_split[w_1 * 32 + 3] = (((mask_word >> 3 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 3]); + sv_split[w_1 * 32 + 4] = (((mask_word >> 4 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 4]); + sv_split[w_1 * 32 + 5] = (((mask_word >> 5 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 5]); + sv_split[w_1 * 32 + 6] = (((mask_word >> 6 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 6]); + sv_split[w_1 * 32 + 7] = (((mask_word >> 7 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 7]); + sv_split[w_1 * 32 + 8] = (((mask_word >> 8 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 8]); + sv_split[w_1 * 32 + 9] = (((mask_word >> 9 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 9]); + sv_split[w_1 * 32 + 10] = (((mask_word >> 10 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 10]); + sv_split[w_1 * 32 + 11] = (((mask_word >> 11 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 11]); + sv_split[w_1 * 32 + 12] = (((mask_word >> 12 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 12]); + sv_split[w_1 * 32 + 13] = (((mask_word >> 13 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 13]); + sv_split[w_1 * 32 + 14] = (((mask_word >> 14 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 14]); + sv_split[w_1 * 32 + 15] = (((mask_word >> 15 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 15]); + sv_split[w_1 * 32 + 16] = (((mask_word >> 16 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 16]); + sv_split[w_1 * 32 + 17] = (((mask_word >> 17 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 17]); + sv_split[w_1 * 32 + 18] = (((mask_word >> 18 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 18]); + sv_split[w_1 * 32 + 19] = (((mask_word >> 19 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 19]); + sv_split[w_1 * 32 + 20] = (((mask_word >> 20 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 20]); + sv_split[w_1 * 32 + 21] = (((mask_word >> 21 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 21]); + sv_split[w_1 * 32 + 22] = (((mask_word >> 22 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 22]); + sv_split[w_1 * 32 + 23] = (((mask_word >> 23 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 23]); + sv_split[w_1 * 32 + 24] = (((mask_word >> 24 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 24]); + sv_split[w_1 * 32 + 25] = (((mask_word >> 25 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 25]); + sv_split[w_1 * 32 + 26] = (((mask_word >> 26 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 26]); + sv_split[w_1 * 32 + 27] = (((mask_word >> 27 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 27]); + sv_split[w_1 * 32 + 28] = (((mask_word >> 28 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 28]); + sv_split[w_1 * 32 + 29] = (((mask_word >> 29 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 29]); + sv_split[w_1 * 32 + 30] = (((mask_word >> 30 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 30]); + sv_split[w_1 * 32 + 31] = (((mask_word >> 31 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 31]); + } + } + float2 _reg_reduce_max2_2 = {-CAKE_INF, -CAKE_INF}; + row_max_x32_accum(&sv_split[0], _reg_reduce_max2_2); + row_max_x32_accum(&sv_split[32], _reg_reduce_max2_2); + float sv_split_max = row_max_reduce(_reg_reduce_max2_2); + half_max = sv_split_max; + } + { + float _max_1 = max_noftz(new_max, half_max); + new_max = _max_1; + smem_softmax_warp_pair_exchange[phase * 128 + stats_row] = new_max; + asm volatile("barrier.sync %0, 64;" :: "r"(2 + warp % 2) : "memory"); + float _max_2 = max_noftz(new_max, smem_softmax_warp_pair_exchange[phase * 128 + (stats_row ^ 64)]); + new_max = _max_2; + } + float _fma_0 = __fmaf_rn(row_max_val, softmax_scale_log2, (-new_max) * softmax_scale_log2); + float delta = _fma_0; + float _exp2_0 = approx_exp2(delta); + float exp_delta = _exp2_0; + float acc_scale = ((row_max_val > -CAKE_INF) ? exp_delta : 1.0f); + smem_stats_max[phase * 128 + stats_row] = acc_scale; + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(stats_addr + (phase) * 8); + row_max_val = new_max; + float safe_max = ((new_max == -CAKE_INF) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + float block_sum = 0.0f; + { + { + const float2 _fma_b2_3 = {softmax_scale_log2, softmax_scale_log2}; + const float2 _fma_c2_4 = {-max_scaled, -max_scaled}; + #pragma unroll + for (int _lf = 0; _lf < 32; _lf++) + fma_f32x2_inplace(&reinterpret_cast(sv_split)[_lf], _fma_b2_3, _fma_c2_4); + #pragma unroll + for (int _le = 0; _le < 64; _le++) { + sv_split[_le] = approx_exp2(sv_split[_le]); + } + { + const float2* _reg_reduce_src2_5 = reinterpret_cast(&sv_split[0]); + float2 _reg_reduce_sum2_5_0 = make_float2(0.0f, 0.0f); + float2 _reg_reduce_sum2_5_1 = make_float2(0.0f, 0.0f); + float2 _reg_reduce_sum2_5_2 = make_float2(0.0f, 0.0f); + float2 _reg_reduce_sum2_5_3 = make_float2(0.0f, 0.0f); + #pragma unroll + for (int _rr = 0; _rr < 8; _rr++) { + add_f32x2_inplace(&_reg_reduce_sum2_5_0, _reg_reduce_src2_5[_rr * 4 + 0]); + add_f32x2_inplace(&_reg_reduce_sum2_5_1, _reg_reduce_src2_5[_rr * 4 + 1]); + add_f32x2_inplace(&_reg_reduce_sum2_5_2, _reg_reduce_src2_5[_rr * 4 + 2]); + add_f32x2_inplace(&_reg_reduce_sum2_5_3, _reg_reduce_src2_5[_rr * 4 + 3]); + } + add_f32x2_inplace(&_reg_reduce_sum2_5_0, _reg_reduce_sum2_5_1); + add_f32x2_inplace(&_reg_reduce_sum2_5_2, _reg_reduce_sum2_5_3); + add_f32x2_inplace(&_reg_reduce_sum2_5_0, _reg_reduce_sum2_5_2); + float sv_split_sum = _reg_reduce_sum2_5_0.x + _reg_reduce_sum2_5_0.y; + block_sum = sv_split_sum; + } + uint32_t _fp8_0[16]; + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(sv_split[0]), "f"(sv_split[1]), + "f"(sv_split[2]), "f"(sv_split[3])); + _fp8_0[0] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(sv_split[4]), "f"(sv_split[5]), + "f"(sv_split[6]), "f"(sv_split[7])); + _fp8_0[1] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(sv_split[8]), "f"(sv_split[9]), + "f"(sv_split[10]), "f"(sv_split[11])); + _fp8_0[2] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(sv_split[12]), "f"(sv_split[13]), + "f"(sv_split[14]), "f"(sv_split[15])); + _fp8_0[3] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(sv_split[16]), "f"(sv_split[17]), + "f"(sv_split[18]), "f"(sv_split[19])); + _fp8_0[4] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(sv_split[20]), "f"(sv_split[21]), + "f"(sv_split[22]), "f"(sv_split[23])); + _fp8_0[5] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(sv_split[24]), "f"(sv_split[25]), + "f"(sv_split[26]), "f"(sv_split[27])); + _fp8_0[6] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(sv_split[28]), "f"(sv_split[29]), + "f"(sv_split[30]), "f"(sv_split[31])); + _fp8_0[7] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(sv_split[32]), "f"(sv_split[33]), + "f"(sv_split[34]), "f"(sv_split[35])); + _fp8_0[8] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(sv_split[36]), "f"(sv_split[37]), + "f"(sv_split[38]), "f"(sv_split[39])); + _fp8_0[9] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(sv_split[40]), "f"(sv_split[41]), + "f"(sv_split[42]), "f"(sv_split[43])); + _fp8_0[10] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(sv_split[44]), "f"(sv_split[45]), + "f"(sv_split[46]), "f"(sv_split[47])); + _fp8_0[11] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(sv_split[48]), "f"(sv_split[49]), + "f"(sv_split[50]), "f"(sv_split[51])); + _fp8_0[12] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(sv_split[52]), "f"(sv_split[53]), + "f"(sv_split[54]), "f"(sv_split[55])); + _fp8_0[13] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(sv_split[56]), "f"(sv_split[57]), + "f"(sv_split[58]), "f"(sv_split[59])); + _fp8_0[14] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(sv_split[60]), "f"(sv_split[61]), + "f"(sv_split[62]), "f"(sv_split[63])); + _fp8_0[15] = _packed; + } + int p_stage_base = smem_p_fp8_addr + (unsigned int)(phase * 8192); + #pragma unroll + for (int p_vec = 0; p_vec < 4; p_vec++) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_stage_base + (my_row * 128 + (n_half * 4 + p_vec) * 16 ^ (my_row * 128 + (n_half * 4 + p_vec) * 16 >> 7 & 7) << 4))), "r"(_fp8_0[p_vec * 4]), "r"(_fp8_0[p_vec * 4 + 1]), "r"(_fp8_0[p_vec * 4 + 2]), "r"(_fp8_0[p_vec * 4 + 3]) : "memory"); + } + } + } + { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + { + asm volatile("" ::: "memory"); + } + mbarrier_arrive(p_full_addr + (phase) * 8); + { + asm volatile("" ::: "memory"); + float _fma_3 = __fmaf_rn(row_sum_val, acc_scale, block_sum); + row_sum_val = _fma_3; + } + } + smem_stats_sum[stats_row] = row_sum_val; + smem_stats_final_max[stats_row] = row_max_val; + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(sum_ready_addr); + softmax_tile_cursor = softmax_tile_cursor + (num_kv_tiles_1 - bid % 2 + 1) / 2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_dealloc_addr); + } + } + // ---- Role: softmax_padding ---- + if (warp == 1 || warp == 3) { + { // softmax_padding_main + const int lowhead_softmax_padding = 0; + } + } + // ---- Role: correction_wg ---- + if (warp == 4 || warp == 6) { + { // correction_wg_main + float softmax_scale_log2_1 = bmm1_scale[0] * 1.4426950408889634f; + float output_scale = bmm2_scale[0]; + const int wg_dummy_inc_1 = 0; + int all_num_kv_tiles_3 = (sparse_topk + 128 - 1) / 128; + const int tmem_row_base_1 = ((1) ? warp % 2 * 32 : warp % 4 * 32); + const int n_half_1 = ((1) ? (int)(warp % 4 / 2) : 0); + const int my_row_1 = tmem_row_base_1 + lane; + const int stats_row_1 = n_half_1 * 64 + my_row_1; + const int corr_row = tmem_row_base_1 << 16; + int correction_tile_cursor = 0; + unsigned int _phase_q_full_0_2 = 0; + unsigned int _phase_pv_done_0 = 0; + unsigned int _phase_o_done_0 = 0; + unsigned int _phase_sum_ready_0 = 0; + unsigned int _phase_merge_ready_0 = 0; + unsigned int _phase_merge_full_0 = 0; + #pragma unroll + for (unsigned int work_idx_2 = 0; work_idx_2 < 1; work_idx_2++) { + int split_idx_3 = 0; + int query_idx_3 = bid / 2 >> 1; + int v_chunk_1 = bid / 2 & 1; + { + split_idx_3 = 0; + query_idx_3 = bid / 2; + v_chunk_1 = bid % 2; + } + int tiles_per_split_3 = all_num_kv_tiles_3 + 1 - 1; + int first_tile_3 = split_idx_3 * tiles_per_split_3; + int num_kv_tiles_2 = tiles_per_split_3; + #pragma unroll 1 + for (int tile_1 = 0; tile_1 < (num_kv_tiles_2 - bid % 2 + 1) / 2; tile_1++) { + int pipeline_tile_1 = correction_tile_cursor + tile_1; + int phase_1 = pipeline_tile_1 & 1; + int stats_wait_phase = pipeline_tile_1 >> 1 & 1; + mbarrier_wait(stats_addr + (phase_1) * 8, stats_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + float acc_scale_1 = smem_stats_max[phase_1 * 128 + stats_row_1]; + if (tile_1 > 0) { + mbarrier_wait(pv_done_addr, _phase_pv_done_0); + _phase_pv_done_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + } + if (tile_1 > 0) { + int _vote_1 = __any_sync(0xFFFFFFFF, acc_scale_1 < 1.0f); + int any_rescale = _vote_1; + if (any_rescale != 0) { + #pragma unroll + for (int vs = 0; vs < 2; vs++) { + int o_base = taddr + 256 + (unsigned int)(vs * 128) + (unsigned int)corr_row; + #pragma unroll + for (int c = 0; c < 128; c += 16) { + float _tmem_load_3[16]; + tmem_ld_x16(&_tmem_load_3[0], o_base + c); + const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + mul_f32x2_inplace(&reinterpret_cast(_tmem_load_3)[_ls], _scale2_0); + tmem_st_x16_f32(o_base + c, _tmem_load_3); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + } + } + } + } + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(corr_done_addr + (phase_1) * 8); + } + mbarrier_wait(o_done_addr, _phase_o_done_0); + _phase_o_done_0 ^= 1; + mbarrier_wait(sum_ready_addr, _phase_sum_ready_0); + _phase_sum_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + float wt_sum = smem_stats_sum[stats_row_1]; + float wt_max = smem_stats_final_max[stats_row_1]; + smem_corr_warp_pair_exchange[stats_row_1] = wt_sum; + asm volatile("barrier.sync %0, 64;" :: "r"(4 + warp % 2) : "memory"); + wt_sum = wt_sum + smem_corr_warp_pair_exchange[stats_row_1 ^ 64]; + int wt_peer_rank = bid % 2 ^ 1; + int wt_quadrant = warp - 4; + if (warp == 4) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(merge_full_addr, 66048); + } + } + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cluster.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(merge_ready_addr), "r"(wt_peer_rank) : "memory"); + if (n_half_1 == 0) { + smem_merge_stats_send[my_row_1] = wt_max; + smem_merge_stats_send[64 + my_row_1] = wt_sum; + } + int wt_send_block0 = (bid % 2 ^ 1) * 2; + #pragma unroll + for (int wt_b = 0; wt_b < 2; wt_b++) { + #pragma unroll + for (int wt_c = 0; wt_c < 64; wt_c += 32) { + float _tmem_load_4[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_4[0]), "=f"(_tmem_load_4[1]), "=f"(_tmem_load_4[2]), "=f"(_tmem_load_4[3]), "=f"(_tmem_load_4[4]), "=f"(_tmem_load_4[5]), "=f"(_tmem_load_4[6]), "=f"(_tmem_load_4[7]), "=f"(_tmem_load_4[8]), "=f"(_tmem_load_4[9]), "=f"(_tmem_load_4[10]), "=f"(_tmem_load_4[11]), "=f"(_tmem_load_4[12]), "=f"(_tmem_load_4[13]), "=f"(_tmem_load_4[14]), "=f"(_tmem_load_4[15]), "=f"(_tmem_load_4[16]), "=f"(_tmem_load_4[17]), "=f"(_tmem_load_4[18]), "=f"(_tmem_load_4[19]), "=f"(_tmem_load_4[20]), "=f"(_tmem_load_4[21]), "=f"(_tmem_load_4[22]), "=f"(_tmem_load_4[23]), "=f"(_tmem_load_4[24]), "=f"(_tmem_load_4[25]), "=f"(_tmem_load_4[26]), "=f"(_tmem_load_4[27]), "=f"(_tmem_load_4[28]), "=f"(_tmem_load_4[29]), "=f"(_tmem_load_4[30]), "=f"(_tmem_load_4[31]) + : "r"(taddr + 256 + (unsigned int)((wt_send_block0 + wt_b) * 64) + (unsigned int)wt_c + (unsigned int)corr_row)); + #pragma unroll + for (int wt_k = 0; wt_k < 32; wt_k += 4) { + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(smem_merge_send_addr + (unsigned int)((((wt_b * 64 + wt_c + wt_k) / 4 * 4 + wt_quadrant) * 32 + lane) * 16)), "f"(_tmem_load_4[wt_k]), "f"(_tmem_load_4[wt_k + 1]), "f"(_tmem_load_4[wt_k + 2]), "f"(_tmem_load_4[wt_k + 3]) : "memory"); + } + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("barrier.sync 6, 64;" ::: "memory"); + if (warp == 4) { + if (elect_sync()) { + mbarrier_wait(merge_ready_addr, _phase_merge_ready_0); + _phase_merge_ready_0 ^= 1; + uint32_t _mapa_0; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mapa_0) : "r"(smem_merge_recv_addr), "r"(wt_peer_rank)); + uint32_t _mapa_1; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mapa_1) : "r"(merge_full_addr), "r"(wt_peer_rank)); + asm volatile( + "cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes" + " [%0], [%1], %2, [%3];" + :: "r"(_mapa_0), "r"(smem_merge_send_addr), "r"((uint32_t)(65536)), "r"(_mapa_1) + : "memory"); + uint32_t _mapa_2; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mapa_2) : "r"(smem_merge_stats_recv_addr), "r"(wt_peer_rank)); + uint32_t _mapa_3; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mapa_3) : "r"(merge_full_addr), "r"(wt_peer_rank)); + asm volatile( + "cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes" + " [%0], [%1], %2, [%3];" + :: "r"(_mapa_2), "r"(smem_merge_stats_send_addr), "r"((uint32_t)(512)), "r"(_mapa_3) + : "memory"); + } + } + mbarrier_wait(merge_full_addr, _phase_merge_full_0); + _phase_merge_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + float wt_peer_max = smem_merge_stats_recv[my_row_1]; + float wt_peer_sum = smem_merge_stats_recv[64 + my_row_1]; + float _max_5 = max_noftz(wt_max, wt_peer_max); + float wt_merged_max = _max_5; + float _exp2_1 = approx_exp2((wt_max - wt_merged_max) * softmax_scale_log2_1); + float wt_s_own = ((wt_max > -CAKE_INF) ? _exp2_1 : 0.0f); + float _exp2_2 = approx_exp2((wt_peer_max - wt_merged_max) * softmax_scale_log2_1); + float wt_s_peer = ((wt_peer_max > -CAKE_INF) ? _exp2_2 : 0.0f); + float wt_merged_sum = wt_sum * wt_s_own + wt_peer_sum * wt_s_peer; + float _rcp_0 = approx_rcp(wt_merged_sum); + float wt_inv_sum = ((wt_merged_sum > 0.0f) ? _rcp_0 : 0.0f); + int wt_head_idx = my_row_1; + int wt_o_base = (query_idx_3 * num_heads + wt_head_idx) * 512; + if (bid % 2 == 0 && n_half_1 == 0 && wt_head_idx < num_heads) { + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(wt_merged_sum)); + partial_lse[query_idx_3 * num_heads + wt_head_idx + split_idx_3] = ((wt_merged_sum > 0.0f) ? wt_merged_max * softmax_scale_log2_1 + _log2_0 : -CAKE_INF); + } + int wt_keep_block0 = bid % 2 * 2; + #pragma unroll + for (int wt_b_1 = 0; wt_b_1 < 2; wt_b_1++) { + #pragma unroll + for (int wt_c_1 = 0; wt_c_1 < 64; wt_c_1 += 32) { + float _tmem_load_5[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_5[0]), "=f"(_tmem_load_5[1]), "=f"(_tmem_load_5[2]), "=f"(_tmem_load_5[3]), "=f"(_tmem_load_5[4]), "=f"(_tmem_load_5[5]), "=f"(_tmem_load_5[6]), "=f"(_tmem_load_5[7]), "=f"(_tmem_load_5[8]), "=f"(_tmem_load_5[9]), "=f"(_tmem_load_5[10]), "=f"(_tmem_load_5[11]), "=f"(_tmem_load_5[12]), "=f"(_tmem_load_5[13]), "=f"(_tmem_load_5[14]), "=f"(_tmem_load_5[15]), "=f"(_tmem_load_5[16]), "=f"(_tmem_load_5[17]), "=f"(_tmem_load_5[18]), "=f"(_tmem_load_5[19]), "=f"(_tmem_load_5[20]), "=f"(_tmem_load_5[21]), "=f"(_tmem_load_5[22]), "=f"(_tmem_load_5[23]), "=f"(_tmem_load_5[24]), "=f"(_tmem_load_5[25]), "=f"(_tmem_load_5[26]), "=f"(_tmem_load_5[27]), "=f"(_tmem_load_5[28]), "=f"(_tmem_load_5[29]), "=f"(_tmem_load_5[30]), "=f"(_tmem_load_5[31]) + : "r"(taddr + 256 + (unsigned int)((wt_keep_block0 + wt_b_1) * 64) + (unsigned int)wt_c_1 + (unsigned int)corr_row)); + float wt_peer[32]; + #pragma unroll + for (int wt_k_1 = 0; wt_k_1 < 32; wt_k_1 += 4) { + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&wt_peer[wt_k_1])), "=r"(*reinterpret_cast(&wt_peer[(wt_k_1) + 1])), "=r"(*reinterpret_cast(&wt_peer[(wt_k_1) + 2])), "=r"(*reinterpret_cast(&wt_peer[(wt_k_1) + 3])) + : "r"(smem_merge_recv_addr + (unsigned int)((((wt_b_1 * 64 + wt_c_1 + wt_k_1) / 4 * 4 + wt_quadrant) * 32 + lane) * 16))); + } + const float2 _scale2_1 = {wt_s_own, wt_s_own}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + mul_f32x2_inplace(&reinterpret_cast(_tmem_load_5)[_ls], _scale2_1); + const float2 _scale2_2 = {wt_s_peer, wt_s_peer}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + mul_f32x2_inplace(&reinterpret_cast(wt_peer)[_ls], _scale2_2); + #pragma unroll + for (int _la = 0; _la < 16; _la++) + add_f32x2_inplace(&reinterpret_cast(_tmem_load_5)[_la], reinterpret_cast(wt_peer)[_la]); + int wt_gmem = wt_o_base + (wt_keep_block0 + wt_b_1) * 128 + n_half_1 * 64 + wt_c_1; + if (wt_head_idx < num_heads) { + #pragma unroll + for (int wt_j = 0; wt_j < 32; wt_j += 8) { + { + const float2 _prescale2_3 = {wt_inv_sum * output_scale, wt_inv_sum * output_scale}; + #if __CUDA_ARCH__ >= 1000 + #pragma unroll + for (int _ps = 0; _ps < 4; _ps++) + mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_5[wt_j])[_ps], _prescale2_3); + #else + #pragma unroll + for (int _ps = 0; _ps < 8; _ps++) + _tmem_load_5[wt_j + _ps] *= wt_inv_sum * output_scale; + #endif + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_tmem_load_5[wt_j + 0], _tmem_load_5[wt_j + 1]); + _pk[1] = __floats2bfloat162_rn(_tmem_load_5[wt_j + 2], _tmem_load_5[wt_j + 3]); + _pk[2] = __floats2bfloat162_rn(_tmem_load_5[wt_j + 4], _tmem_load_5[wt_j + 5]); + _pk[3] = __floats2bfloat162_rn(_tmem_load_5[wt_j + 6], _tmem_load_5[wt_j + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(O + (wt_gmem + wt_j)))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + } + } + correction_tile_cursor = correction_tile_cursor + (num_kv_tiles_2 - bid % 2 + 1) / 2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_dealloc_addr); + } + } + // ---- Role: correction_padding ---- + if (warp == 5 || warp == 7) { + { // correction_padding_main + const int lowhead_correction_padding = 0; + } + } + // ---- Role: mma_warp ---- + if (warp == 8) { + { // mma_warp_main + const int wg2_dummy_1 = 0; + int all_num_kv_tiles_4 = (sparse_topk + 128 - 1) / 128; + unsigned int mma_k_stage = 0; + unsigned int mma_v_stage = 0; + unsigned int mma_kv_stage = 0; + int mma_tile_cursor = 0; + unsigned int _phase_q_full_0_3 = 0; + unsigned int _phase_k_full = 0; + #pragma unroll + for (unsigned int work_idx_3 = 0; work_idx_3 < 1; work_idx_3++) { + int split_idx_4 = 0; + { + split_idx_4 = 0; + } + int tiles_per_split_4 = all_num_kv_tiles_4 + 1 - 1; + int first_tile_4 = split_idx_4 * tiles_per_split_4; + int num_kv_tiles_3 = tiles_per_split_4; + int wt_num_local_tiles_1 = (num_kv_tiles_3 - bid % 2 + 1) / 2; + mbarrier_wait(q_full_addr, _phase_q_full_0_3); + _phase_q_full_0_3 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int wt_first_pv = 1; + #pragma unroll 1 + for (int wt_tile_1 = 0; wt_tile_1 < wt_num_local_tiles_1; wt_tile_1++) { + int wt_pipeline_tile = mma_tile_cursor + wt_tile_1; + int wt_phase = wt_pipeline_tile & 1; + int wt_score_col = ((wt_phase != 0) ? 128 : 0); + mbarrier_wait(k_full_addr + (mma_k_stage) * 8, _phase_k_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_0 = make_warp_uniform((((smem_q_full_addr) >> 4) & 0x3FFF) + (0) * 2048); + int _mma_b_lo_0 = make_warp_uniform((((smem_k_full_addr) >> 4) & 0x3FFF) + (mma_k_stage) * 4096); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 69206032;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (wt_score_col))), "r"(0)); + int _mma_a_lo_1 = make_warp_uniform((((smem_q_full_addr + 8192) >> 4) & 0x3FFF) + (0) * 2048); + int _mma_b_lo_1 = make_warp_uniform((((smem_k_full_addr + 16384) >> 4) & 0x3FFF) + (mma_k_stage) * 4096); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 69206032;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_tmem_scratch + (wt_score_col))), "r"(1)); + int _mma_a_lo_2 = make_warp_uniform((((smem_q_full_addr + 16384) >> 4) & 0x3FFF) + (0) * 2048); + int _mma_b_lo_2 = make_warp_uniform((((smem_k_full_addr + 32768) >> 4) & 0x3FFF) + (mma_k_stage) * 4096); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 69206032;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (wt_score_col))), "r"(1)); + int _mma_a_lo_3 = make_warp_uniform((((smem_q_full_addr + 24576) >> 4) & 0x3FFF) + (0) * 2048); + int _mma_b_lo_3 = make_warp_uniform((((smem_k_full_addr + 49152) >> 4) & 0x3FFF) + (mma_k_stage) * 4096); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 69206032;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"((tmem_tmem_scratch + (wt_score_col))), "r"(1)); + elect_commit(s_full_addr + (wt_phase) * 8); + if (wt_tile_1 > 0) { + int wt_prev_pipeline_tile = wt_pipeline_tile - 1; + int wt_prev_phase = wt_prev_pipeline_tile & 1; + int wt_pv_wait_phase = wt_prev_pipeline_tile >> 1 & 1; + mbarrier_wait(p_full_addr + (wt_prev_phase) * 8, wt_pv_wait_phase); + mbarrier_wait(corr_done_addr + (wt_prev_phase) * 8, wt_pv_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + unsigned int wt_prev_kv_stage = mma_k_stage ^ 1; + int _mma_a_lo_4 = make_warp_uniform((((smem_p_fp8_addr) >> 4) & 0x3FFF) + (wt_prev_phase) * 512); + int _mma_b_lo_4 = make_warp_uniform(((((smem_v_full_addr) >> 4) & 0x3FFF) | 0x4000000) + (wt_prev_kv_stage) * 4096); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 69271568;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"((tmem_tmem_scratch + (256))), "r"(((wt_first_pv) ? 0 : 1))); + int _mma_a_lo_5 = make_warp_uniform((((smem_p_fp8_addr) >> 4) & 0x3FFF) + (wt_prev_phase) * 512); + int _mma_b_lo_5 = make_warp_uniform(((((smem_v_full_addr + 16384) >> 4) & 0x3FFF) | 0x4000000) + (wt_prev_kv_stage) * 4096); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 69271568;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_5), "r"(_mma_b_lo_5), "r"((tmem_tmem_scratch + (320))), "r"(((wt_first_pv) ? 0 : 1))); + int _mma_a_lo_6 = make_warp_uniform((((smem_p_fp8_addr) >> 4) & 0x3FFF) + (wt_prev_phase) * 512); + int _mma_b_lo_6 = make_warp_uniform(((((smem_v_full_addr + 32768) >> 4) & 0x3FFF) | 0x4000000) + (wt_prev_kv_stage) * 4096); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 69271568;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_6), "r"(_mma_b_lo_6), "r"((tmem_tmem_scratch + (384))), "r"(((wt_first_pv) ? 0 : 1))); + int _mma_a_lo_7 = make_warp_uniform((((smem_p_fp8_addr) >> 4) & 0x3FFF) + (wt_prev_phase) * 512); + int _mma_b_lo_7 = make_warp_uniform(((((smem_v_full_addr + 49152) >> 4) & 0x3FFF) | 0x4000000) + (wt_prev_kv_stage) * 4096); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 69271568;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_7), "r"(_mma_b_lo_7), "r"((tmem_tmem_scratch + (448))), "r"(((wt_first_pv) ? 0 : 1))); + wt_first_pv = 0; + elect_commit(k_empty_addr + (wt_prev_kv_stage) * 8); + elect_commit(pv_done_addr); + } + mma_k_stage += 1; + if (mma_k_stage == 2) { mma_k_stage = 0; _phase_k_full ^= 1; } + } + if (wt_num_local_tiles_1 > 0) { + int wt_last_pipeline_tile = mma_tile_cursor + wt_num_local_tiles_1 - 1; + int wt_last_phase = wt_last_pipeline_tile & 1; + int wt_drain_wait_phase = wt_last_pipeline_tile >> 1 & 1; + mbarrier_wait(p_full_addr + (wt_last_phase) * 8, wt_drain_wait_phase); + mbarrier_wait(corr_done_addr + (wt_last_phase) * 8, wt_drain_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + unsigned int wt_last_kv_stage = mma_k_stage ^ 1; + int _mma_a_lo_8 = make_warp_uniform((((smem_p_fp8_addr) >> 4) & 0x3FFF) + (wt_last_phase) * 512); + int _mma_b_lo_8 = make_warp_uniform(((((smem_v_full_addr) >> 4) & 0x3FFF) | 0x4000000) + (wt_last_kv_stage) * 4096); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 69271568;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_8), "r"(_mma_b_lo_8), "r"((tmem_tmem_scratch + (256))), "r"(((wt_first_pv) ? 0 : 1))); + int _mma_a_lo_9 = make_warp_uniform((((smem_p_fp8_addr) >> 4) & 0x3FFF) + (wt_last_phase) * 512); + int _mma_b_lo_9 = make_warp_uniform(((((smem_v_full_addr + 16384) >> 4) & 0x3FFF) | 0x4000000) + (wt_last_kv_stage) * 4096); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 69271568;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_9), "r"(_mma_b_lo_9), "r"((tmem_tmem_scratch + (320))), "r"(((wt_first_pv) ? 0 : 1))); + int _mma_a_lo_10 = make_warp_uniform((((smem_p_fp8_addr) >> 4) & 0x3FFF) + (wt_last_phase) * 512); + int _mma_b_lo_10 = make_warp_uniform(((((smem_v_full_addr + 32768) >> 4) & 0x3FFF) | 0x4000000) + (wt_last_kv_stage) * 4096); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 69271568;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_10), "r"(_mma_b_lo_10), "r"((tmem_tmem_scratch + (384))), "r"(((wt_first_pv) ? 0 : 1))); + int _mma_a_lo_11 = make_warp_uniform((((smem_p_fp8_addr) >> 4) & 0x3FFF) + (wt_last_phase) * 512); + int _mma_b_lo_11 = make_warp_uniform(((((smem_v_full_addr + 49152) >> 4) & 0x3FFF) | 0x4000000) + (wt_last_kv_stage) * 4096); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 69271568;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_11), "r"(_mma_b_lo_11), "r"((tmem_tmem_scratch + (448))), "r"(((wt_first_pv) ? 0 : 1))); + elect_commit(k_empty_addr + (wt_last_kv_stage) * 8); + } + elect_commit(q_empty_addr); + elect_commit(o_done_addr); + mma_tile_cursor = mma_tile_cursor + wt_num_local_tiles_1; + } + mbarrier_arrive(tmem_dealloc_addr); + unsigned int _phase_tmem_dealloc_0 = 0; + mbarrier_wait(tmem_dealloc_addr, _phase_tmem_dealloc_0); + _phase_tmem_dealloc_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + } + // ---- Role: empty1 ---- + if (warp >= 10 && warp <= 11) { + { // empty1_main + const int wg2_dummy_2 = 0; + unsigned int _phase_q_full_0_4 = 0; + } + } + + // Cleanup +} + +} // extern "C" diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_c1d444bba4cdb8273c7c_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_c1d444bba4cdb8273c7c_binding.cu new file mode 100644 index 00000000000..b03ef9f4476 --- /dev/null +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_c1d444bba4cdb8273c7c_binding.cu @@ -0,0 +1,526 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +// Generated by the Cake source exporter — do not edit. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_c1d444bba4cdb8273c7c'. +#include +#include +#include +#include + +#include "tvm_ffi_utils.h" + +#include +#include +#include +#include +#include +#include +#include + +struct CakeTensorMap; +extern "C" __global__ void kernel_cake_dsv4_c1d444bba4cdb8273c7c(CakeTensorMap const* tmap_q, CakeTensorMap const* tmap_swa_kv, CakeTensorMap const* tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int max_q_len, int batch_size, int has_sinks); + + +namespace cake_host_shim_1a24ae3a58f0bd58 { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t CakeDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +// Explicit immutable caller storage: initialize once, outside graph capture, +// exactly as the production TmaDeviceSlot path. The owner must retain the +// allocation and must never expose it as mutable scratch or rewrite its bytes. +template +static inline void PrepareImmutableCallerTmaWorkspace( + const tvm::ffi::Tensor& retained_workspace, + const CUtensorMap (&maps)[N], cudaStream_t stream) { + void* workspace = retained_workspace.data_ptr(); + CUcontext context = nullptr; + CUresult result = cuCtxGetCurrent(&context); + TVM_FFI_CHECK(result == CUDA_SUCCESS && context != nullptr, RuntimeError) + << "immutable TMA workspace requires an active CUDA context"; + CUcontext allocation_context = nullptr; + result = cuPointerGetAttribute( + &allocation_context, CU_POINTER_ATTRIBUTE_CONTEXT, + reinterpret_cast(workspace)); + TVM_FFI_CHECK(result == CUDA_SUCCESS && allocation_context == context, RuntimeError) + << "immutable TMA workspace must belong to the active CUDA context"; + unsigned long long buffer_id = 0; + result = cuPointerGetAttribute( + &buffer_id, CU_POINTER_ATTRIBUTE_BUFFER_ID, + reinterpret_cast(workspace)); + TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) + << "querying immutable TMA allocation identity failed"; + std::string key = std::to_string(reinterpret_cast(context)); + key += ":" + std::to_string(reinterpret_cast(workspace)); + key += ":" + std::to_string(buffer_id); + std::string descriptor_bytes(reinterpret_cast(maps), sizeof(maps)); + struct InitializedWorkspace { + tvm::ffi::Tensor owner; + std::string descriptor_bytes; + }; + static std::mutex mu; + // Keep the actual Tensor allocation alive: a caching allocator may reuse a + // freed suballocation without changing its enclosing CUDA buffer identity. + static auto* initialized = new std::unordered_map(); + std::lock_guard lock(mu); + auto it = initialized->find(key); + if (it != initialized->end()) { + TVM_FFI_CHECK(it->second.descriptor_bytes == descriptor_bytes, ValueError) + << "immutable TMA descriptors changed; provide new private descriptor storage"; + return; + } + CUstreamCaptureStatus capture_status = CU_STREAM_CAPTURE_STATUS_NONE; + result = cuStreamIsCapturing(reinterpret_cast(stream), &capture_status); + TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) + << "querying immutable TMA initialization capture state failed"; + TVM_FFI_CHECK(capture_status == CU_STREAM_CAPTURE_STATUS_NONE, RuntimeError) + << "immutable TMA storage must be initialized before CUDA Graph capture"; + // Synchronous host-to-device initialization completes before publishing the + // slot to any thread/stream, matching TmaDeviceSlot. The consumer retains its + // generated tensor-map acquire; the mutable upload's release fence is unchanged. + result = cuMemcpyHtoD(reinterpret_cast(workspace), maps, sizeof(maps)); + TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) + << "initializing immutable caller TMA descriptors failed"; + initialized->emplace( + std::move(key), InitializedWorkspace{retained_workspace, std::move(descriptor_bytes)}); +} + +static inline void* CallerTmaWorkspaceSlot( + const tvm::ffi::TensorView& workspace, size_t slot) { + return static_cast(workspace.data_ptr()) + slot * sizeof(CUtensorMap); +} + +// 4D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + int64_t d3 = t.size(t.ndim() - 3); + TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + TVM_FFI_CHECK(d1 % 128 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 128; + uint64_t global_dim[4] = {(uint64_t)(128), (uint64_t)(d2), (uint64_t)((d1 / 128)), (uint64_t)(d3)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (128, 8, 1, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 128; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + (uint64_t)((carrier_stride_1 * 8) / 8), + (uint64_t)((carrier_stride_2 * 8) / 8), + }; + uint32_t box_dim[4] = {128u, 8u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_swa_kv' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_swa_kv' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; + TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (128, 1) exceeds resolved global dims for 'tmap_swa_kv'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + }; + uint32_t box_dim[2] = {128u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_compressed_kv' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_compressed_kv(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_compressed_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_compressed_kv' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_compressed_kv' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_compressed_kv"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_compressed_kv' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved a non-positive global dim"; + TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (128, 1) exceeds resolved global dims for 'tmap_compressed_kv'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + }; + uint32_t box_dim[2] = {128u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_kv') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_compressed_kv, TensorView arg_O, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_seq_lens, TensorView arg_cum_seq_lens_q, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_max_q_len, int64_t arg_batch_size, int64_t arg_has_sinks, tvm::ffi::Tensor arg_tma_descriptor_workspace, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 1, 8, 1); + CheckContiguous(arg_tmap_q, "tmap_q"); + CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); + CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 1, 8, 1); + CheckCudaTensor(arg_tmap_compressed_kv, "tmap_compressed_kv"); + CheckDtype(arg_tmap_compressed_kv, "tmap_compressed_kv", 1, 8, 1); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_swa_indices, "swa_indices"); + CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); + CheckContiguous(arg_swa_indices, "swa_indices"); + CheckCudaTensor(arg_compressed_indices, "compressed_indices"); + CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); + CheckContiguous(arg_compressed_indices, "compressed_indices"); + CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); + CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckCudaTensor(arg_seq_lens, "seq_lens"); + CheckDtype(arg_seq_lens, "seq_lens", 0, 32, 1); + CheckContiguous(arg_seq_lens, "seq_lens"); + CheckCudaTensor(arg_cum_seq_lens_q, "cum_seq_lens_q"); + CheckDtype(arg_cum_seq_lens_q, "cum_seq_lens_q", 0, 32, 1); + CheckContiguous(arg_cum_seq_lens_q, "cum_seq_lens_q"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); + CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); + CheckContiguous(arg_bmm1_scale, "bmm1_scale"); + CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); + CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); + CheckContiguous(arg_bmm2_scale, "bmm2_scale"); + TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) + << "scalar 'swa_index_stride' value " << arg_swa_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) + << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) + << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) + << "scalar 'sparse_topk' value " << arg_sparse_topk + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_q_len >= -2147483648LL && arg_max_q_len <= 2147483647LL, ValueError) + << "scalar 'max_q_len' value " << arg_max_q_len + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_batch_size >= -2147483648LL && arg_batch_size <= 2147483647LL, ValueError) + << "scalar 'batch_size' value " << arg_batch_size + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) + << "scalar 'has_sinks' value " << arg_has_sinks + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); + CheckSameCudaDevice(arg_tmap_compressed_kv, arg_tmap_q, "tmap_compressed_kv", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); + CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); + CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); + CheckSameCudaDevice(arg_seq_lens, arg_tmap_q, "seq_lens", "tmap_q"); + CheckSameCudaDevice(arg_cum_seq_lens_q, arg_tmap_q, "cum_seq_lens_q", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); + CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); + CheckCudaTensor(arg_tma_descriptor_workspace, "tma_descriptor_workspace"); + CheckSameCudaDevice(arg_tma_descriptor_workspace, arg_tmap_q, "tma_descriptor_workspace", "tmap_q"); + CheckContiguous(arg_tma_descriptor_workspace, "tma_descriptor_workspace"); + CheckDtype(arg_tma_descriptor_workspace, "tma_descriptor_workspace", 1, 8, 1); + TVM_FFI_CHECK(arg_tma_descriptor_workspace.numel() >= 384, ValueError) + << "tma_descriptor_workspace requires at least 384 bytes"; + TVM_FFI_CHECK( + reinterpret_cast(arg_tma_descriptor_workspace.data_ptr()) % 128 == 0, + ValueError) + << "tma_descriptor_workspace must be 128-byte aligned"; + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + + cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); + CUtensorMap h_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CakeTensorMap const* p_tmap_q = static_cast(CallerTmaWorkspaceSlot(arg_tma_descriptor_workspace, 0u)); + CUtensorMap h_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); + CakeTensorMap const* p_tmap_swa_kv = static_cast(CallerTmaWorkspaceSlot(arg_tma_descriptor_workspace, 1u)); + CUtensorMap h_tmap_compressed_kv = EncodeTma_tmap_compressed_kv(arg_tmap_compressed_kv); + CakeTensorMap const* p_tmap_compressed_kv = static_cast(CallerTmaWorkspaceSlot(arg_tma_descriptor_workspace, 2u)); + __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); + int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); + int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); + int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); + int* p_seq_lens = static_cast(arg_seq_lens.data_ptr()); + int* p_cum_seq_lens_q = static_cast(arg_cum_seq_lens_q.data_ptr()); + float* p_sinks = static_cast(arg_sinks.data_ptr()); + float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); + float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); + int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; + int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; + int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; + int32_t v_sparse_topk = (int32_t)arg_sparse_topk; + int32_t v_max_q_len = (int32_t)arg_max_q_len; + int32_t v_batch_size = (int32_t)arg_batch_size; + int32_t v_has_sinks = (int32_t)arg_has_sinks; + void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_tmap_compressed_kv, &p_O, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_seq_lens, &p_cum_seq_lens_q, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_max_q_len, &v_batch_size, &v_has_sinks}; + + static std::mutex smem_mu; + static auto* smem_status_by_device = + new std::unordered_map(); + cudaError_t smem_status = cudaSuccess; + { + std::lock_guard lock(smem_mu); + auto it = smem_status_by_device->find(dev.device_id); + if (it == smem_status_by_device->end()) { + int current_device = -1; + smem_status = cudaGetDevice(¤t_device); + if (smem_status == cudaSuccess && current_device != dev.device_id) { + smem_status = cudaErrorInvalidDevice; + } + if (smem_status == cudaSuccess) { + smem_status = cudaFuncSetAttribute( + reinterpret_cast(kernel_cake_dsv4_c1d444bba4cdb8273c7c), + cudaFuncAttributeMaxDynamicSharedMemorySize, + 147072); + } + smem_status_by_device->emplace(dev.device_id, smem_status); + } else { + smem_status = it->second; + } + } + TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) + << "cudaFuncSetAttribute for kernel_cake_dsv4_c1d444bba4cdb8273c7c failed: " + << cudaGetErrorString(smem_status); + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(512u, 1u, 1u); + const CUtensorMap workspace_maps[] = {h_tmap_q, h_tmap_swa_kv, h_tmap_compressed_kv}; + PrepareImmutableCallerTmaWorkspace(arg_tma_descriptor_workspace, workspace_maps, stream); + + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes = 147072u; + config.stream = stream; + cudaError_t launch_status = cudaLaunchKernelExC( + &config, reinterpret_cast(kernel_cake_dsv4_c1d444bba4cdb8273c7c), kargs); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernelExC for kernel_cake_dsv4_c1d444bba4cdb8273c7c failed: " + << cudaGetErrorString(launch_status); +} + +} // namespace cake_host_shim_1a24ae3a58f0bd58 + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_1a24ae3a58f0bd58::Run); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_c1d444bba4cdb8273c7c_kernel.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_c1d444bba4cdb8273c7c_kernel.cu new file mode 100644 index 00000000000..f2cfc0f9c4f --- /dev/null +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_c1d444bba4cdb8273c7c_kernel.cu @@ -0,0 +1,2655 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); +template +struct __align__(128) CakeTensorMapPack { CakeTensorMap maps[N]; }; + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +#define CAKE_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SOURCE_EXACT_SCRATCH_OFFSET 0 +#define NUM_Q_PIPE_STAGES 1 +#define NUM_KV_PIPE_STAGES 8 +#define NUM_PAGE_PIPE_STAGES 6 +#define NUM_SCORE_PIPE_STAGES 2 +#define NUM_STATS_PIPE_STAGES 2 +#define NUM_P_PIPE_STAGES 2 +#define NUM_O_PIPE_STAGES 1 +#define NUM_WORK_PIPE_STAGES 2 +#define NUM_THROTTLE_PIPE_STAGES 2 +#define SMEM_SMEM_Q_OFF 1024 +#define SMEM_SMEM_Q_STAGE_BYTES 1024 +#define SMEM_SMEM_Q_STRIDE 1024 +#define SMEM_SMEM_KV_OFF 5120 +#define SMEM_SMEM_KV_STAGE_BYTES 16384 +#define SMEM_SMEM_KV_STRIDE 16384 +#define SMEM_SMEM_V_OFF 5120 +#define SMEM_SMEM_V_STAGE_BYTES 16384 +#define SMEM_SMEM_V_STRIDE 16384 +#define SMEM_SMEM_P_OFF 136192 +#define SMEM_SMEM_P_STAGE_BYTES 1024 +#define SMEM_SMEM_P_STRIDE 1024 +#define SMEM_SMEM_PAGE_OFFSETS_OFF 138240 +#define SMEM_SMEM_PAGE_OFFSETS_STAGE_BYTES 1024 +#define SMEM_SMEM_PAGE_OFFSETS_STRIDE 1024 +#define SMEM_SMEM_O_OFF 144384 +#define SMEM_SMEM_O_STAGE_BYTES 2048 +#define SMEM_SMEM_O_STRIDE 2048 +#define SMEM_SMEM_SOFTMAX_OFF 146432 +#define SMEM_SMEM_SOFTMAX_STAGE_BYTES 128 +#define SMEM_SMEM_SOFTMAX_STRIDE 128 +#define SMEM_SMEM_SOFTMAX_U32_OFF 146432 +#define SMEM_SMEM_SOFTMAX_U32_STAGE_BYTES 128 +#define SMEM_SMEM_SOFTMAX_U32_STRIDE 128 +#define SMEM_SMEM_SOFTMAX_UNUSED_OFF 146560 +#define SMEM_SMEM_SOFTMAX_UNUSED_STAGE_BYTES 128 +#define SMEM_SMEM_SOFTMAX_UNUSED_STRIDE 128 +#define SMEM_SMEM_CORR_OFF 146688 +#define SMEM_SMEM_CORR_STAGE_BYTES 128 +#define SMEM_SMEM_CORR_STRIDE 128 +#define SMEM_SMEM_STATS_OFF 146816 +#define SMEM_SMEM_STATS_STAGE_BYTES 96 +#define SMEM_SMEM_STATS_STRIDE 96 +#define SMEM_WORK_RESPONSE_OFF 147008 +#define SMEM_WORK_RESPONSE_STAGE_BYTES 16 +#define SMEM_WORK_RESPONSE_STRIDE 16 +#define SMEM_TOTAL 147072 +#define THREADS 512 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count) : "memory"); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f8f6f4( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x4(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x4.b32" + " {%0, %1, %2, %3}, [%4];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_2d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3}], [%4];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr) { + // Canonical .shared::cta form for non-multicast gather4, matching + // trtllm-gen / cuda_ptx and the PTX ISA qualifier order + // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, + // .shared::cluster is reserved for the multicast variant (ctaMask). + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x4_wait(float* dst, int addr) { + tmem_ld_x4(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(512, 1) void +kernel_cake_dsv4_c1d444bba4cdb8273c7c(CakeTensorMap const* tmap_q, CakeTensorMap const* tmap_swa_kv, CakeTensorMap const* tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int max_q_len, int batch_size, int has_sinks) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int mbar_base = smem; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define kv_full_addr (mbar_base + 16) + #define kv_empty_addr (mbar_base + 80) + #define page_full_addr (mbar_base + 144) + #define page_empty_addr (mbar_base + 192) + #define score_full_addr (mbar_base + 240) + #define score_empty_addr (mbar_base + 256) + #define stats_full_addr (mbar_base + 272) + #define stats_empty_addr (mbar_base + 288) + #define p_full_addr (mbar_base + 304) + #define p_empty_addr (mbar_base + 320) + #define o_full_addr (mbar_base + 336) + #define o_empty_addr (mbar_base + 344) + #define work_full_addr (mbar_base + 352) + #define work_empty_addr (mbar_base + 368) + #define throttle_full_addr (mbar_base + 384) + #define throttle_empty_addr (mbar_base + 400) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + if (tid == 0) { + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_q)) : "memory"); + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_swa_kv)) : "memory"); + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_compressed_kv)) : "memory"); + } + __syncthreads(); + + + // Kernel setup ops + uint8_t* smem_q = reinterpret_cast(smem_raw + 1024); + const int smem_q_addr = smem + 1024; + uint8_t* smem_kv = reinterpret_cast(smem_raw + 5120); + const int smem_kv_addr = smem + 5120; + uint8_t* smem_v = reinterpret_cast(smem_raw + 5120); + const int smem_v_addr = smem + 5120; + uint8_t* smem_p = reinterpret_cast(smem_raw + 136192); + const int smem_p_addr = smem + 136192; + int* smem_page_offsets = reinterpret_cast(smem_raw + 138240); + const int smem_page_offsets_addr = smem + 138240; + __nv_bfloat16* smem_o = reinterpret_cast<__nv_bfloat16*>(smem_raw + 144384); + const int smem_o_addr = smem + 144384; + float* smem_softmax = reinterpret_cast(smem_raw + 146432); + const int smem_softmax_addr = smem + 146432; + unsigned int* smem_softmax_u32 = reinterpret_cast(smem_raw + 146432); + const int smem_softmax_u32_addr = smem + 146432; + float* smem_softmax_unused = reinterpret_cast(smem_raw + 146560); + const int smem_softmax_unused_addr = smem + 146560; + float* smem_corr = reinterpret_cast(smem_raw + 146688); + const int smem_corr_addr = smem + 146688; + float* smem_stats = reinterpret_cast(smem_raw + 146816); + const int smem_stats_addr = smem + 146816; + unsigned int* work_response = reinterpret_cast(smem_raw + 147008); + const int work_response_addr = smem + 147008; + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)(tmap_q)) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)(tmap_swa_kv)) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)(tmap_compressed_kv)) : "memory"); + + // Mbarrier init (18 pipeline groups, 0 ordered-sequence groups, 52 barriers) + // Mbarriers at smem_raw[0..416) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // --- pipeline 'q_pipe' --- + // q_full: 1 barriers, init_count=1 + mbarrier_init(smem + 0, 1); + // q_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 8, 1); + // --- pipeline 'kv_pipe' --- + // kv_full: 8 barriers, init_count=1 + mbarrier_init(smem + 16, 1); + mbarrier_init(smem + 24, 1); + mbarrier_init(smem + 32, 1); + mbarrier_init(smem + 40, 1); + mbarrier_init(smem + 48, 1); + mbarrier_init(smem + 56, 1); + mbarrier_init(smem + 64, 1); + mbarrier_init(smem + 72, 1); + // kv_empty: 8 barriers, init_count=1 + mbarrier_init(smem + 80, 1); + mbarrier_init(smem + 88, 1); + mbarrier_init(smem + 96, 1); + mbarrier_init(smem + 104, 1); + mbarrier_init(smem + 112, 1); + mbarrier_init(smem + 120, 1); + mbarrier_init(smem + 128, 1); + mbarrier_init(smem + 136, 1); + // --- pipeline 'page_pipe' --- + // page_full: 6 barriers, init_count=32 + mbarrier_init(smem + 144, 32); + mbarrier_init(smem + 152, 32); + mbarrier_init(smem + 160, 32); + mbarrier_init(smem + 168, 32); + mbarrier_init(smem + 176, 32); + mbarrier_init(smem + 184, 32); + // page_empty: 6 barriers, init_count=128 + mbarrier_init(smem + 192, 128); + mbarrier_init(smem + 200, 128); + mbarrier_init(smem + 208, 128); + mbarrier_init(smem + 216, 128); + mbarrier_init(smem + 224, 128); + mbarrier_init(smem + 232, 128); + // --- pipeline 'score_pipe' --- + // score_full: 2 barriers, init_count=1 + mbarrier_init(smem + 240, 1); + mbarrier_init(smem + 248, 1); + // score_empty: 2 barriers, init_count=128 + mbarrier_init(smem + 256, 128); + mbarrier_init(smem + 264, 128); + // --- pipeline 'stats_pipe' --- + // stats_full: 2 barriers, init_count=128 + mbarrier_init(smem + 272, 128); + mbarrier_init(smem + 280, 128); + // stats_empty: 2 barriers, init_count=128 + mbarrier_init(smem + 288, 128); + mbarrier_init(smem + 296, 128); + // --- pipeline 'p_pipe' --- + // p_full: 2 barriers, init_count=128 + mbarrier_init(smem + 304, 128); + mbarrier_init(smem + 312, 128); + // p_empty: 2 barriers, init_count=1 + mbarrier_init(smem + 320, 1); + mbarrier_init(smem + 328, 1); + // --- pipeline 'o_pipe' --- + // o_full: 1 barriers, init_count=1 + mbarrier_init(smem + 336, 1); + // o_empty: 1 barriers, init_count=128 + mbarrier_init(smem + 344, 128); + // --- pipeline 'work_pipe' --- + // work_full: 2 barriers, init_count=1 + mbarrier_init(smem + 352, 1); + mbarrier_init(smem + 360, 1); + // work_empty: 2 barriers, init_count=512 + mbarrier_init(smem + 368, 512); + mbarrier_init(smem + 376, 512); + // --- pipeline 'throttle_pipe' --- + // throttle_full: 2 barriers, init_count=128 + mbarrier_init(smem + 384, 128); + mbarrier_init(smem + 392, 128); + // throttle_empty: 2 barriers, init_count=32 + mbarrier_init(smem + 400, 32); + mbarrier_init(smem + 408, 32); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncwarp(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 416); + if (warp == 0) { + int _tmem_hold = smem + 416; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_source_exact_scratch = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Inc phase consumes the registers released above. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 136;"); + } + + // ---- Role: softmax ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 152;"); + { // softmax_main + unsigned int work_stage = 0; + unsigned int score_stage = 0; + unsigned int score_phase = 0; + unsigned int stats_stage = 0; + unsigned int stats_phase = 1; + unsigned int p_stage = 0; + unsigned int p_phase = 1; + unsigned int work_x = blockIdx.x; + unsigned int work_y = blockIdx.y; + unsigned int work_z = blockIdx.z; + const int warp_in_wg = warp; + const int wg_tid = warp_in_wg * 32 + lane; + const int col_pair_base = wg_tid % 4 * 2; + const int ldtm_row_base = warp_in_wg * 32 + lane / 4; + float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; + unsigned int _phase_work_full = 0; + #pragma unroll 1 + for (unsigned int _work = 0; _work < 120; _work++) { + int q_begin = cum_seq_lens_q[work_z]; + int q_end = cum_seq_lens_q[work_z + 1]; + int q_len = q_end - q_begin; + int query_idx = (unsigned int)q_begin + work_x; + int active_topk = 0; + if (work_x < (unsigned int)q_len) { + int _max_4 = ((sparse_topk_lens[query_idx] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx] + sparse_topk_lens_offset) : (0)); + int _min_7 = ((_max_4) < (sparse_topk) ? (_max_4) : (sparse_topk)); + active_topk = _min_7; + } + int num_steps = (active_topk + 128 - 1) / 128; + float row_max_pair[2]; + float row_sum_pair[2]; + row_max_pair[0] = -CAKE_INF; + row_max_pair[1] = -CAKE_INF; + row_sum_pair[0] = 0.0f; + row_sum_pair[1] = 0.0f; + asm volatile("barrier.sync 1, 128;" ::: "memory"); + uint32_t _amf_u_0 = __float_as_uint(-3.4028235e+38f); + uint32_t _amf_mask_0 = -int32_t(_amf_u_0 >> 31) | 0x80000000u; + unsigned int _amf_enc_0 = _amf_u_0 ^ _amf_mask_0; + if (wg_tid < 8) { + smem_softmax_u32[wg_tid] = _amf_enc_0; + } + asm volatile("barrier.sync 1, 128;" ::: "memory"); + #pragma unroll 1 + for (int tile = 0; tile < num_steps; tile++) { + mbarrier_wait(score_full_addr + (score_stage) * 8, score_phase); + int score_col = ((score_stage == 0) ? 0 : 8); + float scores[8]; + float score_lo[4]; + float score_hi[4]; + int tile_end = (tile + 1) * 128; + int is_full_tile = ((active_topk % 128 == 0) ? 1 : 0); + is_full_tile = is_full_tile | ((tile_end < active_topk) ? 1 : 0); + if (is_full_tile != 0) { + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x1.b32" + " {%0, %1, %2, %3}, [%4];" + : "=r"(*reinterpret_cast(&score_lo[0])), "=r"(*reinterpret_cast(&score_lo[1])), "=r"(*reinterpret_cast(&score_lo[2])), "=r"(*reinterpret_cast(&score_lo[3])) + : "r"(taddr + (unsigned int)score_col)); + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x1.b32" + " {%0, %1, %2, %3}, [%4];" + : "=r"(*reinterpret_cast(&score_hi[0])), "=r"(*reinterpret_cast(&score_hi[1])), "=r"(*reinterpret_cast(&score_hi[2])), "=r"(*reinterpret_cast(&score_hi[3])) + : "r"(taddr + (unsigned int)score_col + 1048576)); + #pragma unroll + for (int c = 0; c < 4; c++) { + scores[c] = score_lo[c]; + scores[c + 4] = score_hi[c]; + } + } else { + asm volatile("" ::: "memory"); + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x1.b32" + " {%0, %1, %2, %3}, [%4];" + : "=r"(*reinterpret_cast(&score_lo[0])), "=r"(*reinterpret_cast(&score_lo[1])), "=r"(*reinterpret_cast(&score_lo[2])), "=r"(*reinterpret_cast(&score_lo[3])) + : "r"(taddr + (unsigned int)score_col)); + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x1.b32" + " {%0, %1, %2, %3}, [%4];" + : "=r"(*reinterpret_cast(&score_hi[0])), "=r"(*reinterpret_cast(&score_hi[1])), "=r"(*reinterpret_cast(&score_hi[2])), "=r"(*reinterpret_cast(&score_hi[3])) + : "r"(taddr + (unsigned int)score_col + 1048576)); + #pragma unroll + for (int c_1 = 0; c_1 < 4; c_1++) { + scores[c_1] = score_lo[c_1]; + scores[c_1 + 4] = score_hi[c_1]; + } + if (active_topk <= tile * 128 + ldtm_row_base) { + scores[0] = -CAKE_INF; + scores[1] = -CAKE_INF; + } + if (active_topk <= tile * 128 + ldtm_row_base + 8) { + scores[2] = -CAKE_INF; + scores[3] = -CAKE_INF; + } + if (active_topk <= tile * 128 + ldtm_row_base + 16) { + scores[4] = -CAKE_INF; + scores[5] = -CAKE_INF; + } + if (active_topk <= tile * 128 + ldtm_row_base + 24) { + scores[6] = -CAKE_INF; + scores[7] = -CAKE_INF; + } + asm volatile("" ::: "memory"); + } + float pair_max[2]; + float _max_5 = max_noftz(scores[0], scores[2]); + float _max_6 = max_noftz(scores[4], scores[6]); + float _max_7 = max_noftz(_max_5, _max_6); + pair_max[0] = _max_7; + float _max_8 = max_noftz(scores[1], scores[3]); + float _max_9 = max_noftz(scores[5], scores[7]); + float _max_10 = max_noftz(_max_8, _max_9); + pair_max[1] = _max_10; + #pragma unroll + for (int c_2 = 0; c_2 < 2; c_2++) { + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, pair_max[c_2], 16); + float _max_11 = max_noftz(pair_max[c_2], _shfl_xor_0); + pair_max[c_2] = _max_11; + float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, pair_max[c_2], 8); + float _max_12 = max_noftz(pair_max[c_2], _shfl_xor_1); + pair_max[c_2] = _max_12; + } + float old_max_pair[2]; + float new_max_pair[2]; + #pragma unroll + for (int c_3 = 0; c_3 < 2; c_3++) { + old_max_pair[c_3] = row_max_pair[c_3]; + float _max_13 = max_noftz(row_max_pair[c_3], pair_max[c_3]); + new_max_pair[c_3] = _max_13; + } + if (lane < 8) { + uint32_t _amf_u_1 = __float_as_uint(new_max_pair[0]); + uint32_t _amf_mask_1 = -int32_t(_amf_u_1 >> 31) | 0x80000000u; + unsigned int _amf_enc_1 = _amf_u_1 ^ _amf_mask_1; + uint32_t _amf_u_2 = __float_as_uint(new_max_pair[1]); + uint32_t _amf_mask_2 = -int32_t(_amf_u_2 >> 31) | 0x80000000u; + unsigned int _amf_enc_2 = _amf_u_2 ^ _amf_mask_2; + atomicMax(&smem_softmax_u32[col_pair_base], _amf_enc_1); + atomicMax(&smem_softmax_u32[col_pair_base + 1], _amf_enc_2); + } + asm volatile("barrier.sync 1, 128;" ::: "memory"); + uint32_t _amf_u_3 = smem_softmax_u32[col_pair_base]; + uint32_t _amf_mask_3 = ((_amf_u_3 >> 31) - 1u) | 0x80000000u; + float _amf_dec_0 = __uint_as_float(_amf_u_3 ^ _amf_mask_3); + new_max_pair[0] = _amf_dec_0; + uint32_t _amf_u_4 = smem_softmax_u32[col_pair_base + 1]; + uint32_t _amf_mask_4 = ((_amf_u_4 >> 31) - 1u) | 0x80000000u; + float _amf_dec_1 = __uint_as_float(_amf_u_4 ^ _amf_mask_4); + new_max_pair[1] = _amf_dec_1; + mbarrier_wait(stats_empty_addr + (stats_stage) * 8, stats_phase); + float stats_pair[4]; + stats_pair[0] = old_max_pair[0]; + stats_pair[1] = old_max_pair[1]; + stats_pair[2] = new_max_pair[0]; + stats_pair[3] = new_max_pair[1]; + int stats_col = ((stats_stage == 0) ? 16 : 48); + tmem_st_x4_f32(taddr + (unsigned int)stats_col, stats_pair); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(stats_full_addr + (stats_stage) * 8); + stats_stage += 1; + if (stats_stage == 2) { stats_stage = 0; stats_phase ^= 1; } + float acc_scale_pair[2]; + #pragma unroll + for (int c_4 = 0; c_4 < 2; c_4++) { + float delta = softmax_scale_log2 * (old_max_pair[c_4] - new_max_pair[c_4]); + float _exp2_0 = approx_exp2(delta); + acc_scale_pair[c_4] = ((old_max_pair[c_4] > -CAKE_INF) ? _exp2_0 : 1.0f); + row_max_pair[c_4] = new_max_pair[c_4]; + } + asm volatile("" ::: "memory"); + float exp_values[8]; + #pragma unroll + for (int c_5 = 0; c_5 < 8; c_5++) { + float safe_max = ((new_max_pair[c_5 % 2] == -CAKE_INF) ? 0.0f : new_max_pair[c_5 % 2]); + float _exp2_1 = approx_exp2(scores[c_5] * softmax_scale_log2 - safe_max * softmax_scale_log2); + exp_values[c_5] = _exp2_1; + } + mbarrier_wait(p_empty_addr + (p_stage) * 8, p_phase); + float p_values[8]; + #pragma unroll + for (int c_6 = 0; c_6 < 8; c_6++) { + p_values[c_6] = exp_values[c_6] * 448.0f; + } + unsigned int regs_p[2]; + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(p_values[0]), "f"(p_values[1]), + "f"(p_values[2]), "f"(p_values[3])); + regs_p[0] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(p_values[4]), "f"(p_values[5]), + "f"(p_values[6]), "f"(p_values[7])); + regs_p[1] = _packed; + } + int mtx_idx = lane / 8; + int thr_row_idx = lane % 8; + int seg_col_idx = warp_in_wg * 2 + mtx_idx ^ thr_row_idx; + int stsm_offset = thr_row_idx * 128 + seg_col_idx * 16; + const void* _stmatrix_b8_ptr_5 = reinterpret_cast(reinterpret_cast(smem_p) + (p_stage * 1024 + (unsigned int)stsm_offset)); + uint64_t _stmatrix_b8_addr64_5; + asm volatile("cvta.to.shared.u64 %0, %1;" : "=l"(_stmatrix_b8_addr64_5) : "l"(_stmatrix_b8_ptr_5)); + uint32_t _stmatrix_b8_addr_5; + asm volatile("cvt.u32.u64 %0, %1;" : "=r"(_stmatrix_b8_addr_5) : "l"(_stmatrix_b8_addr64_5)); + asm volatile("stmatrix.sync.aligned.m16n8.x2.trans.shared.b8 [%0], {%1, %2};\n" + :: "r"(_stmatrix_b8_addr_5), "r"(regs_p[0]), "r"(regs_p[1]) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(p_full_addr + (p_stage) * 8); + p_stage += 1; + if (p_stage == 2) { p_stage = 0; p_phase ^= 1; } + float2 _f2_0 = make_float2(exp_values[0], exp_values[1]); + float2 _f2_1 = make_float2(exp_values[2], exp_values[3]); + float2 _f2_2 = make_float2(exp_values[4], exp_values[5]); + float2 _f2_3 = make_float2(exp_values[6], exp_values[7]); + float2 _f2_4 = make_float2(row_sum_pair[0], row_sum_pair[1]); + float2 _f2_5 = make_float2(acc_scale_pair[0], acc_scale_pair[1]); + float2 sum_p0_pair = fma_f32x2_rn_ftz(_f2_4, _f2_5, _f2_0); + float2 sum_p1_pair = add_f32x2(sum_p0_pair, _f2_1); + float2 sum_p2_pair = add_f32x2(sum_p1_pair, _f2_2); + float2 next_sum_pair = add_f32x2(sum_p2_pair, _f2_3); + row_sum_pair[0] = next_sum_pair.x; + row_sum_pair[1] = next_sum_pair.y; + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + mbarrier_arrive(score_empty_addr + (score_stage) * 8); + score_stage += 1; + if (score_stage == 2) { score_stage = 0; score_phase ^= 1; } + } + if (num_steps > 0) { + mbarrier_wait(stats_empty_addr + (stats_stage) * 8, stats_phase); + float final_stats_pair[4]; + final_stats_pair[0] = row_sum_pair[0]; + final_stats_pair[1] = row_sum_pair[1]; + final_stats_pair[2] = row_max_pair[0]; + final_stats_pair[3] = row_max_pair[1]; + int final_stats_col = ((stats_stage == 0) ? 16 : 48); + tmem_st_x4_f32(taddr + (unsigned int)final_stats_col, final_stats_pair); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(stats_full_addr + (stats_stage) * 8); + stats_stage += 1; + if (stats_stage == 2) { stats_stage = 0; stats_phase ^= 1; } + } + mbarrier_wait(work_full_addr + (work_stage) * 8, _phase_work_full); + uint32_t _clc_valid_3 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_3) + : "r"(work_response_addr + work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_9 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_9) + : "r"(work_response_addr + work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_10 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_10) + : "r"(work_response_addr + work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_11 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_11) + : "r"(work_response_addr + work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + work_stage * 8), "r"(0) : "memory"); + work_stage += 1; + if (work_stage == 2) { work_stage = 0; _phase_work_full ^= 1; } + if (_clc_valid_3 == 0) { + break; + } + work_x = _clc_ctaid_9; + work_y = _clc_ctaid_10; + work_z = _clc_ctaid_11; + } + } + } + // ---- Role: correction ---- + if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 136;"); + { // correction_main + unsigned int work_stage_1 = 0; + unsigned int stats_stage_1 = 0; + unsigned int stats_phase_1 = 0; + unsigned int work_x_1 = blockIdx.x; + unsigned int work_y_1 = blockIdx.y; + unsigned int work_z_1 = blockIdx.z; + const int corr_rank = warp - 4; + const int corr_tid = corr_rank * 32 + lane; + const int col_pair_base_1 = corr_tid % 4 * 2; + float softmax_scale_log2_1 = bmm1_scale[0] * 1.4426950408889634f; + float output_scale = bmm2_scale[0]; + unsigned int _phase_o_full_0 = 0; + unsigned int _phase_work_full_1 = 0; + #pragma unroll 1 + for (unsigned int _work_1 = 0; _work_1 < 120; _work_1++) { + int q_begin_1 = cum_seq_lens_q[work_z_1]; + int q_end_1 = cum_seq_lens_q[work_z_1 + 1]; + int q_len_1 = q_end_1 - q_begin_1; + int query_idx_1 = (unsigned int)q_begin_1 + work_x_1; + int active_topk_1 = 0; + if (work_x_1 < (unsigned int)q_len_1) { + int _max_14 = ((sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) : (0)); + int _min_8 = ((_max_14) < (sparse_topk) ? (_max_14) : (sparse_topk)); + active_topk_1 = _min_8; + } + int num_steps_1 = (active_topk_1 + 128 - 1) / 128; + int head_group = work_y_1 / 4; + int value_quarter = work_y_1 % 4; + if (num_steps_1 > 0) { + mbarrier_wait(stats_full_addr + (stats_stage_1) * 8, stats_phase_1); + mbarrier_arrive(stats_empty_addr + (stats_stage_1) * 8); + stats_stage_1 += 1; + if (stats_stage_1 == 2) { stats_stage_1 = 0; stats_phase_1 ^= 1; } + } + #pragma unroll 1 + for (int tile_1 = 1; tile_1 < num_steps_1; tile_1++) { + mbarrier_wait(stats_full_addr + (stats_stage_1) * 8, stats_phase_1); + int stats_col_1 = ((stats_stage_1 == 0) ? 16 : 48); + float _tmem_load_0[4]; + tmem_ld_x4(&_tmem_load_0[0], taddr + (unsigned int)stats_col_1); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + float scale_pair[2]; + #pragma unroll + for (int c_7 = 0; c_7 < 2; c_7++) { + float max_diff = _tmem_load_0[c_7] - _tmem_load_0[c_7 + 2]; + float _exp2_2 = approx_exp2(softmax_scale_log2_1 * max_diff); + scale_pair[c_7] = ((max_diff != 0.0f) ? _exp2_2 : 1.0f); + } + mbarrier_arrive(stats_empty_addr + (stats_stage_1) * 8); + stats_stage_1 += 1; + if (stats_stage_1 == 2) { stats_stage_1 = 0; stats_phase_1 ^= 1; } + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int needs_scale = ((scale_pair[0] != 1.0f) ? 1 : 0); + needs_scale = needs_scale | ((scale_pair[1] != 1.0f) ? 1 : 0); + int _vote_0 = __any_sync(0xFFFFFFFF, needs_scale != 0); + needs_scale = _vote_0; + if (needs_scale != 0) { + float values_lo[4]; + float values_hi[4]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x1.b32" + " {%0, %1, %2, %3}, [%4];" + : "=r"(*reinterpret_cast(&values_lo[0])), "=r"(*reinterpret_cast(&values_lo[1])), "=r"(*reinterpret_cast(&values_lo[2])), "=r"(*reinterpret_cast(&values_lo[3])) + : "r"(taddr + 80)); + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x1.b32" + " {%0, %1, %2, %3}, [%4];" + : "=r"(*reinterpret_cast(&values_hi[0])), "=r"(*reinterpret_cast(&values_hi[1])), "=r"(*reinterpret_cast(&values_hi[2])), "=r"(*reinterpret_cast(&values_hi[3])) + : "r"(taddr + 80 + 1048576)); + #pragma unroll + for (int pair = 0; pair < 4; pair++) { + asm volatile("" ::: "memory"); + values_lo[pair] = values_lo[pair] * scale_pair[pair % 2]; + values_hi[pair] = values_hi[pair] * scale_pair[pair % 2]; + } + asm volatile( + "tcgen05.st.sync.aligned.16x256b.x1.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(taddr + 80), "r"(*reinterpret_cast(&values_lo[0])), "r"(*reinterpret_cast(&values_lo[1])), "r"(*reinterpret_cast(&values_lo[2])), "r"(*reinterpret_cast(&values_lo[3]))); + asm volatile( + "tcgen05.st.sync.aligned.16x256b.x1.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(taddr + 80 + 1048576), "r"(*reinterpret_cast(&values_hi[0])), "r"(*reinterpret_cast(&values_hi[1])), "r"(*reinterpret_cast(&values_hi[2])), "r"(*reinterpret_cast(&values_hi[3]))); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + } + mbarrier_arrive(o_empty_addr); + } + if (num_steps_1 > 0) { + mbarrier_wait(stats_full_addr + (stats_stage_1) * 8, stats_phase_1); + int final_stats_col_1 = ((stats_stage_1 == 0) ? 16 : 48); + float _tmem_load_1[4]; + tmem_ld_x4(&_tmem_load_1[0], taddr + (unsigned int)final_stats_col_1); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + mbarrier_arrive(stats_empty_addr + (stats_stage_1) * 8); + stats_stage_1 += 1; + if (stats_stage_1 == 2) { stats_stage_1 = 0; stats_phase_1 ^= 1; } + float pair_sum[2]; + pair_sum[0] = _tmem_load_1[0]; + pair_sum[1] = _tmem_load_1[1]; + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, pair_sum[0], 16); + pair_sum[0] = pair_sum[0] + _shfl_xor_2; + float _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, pair_sum[1], 16); + pair_sum[1] = pair_sum[1] + _shfl_xor_3; + float _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, pair_sum[0], 8); + pair_sum[0] = pair_sum[0] + _shfl_xor_4; + float _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, pair_sum[1], 8); + pair_sum[1] = pair_sum[1] + _shfl_xor_5; + float _shfl_xor_6 = __shfl_xor_sync(0xFFFFFFFF, pair_sum[0], 4); + pair_sum[0] = pair_sum[0] + _shfl_xor_6; + float _shfl_xor_7 = __shfl_xor_sync(0xFFFFFFFF, pair_sum[1], 4); + pair_sum[1] = pair_sum[1] + _shfl_xor_7; + if (lane < 4) { + smem_corr[corr_rank * 8 + col_pair_base_1] = pair_sum[0]; + smem_corr[corr_rank * 8 + col_pair_base_1 + 1] = pair_sum[1]; + } + asm volatile("barrier.sync 2, 128;" ::: "memory"); + float sum0 = smem_corr[col_pair_base_1] + smem_corr[8 + col_pair_base_1] + smem_corr[16 + col_pair_base_1] + smem_corr[24 + col_pair_base_1]; + float sum1 = smem_corr[col_pair_base_1 + 1] + smem_corr[8 + col_pair_base_1 + 1] + smem_corr[16 + col_pair_base_1 + 1] + smem_corr[24 + col_pair_base_1 + 1]; + if (has_sinks != 0) { + float _exp2_3 = approx_exp2(sinks[head_group * 8 + col_pair_base_1] * 1.4426950408889634f - _tmem_load_1[2] * softmax_scale_log2_1); + sum0 = sum0 + _exp2_3; + float _exp2_4 = approx_exp2(sinks[head_group * 8 + col_pair_base_1 + 1] * 1.4426950408889634f - _tmem_load_1[3] * softmax_scale_log2_1); + sum1 = sum1 + _exp2_4; + } + float _rcp_0 = approx_rcp(sum0); + float scale0 = ((sum0 > 0.0f) ? output_scale * 0.002232142857142857f * _rcp_0 : 0.0f); + float _rcp_1 = approx_rcp(sum1); + float scale1 = ((sum1 > 0.0f) ? output_scale * 0.002232142857142857f * _rcp_1 : 0.0f); + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + float output_lo[4]; + float output_hi[4]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x1.b32" + " {%0, %1, %2, %3}, [%4];" + : "=r"(*reinterpret_cast(&output_lo[0])), "=r"(*reinterpret_cast(&output_lo[1])), "=r"(*reinterpret_cast(&output_lo[2])), "=r"(*reinterpret_cast(&output_lo[3])) + : "r"(taddr + 80)); + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x1.b32" + " {%0, %1, %2, %3}, [%4];" + : "=r"(*reinterpret_cast(&output_hi[0])), "=r"(*reinterpret_cast(&output_hi[1])), "=r"(*reinterpret_cast(&output_hi[2])), "=r"(*reinterpret_cast(&output_hi[3])) + : "r"(taddr + 80 + 1048576)); + unsigned int packed_o[4]; + #pragma unroll + for (int pair_1 = 0; pair_1 < 4; pair_1++) { + asm volatile("" ::: "memory"); + output_lo[pair_1] = output_lo[pair_1] * ((pair_1 % 2 == 0) ? scale0 : scale1); + output_hi[pair_1] = output_hi[pair_1] * ((pair_1 % 2 == 0) ? scale0 : scale1); + } + float output_values[8]; + #pragma unroll + for (int c_8 = 0; c_8 < 4; c_8++) { + output_values[c_8] = output_lo[c_8]; + output_values[c_8 + 4] = output_hi[c_8]; + } + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(output_values[_lp*2 + 0], output_values[_lp*2+1 + 0])); + packed_o[_lp] = *(uint32_t*)&_bf2; + } + int slice_idx = corr_rank / 2; + int warp_idx_in_slice = corr_rank % 2; + int mtx_idx_1 = lane / 8; + int thr_row_idx_1 = lane % 8; + int seg_col_idx_1 = warp_idx_in_slice * 4 + mtx_idx_1 ^ thr_row_idx_1; + int stsm_offset_1 = slice_idx * 8 * 128 + thr_row_idx_1 * 128 + seg_col_idx_1 * 16; + asm volatile("stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 [%0], {%1, %2, %3, %4};\n" + :: "l"(__cvta_generic_to_shared(reinterpret_cast(smem_o) + stsm_offset_1)), "r"(*reinterpret_cast(&packed_o[0])), "r"(*reinterpret_cast(&packed_o[1])), "r"(*reinterpret_cast(&packed_o[2])), "r"(*reinterpret_cast(&packed_o[3])) + : "memory"); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + int copy_base = corr_tid * 16; + int smem_row = copy_base / 128; + int copy_smem_offset = copy_base ^ smem_row % 8 * 16; + int copy_row = smem_row % 8; + int copy_col = (smem_row / 8 * 128 + copy_base % 128) / 2; + unsigned int copy_vec[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(©_vec[0])), "=r"(*reinterpret_cast(©_vec[(0) + 1])), "=r"(*reinterpret_cast(©_vec[(0) + 2])), "=r"(*reinterpret_cast(©_vec[(0) + 3])) + : "r"(smem_o_addr + (unsigned int)copy_smem_offset)); + int head_groups = gridDim.y / 4; + int output_base = query_idx_1 * head_groups * 8 * 512 + head_group * 8 * 512 + value_quarter * 128; + reinterpret_cast(O + (output_base + copy_row * 512 + copy_col))[0] = reinterpret_cast(copy_vec)[0]; + mbarrier_arrive(o_empty_addr); + } + mbarrier_wait(work_full_addr + (work_stage_1) * 8, _phase_work_full_1); + uint32_t _clc_valid_4 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_4) + : "r"(work_response_addr + work_stage_1 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_12 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_12) + : "r"(work_response_addr + work_stage_1 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_13 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_13) + : "r"(work_response_addr + work_stage_1 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_14 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_14) + : "r"(work_response_addr + work_stage_1 * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + work_stage_1 * 8), "r"(0) : "memory"); + work_stage_1 += 1; + if (work_stage_1 == 2) { work_stage_1 = 0; _phase_work_full_1 ^= 1; } + if (_clc_valid_4 == 0) { + break; + } + work_x_1 = _clc_ctaid_12; + work_y_1 = _clc_ctaid_13; + work_z_1 = _clc_ctaid_14; + } + asm volatile("barrier.sync 11, 128;" ::: "memory"); + if (corr_rank == 0) { + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + } + } + // ---- Role: mma ---- + if (warp == 8) { + { // mma_main + unsigned int work_stage_2 = 0; + unsigned int q_stage = 0; + unsigned int q_phase = 0; + unsigned int kv_stage = 0; + unsigned int kv_phase = 0; + unsigned int score_stage_1 = 0; + unsigned int score_phase_1 = 1; + unsigned int p_stage_1 = 0; + unsigned int p_phase_1 = 0; + unsigned int work_x_2 = blockIdx.x; + unsigned int work_y_2 = blockIdx.y; + unsigned int work_z_2 = blockIdx.z; + unsigned int _phase_o_empty_0 = 1; + unsigned int _phase_work_full_2 = 0; + #pragma unroll 1 + for (unsigned int _work_2 = 0; _work_2 < 120; _work_2++) { + int q_begin_2 = cum_seq_lens_q[work_z_2]; + int q_end_2 = cum_seq_lens_q[work_z_2 + 1]; + int q_len_2 = q_end_2 - q_begin_2; + int query_idx_2 = (unsigned int)q_begin_2 + work_x_2; + int active_topk_2 = 0; + if (work_x_2 < (unsigned int)q_len_2) { + int _max_3 = ((sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) : (0)); + int _min_6 = ((_max_3) < (sparse_topk) ? (_max_3) : (sparse_topk)); + active_topk_2 = _min_6; + } + int num_steps_2 = (active_topk_2 + 128 - 1) / 128; + int value_quarter_1 = work_y_2 % 4; + if (num_steps_2 > 0) { + mbarrier_wait(q_full_addr + (q_stage) * 8, q_phase); + int first_output = 1; + mbarrier_wait(score_empty_addr + (score_stage_1) * 8, score_phase_1); + int score_col_1 = ((score_stage_1 == 0) ? 0 : 8); + #pragma unroll + for (int hs = 0; hs < 3; hs++) { + int hs_sel = ((value_quarter_1 > hs) ? hs : hs + 1); + mbarrier_wait(kv_full_addr + (kv_stage) * 8, kv_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_0 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage) * 1024); + int _mma_b_lo_0 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (hs_sel) * 64); + { + uint64_t _mma_ss_a_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_0); + uint64_t _mma_ss_b_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_0); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134348816, ((hs == 0) ? 0 : 1)); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134348816, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134348816, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134348816, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage) * 8); + kv_stage += 1; + if (kv_stage == 8) { kv_stage = 0; kv_phase ^= 1; } + } + mbarrier_wait(kv_full_addr + (kv_stage) * 8, kv_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_1 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage) * 1024); + int _mma_b_lo_1 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (value_quarter_1) * 64); + { + uint64_t _mma_ss_a_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_1); + uint64_t _mma_ss_b_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_1); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134348816, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134348816, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134348816, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134348816, 1); + } + } + kv_stage += 1; + if (kv_stage == 8) { kv_stage = 0; kv_phase ^= 1; } + elect_commit(score_full_addr + (score_stage_1) * 8); + score_stage_1 += 1; + if (score_stage_1 == 2) { score_stage_1 = 0; score_phase_1 ^= 1; } + #pragma unroll 1 + for (int _tile = 0; _tile < num_steps_2 - 1; _tile++) { + mbarrier_wait(score_empty_addr + (score_stage_1) * 8, score_phase_1); + int score_col_0 = ((score_stage_1 == 0) ? 0 : 8); + #pragma unroll + for (int hs_1 = 0; hs_1 < 3; hs_1++) { + int hs_sel_1 = ((value_quarter_1 > hs_1) ? hs_1 : hs_1 + 1); + mbarrier_wait(kv_full_addr + (kv_stage) * 8, kv_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_2 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage) * 1024); + int _mma_b_lo_2 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (hs_sel_1) * 64); + { + uint64_t _mma_ss_a_desc_2 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_2); + uint64_t _mma_ss_b_desc_2 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_2); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134348816, ((hs_1 == 0) ? 0 : 1)); + } + incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134348816, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134348816, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134348816, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage) * 8); + kv_stage += 1; + if (kv_stage == 8) { kv_stage = 0; kv_phase ^= 1; } + } + mbarrier_wait(kv_full_addr + (kv_stage) * 8, kv_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_3 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage) * 1024); + int _mma_b_lo_3 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (value_quarter_1) * 64); + { + uint64_t _mma_ss_a_desc_3 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_3); + uint64_t _mma_ss_b_desc_3 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_3); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134348816, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134348816, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134348816, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134348816, 1); + } + } + kv_stage += 1; + if (kv_stage == 8) { kv_stage = 0; kv_phase ^= 1; } + elect_commit(score_full_addr + (score_stage_1) * 8); + score_stage_1 += 1; + if (score_stage_1 == 2) { score_stage_1 = 0; score_phase_1 ^= 1; } + unsigned int v_cur = (kv_stage + 3) % 8; + mbarrier_wait(p_full_addr + (p_stage_1) * 8, p_phase_1); + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_4 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (v_cur) * 1024); + int _mma_b_lo_4 = make_warp_uniform(((((smem_p_addr) >> 4) & 0x3FFF) | 0x400000) + (p_stage_1) * 64); + { + uint64_t _mma_ss_a_desc_4 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_4); + uint64_t _mma_ss_b_desc_4 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_4); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134381584, ((first_output) ? 0 : 1)); + } + incr_smem_desc_lo(_mma_ss_a_desc_4, 256U); + incr_smem_desc_lo(_mma_ss_b_desc_4, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134381584, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_4, 256U); + incr_smem_desc_lo(_mma_ss_b_desc_4, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134381584, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_4, 256U); + incr_smem_desc_lo(_mma_ss_b_desc_4, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134381584, 1); + } + } + elect_commit(o_full_addr); + elect_commit(kv_empty_addr + (v_cur) * 8); + elect_commit(p_empty_addr + (p_stage_1) * 8); + p_stage_1 += 1; + if (p_stage_1 == 2) { p_stage_1 = 0; p_phase_1 ^= 1; } + first_output = 0; + } + unsigned int v_tail = (kv_stage + 7) % 8; + mbarrier_wait(p_full_addr + (p_stage_1) * 8, p_phase_1); + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_5 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (v_tail) * 1024); + int _mma_b_lo_5 = make_warp_uniform(((((smem_p_addr) >> 4) & 0x3FFF) | 0x400000) + (p_stage_1) * 64); + { + uint64_t _mma_ss_a_desc_5 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_5); + uint64_t _mma_ss_b_desc_5 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_5); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134381584, ((first_output) ? 0 : 1)); + } + incr_smem_desc_lo(_mma_ss_a_desc_5, 256U); + incr_smem_desc_lo(_mma_ss_b_desc_5, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134381584, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_5, 256U); + incr_smem_desc_lo(_mma_ss_b_desc_5, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134381584, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_5, 256U); + incr_smem_desc_lo(_mma_ss_b_desc_5, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134381584, 1); + } + } + elect_commit(o_full_addr); + elect_commit(kv_empty_addr + (v_tail) * 8); + elect_commit(p_empty_addr + (p_stage_1) * 8); + p_stage_1 += 1; + if (p_stage_1 == 2) { p_stage_1 = 0; p_phase_1 ^= 1; } + elect_commit(q_empty_addr + (q_stage) * 8); + q_phase ^= 1; + } + mbarrier_wait(work_full_addr + (work_stage_2) * 8, _phase_work_full_2); + uint32_t _clc_valid_2 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_2) + : "r"(work_response_addr + work_stage_2 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_6 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_6) + : "r"(work_response_addr + work_stage_2 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_7 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_7) + : "r"(work_response_addr + work_stage_2 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_8 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_8) + : "r"(work_response_addr + work_stage_2 * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + work_stage_2 * 8), "r"(0) : "memory"); + work_stage_2 += 1; + if (work_stage_2 == 2) { work_stage_2 = 0; _phase_work_full_2 ^= 1; } + if (_clc_valid_2 == 0) { + break; + } + work_x_2 = _clc_ctaid_6; + work_y_2 = _clc_ctaid_7; + work_z_2 = _clc_ctaid_8; + } + #pragma unroll + for (int tail = 0; tail < 2; tail++) { + elect_commit(score_full_addr + (tail) * 8); + } + } + } + // ---- Role: page_loader ---- + if (warp == 9) { + { // page_loader_main + unsigned int work_stage_3 = 0; + unsigned int page_stage = 0; + unsigned int page_phase = 1; + unsigned int work_x_3 = blockIdx.x; + unsigned int work_y_3 = blockIdx.y; + unsigned int work_z_3 = blockIdx.z; + unsigned int _phase_work_full_3 = 0; + #pragma unroll 1 + for (unsigned int _work_3 = 0; _work_3 < 120; _work_3++) { + int q_begin_3 = cum_seq_lens_q[work_z_3]; + int q_end_3 = cum_seq_lens_q[work_z_3 + 1]; + int q_len_3 = q_end_3 - q_begin_3; + int query_idx_3 = (unsigned int)q_begin_3 + work_x_3; + int active_topk_3 = 0; + if (work_x_3 < (unsigned int)q_len_3) { + int _max_0 = ((sparse_topk_lens[query_idx_3] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_3] + sparse_topk_lens_offset) : (0)); + int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); + active_topk_3 = _min_0; + } + int num_steps_3 = (active_topk_3 + 128 - 1) / 128; + int num_pairs = (num_steps_3 + 1) / 2; + int _max_1 = ((active_topk_3 - 1 & -4) > (0) ? (active_topk_3 - 1 & -4) : (0)); + int last_aligned = _max_1; + int* swa_row = swa_indices + (query_idx_3 * swa_index_stride); + int* compressed_row = compressed_indices + (query_idx_3 * compressed_index_stride); + #pragma unroll 1 + for (int pair_2 = 0; pair_2 < num_pairs; pair_2++) { + int lane_base = lane * 4; + #pragma unroll + for (int duplicate = 0; duplicate < 2; duplicate++) { + mbarrier_wait(page_empty_addr + (page_stage) * 8, page_phase); + int page0 = pair_2 * 2 * 128 + lane_base; + int page1 = page0 + 128; + int safe0 = ((page0 <= last_aligned) ? page0 : last_aligned); + int safe1 = ((page1 <= last_aligned) ? page1 : last_aligned); + int* src0 = ((safe0 < 128) ? swa_row + safe0 : compressed_row + (safe0 - 128)); + int* src1 = ((safe1 < 128) ? swa_row + safe1 : compressed_row + (safe1 - 128)); + asm volatile("cp.async.ca.shared::cta.global [%0], [%1], 16;" + :: "r"(smem_page_offsets_addr + page_stage * 1024 + (unsigned int)(lane_base * 4)), "l"(src0)); + asm volatile("cp.async.ca.shared::cta.global [%0], [%1], 16;" + :: "r"(smem_page_offsets_addr + page_stage * 1024 + (unsigned int)((128 + lane_base) * 4)), "l"(src1)); + asm volatile( + "{\n\t" + "cp.async.mbarrier.arrive.shared::cta.b64 [%0];\n\t" + "}" + :: "r"(page_full_addr + (page_stage) * 8) : "memory"); + mbarrier_arrive(page_full_addr + (page_stage) * 8); + page_stage += 1; + if (page_stage == 6) { page_stage = 0; page_phase ^= 1; } + } + } + mbarrier_wait(work_full_addr + (work_stage_3) * 8, _phase_work_full_3); + uint32_t _clc_valid_0 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_0) + : "r"(work_response_addr + work_stage_3 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_0 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_0) + : "r"(work_response_addr + work_stage_3 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_1 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_1) + : "r"(work_response_addr + work_stage_3 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_2 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_2) + : "r"(work_response_addr + work_stage_3 * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + work_stage_3 * 8), "r"(0) : "memory"); + work_stage_3 += 1; + if (work_stage_3 == 2) { work_stage_3 = 0; _phase_work_full_3 ^= 1; } + if (_clc_valid_0 == 0) { + break; + } + work_x_3 = _clc_ctaid_0; + work_y_3 = _clc_ctaid_1; + work_z_3 = _clc_ctaid_2; + } + } + } + // ---- Role: scheduler ---- + if (warp == 10) { + { // scheduler_main + asm volatile("griddepcontrol.wait;" ::: "memory"); + unsigned int work_stage_4 = 0; + unsigned int throttle_stage = 0; + unsigned int throttle_phase = 0; + unsigned int _phase_work_empty = 1; + unsigned int _phase_work_full_4 = 0; + #pragma unroll 1 + for (unsigned int _work_4 = 0; _work_4 < 120; _work_4++) { + mbarrier_wait(throttle_full_addr + (throttle_stage) * 8, throttle_phase); + mbarrier_arrive(throttle_empty_addr + (throttle_stage) * 8); + throttle_stage += 1; + if (throttle_stage == 2) { throttle_stage = 0; throttle_phase ^= 1; } + mbarrier_wait(work_empty_addr + (work_stage_4) * 8, _phase_work_empty); + if (elect_sync()) { + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [remAddr32], %2;\n\t" + "}" + :: "r"(work_full_addr + work_stage_4 * 8), "r"(0), "r"((uint32_t)(16)) : "memory"); + asm volatile( + "fence.proxy.async.shared::cta;\n\t" + "clusterlaunchcontrol.try_cancel.async.shared::cta" + ".mbarrier::complete_tx::bytes.b128" + " [%0], [%1];" + :: "r"(work_response_addr + work_stage_4 * 16 + 0 * 16), "r"(work_full_addr + work_stage_4 * 8) + : "memory"); + } + mbarrier_wait(work_full_addr + (work_stage_4) * 8, _phase_work_full_4); + uint32_t _clc_valid_5 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_5) + : "r"(work_response_addr + work_stage_4 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_15 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_15) + : "r"(work_response_addr + work_stage_4 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_16 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_16) + : "r"(work_response_addr + work_stage_4 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_17 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_17) + : "r"(work_response_addr + work_stage_4 * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + work_stage_4 * 8), "r"(0) : "memory"); + work_stage_4 += 1; + if (work_stage_4 == 2) { work_stage_4 = 0; _phase_work_empty ^= 1; _phase_work_full_4 ^= 1; } + if (_clc_valid_5 == 0) { + break; + } + } + } + } + // ---- Role: padding ---- + if (warp == 11) { + { // padding_main + unsigned int work_stage_5 = 0; + unsigned int work_x_4 = blockIdx.x; + unsigned int work_y_4 = blockIdx.y; + unsigned int work_z_4 = blockIdx.z; + unsigned int _phase_work_full_5 = 0; + #pragma unroll 1 + for (unsigned int _work_5 = 0; _work_5 < 120; _work_5++) { + mbarrier_wait(work_full_addr + (work_stage_5) * 8, _phase_work_full_5); + uint32_t _clc_valid_6 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_6) + : "r"(work_response_addr + work_stage_5 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_18 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_18) + : "r"(work_response_addr + work_stage_5 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_19 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_19) + : "r"(work_response_addr + work_stage_5 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_20 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_20) + : "r"(work_response_addr + work_stage_5 * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + work_stage_5 * 8), "r"(0) : "memory"); + work_stage_5 += 1; + if (work_stage_5 == 2) { work_stage_5 = 0; _phase_work_full_5 ^= 1; } + if (_clc_valid_6 == 0) { + break; + } + work_x_4 = _clc_ctaid_18; + work_y_4 = _clc_ctaid_19; + work_z_4 = _clc_ctaid_20; + } + } + } + // ---- Role: loader ---- + if (warp >= 12 && warp <= 15) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 88;"); + { // loader_main + asm volatile("griddepcontrol.wait;" ::: "memory"); + const int load_rank = warp - 12; + unsigned int work_stage_6 = 0; + unsigned int throttle_stage_1 = 0; + unsigned int throttle_phase_1 = 1; + unsigned int q_stage_1 = 0; + unsigned int q_phase_1 = 1; + unsigned int kv_stage_1 = 0; + unsigned int kv_phase_1 = 1; + unsigned int page_k_stage = 0; + unsigned int page_k_phase = 0; + unsigned int page_v_stage = 1; + unsigned int page_v_phase = 0; + unsigned int work_x_5 = blockIdx.x; + unsigned int work_y_5 = blockIdx.y; + unsigned int work_z_5 = blockIdx.z; + unsigned int _phase_work_full_6 = 0; + #pragma unroll 1 + for (unsigned int _work_6 = 0; _work_6 < 120; _work_6++) { + int q_begin_4 = cum_seq_lens_q[work_z_5]; + int q_end_4 = cum_seq_lens_q[work_z_5 + 1]; + int q_len_4 = q_end_4 - q_begin_4; + int query_idx_4 = (unsigned int)q_begin_4 + work_x_5; + int active_topk_4 = 0; + if (work_x_5 < (unsigned int)q_len_4) { + int _max_2 = ((sparse_topk_lens[query_idx_4] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_4] + sparse_topk_lens_offset) : (0)); + int _min_1 = ((_max_2) < (sparse_topk) ? (_max_2) : (sparse_topk)); + active_topk_4 = _min_1; + } + int num_steps_4 = (active_topk_4 + 128 - 1) / 128; + int head_group_1 = work_y_5 / 4; + int value_quarter_2 = work_y_5 % 4; + mbarrier_wait(throttle_empty_addr + (throttle_stage_1) * 8, throttle_phase_1); + mbarrier_arrive(throttle_full_addr + (throttle_stage_1) * 8); + throttle_stage_1 += 1; + if (throttle_stage_1 == 2) { throttle_stage_1 = 0; throttle_phase_1 ^= 1; } + if (num_steps_4 > 0) { + mbarrier_wait(q_empty_addr + (q_stage_1) * 8, q_phase_1); + if (load_rank == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(q_full_addr + (q_stage_1) * 8, 4096); + #pragma unroll + for (int hs_2 = 0; hs_2 < 4; hs_2++) { + tma_4d_gmem2smem(smem_q_addr + (unsigned int)(hs_2 * 1024), tmap_q, 0, head_group_1 * 8, hs_2, query_idx_4, q_full_addr + (q_stage_1) * 8); + } + } + } + mbarrier_wait(page_full_addr + (page_k_stage) * 8, page_k_phase); + int page_base = 0; + int warps_tile_idx = 0; + int _min_2 = ((active_topk_4) < (128) ? (active_topk_4) : (128)); + int warps_valid_rows = _min_2; + int _min_3 = (((warps_valid_rows + 3) / 4) < (4) ? ((warps_valid_rows + 3) / 4) : (4)); + int warps_tile = _min_3; + int warps_ge1 = 4; + int warps_ge2 = 4; + if (warps_tile_idx >= 1) { + warps_ge1 = warps_tile; + } + if (warps_tile_idx >= 2) { + warps_ge2 = warps_tile; + } + #pragma unroll + for (int hs_3 = 0; hs_3 < 4; hs_3++) { + int hs_sel_2 = ((hs_3 == 3) ? value_quarter_2 : ((value_quarter_2 > hs_3) ? hs_3 : hs_3 + 1)); + mbarrier_wait(kv_empty_addr + (kv_stage_1) * 8, kv_phase_1); + if (load_rank == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_1) * 8, ((hs_3 >= 4) ? warps_ge1 : warps_ge2) * 4096); + } + } + asm volatile("barrier.sync 9, 128;" ::: "memory"); + if (load_rank < ((hs_3 >= 4) ? warps_ge1 : warps_ge2)) { + int rows_all[32]; + #pragma unroll + for (int group = 0; group < 8; group++) { + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&rows_all[group * 4])), "=r"(*reinterpret_cast(&rows_all[(group * 4) + 1])), "=r"(*reinterpret_cast(&rows_all[(group * 4) + 2])), "=r"(*reinterpret_cast(&rows_all[(group * 4) + 3])) + : "r"(smem_page_offsets_addr + page_k_stage * 1024 + (unsigned int)((page_base + load_rank * 4 + group * 16) * 4))); + } + if (elect_sync()) { + { + #pragma unroll + for (int group_1 = 0; group_1 < 8; group_1++) { + const int row_base = load_rank * 4 + group_1 * 16; + int dst_k = smem_kv_addr + kv_stage_1 * 16384 + (unsigned int)(row_base * 128); + tma_gather4_gmem2smem(dst_k, tmap_swa_kv, hs_sel_2 * 128, rows_all[group_1 * 4], rows_all[group_1 * 4 + 1], rows_all[group_1 * 4 + 2], rows_all[group_1 * 4 + 3], kv_full_addr + (kv_stage_1) * 8); + } + } + } + } + kv_stage_1 += 1; + if (kv_stage_1 == 8) { kv_stage_1 = 0; kv_phase_1 ^= 1; } + } + #pragma unroll 1 + for (int tile_2 = 0; tile_2 < num_steps_4 - 1; tile_2++) { + int next_tile = tile_2 + 1; + if ((next_tile & 1) == 0) { + mbarrier_wait(page_full_addr + (page_k_stage) * 8, page_k_phase); + } + int page_base_0 = (next_tile & 1) * 128; + int warps_tile_idx_1 = next_tile; + int _min_4 = ((active_topk_4 - next_tile * 128) < (128) ? (active_topk_4 - next_tile * 128) : (128)); + int warps_valid_rows_2 = _min_4; + int _min_5 = (((warps_valid_rows_2 + 3) / 4) < (4) ? ((warps_valid_rows_2 + 3) / 4) : (4)); + int warps_tile_3 = _min_5; + int warps_ge1_4 = 4; + int warps_ge2_5 = 4; + if (warps_tile_idx_1 >= 1) { + warps_ge1_4 = warps_tile_3; + } + if (warps_tile_idx_1 >= 2) { + warps_ge2_5 = warps_tile_3; + } + #pragma unroll + for (int hs_4 = 0; hs_4 < 4; hs_4++) { + int hs_sel_3 = ((hs_4 == 3) ? value_quarter_2 : ((value_quarter_2 > hs_4) ? hs_4 : hs_4 + 1)); + mbarrier_wait(kv_empty_addr + (kv_stage_1) * 8, kv_phase_1); + if (load_rank == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_1) * 8, ((hs_4 >= 4) ? warps_ge1_4 : warps_ge2_5) * 4096); + } + } + asm volatile("barrier.sync 9, 128;" ::: "memory"); + if (load_rank < ((hs_4 >= 4) ? warps_ge1_4 : warps_ge2_5)) { + int rows_all_1[32]; + #pragma unroll + for (int group_2 = 0; group_2 < 8; group_2++) { + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&rows_all_1[group_2 * 4])), "=r"(*reinterpret_cast(&rows_all_1[(group_2 * 4) + 1])), "=r"(*reinterpret_cast(&rows_all_1[(group_2 * 4) + 2])), "=r"(*reinterpret_cast(&rows_all_1[(group_2 * 4) + 3])) + : "r"(smem_page_offsets_addr + page_k_stage * 1024 + (unsigned int)((page_base_0 + load_rank * 4 + group_2 * 16) * 4))); + } + if (elect_sync()) { + if (next_tile == 0) { + #pragma unroll + for (int group_3 = 0; group_3 < 8; group_3++) { + const int row_base_1 = load_rank * 4 + group_3 * 16; + int dst_k_1 = smem_kv_addr + kv_stage_1 * 16384 + (unsigned int)(row_base_1 * 128); + tma_gather4_gmem2smem(dst_k_1, tmap_swa_kv, hs_sel_3 * 128, rows_all_1[group_3 * 4], rows_all_1[group_3 * 4 + 1], rows_all_1[group_3 * 4 + 2], rows_all_1[group_3 * 4 + 3], kv_full_addr + (kv_stage_1) * 8); + } + } else { + #pragma unroll + for (int group_4 = 0; group_4 < 8; group_4++) { + const int row_base_2 = load_rank * 4 + group_4 * 16; + int dst_k_2 = smem_kv_addr + kv_stage_1 * 16384 + (unsigned int)(row_base_2 * 128); + tma_gather4_gmem2smem(dst_k_2, tmap_compressed_kv, hs_sel_3 * 128, rows_all_1[group_4 * 4], rows_all_1[group_4 * 4 + 1], rows_all_1[group_4 * 4 + 2], rows_all_1[group_4 * 4 + 3], kv_full_addr + (kv_stage_1) * 8); + } + } + } + } + kv_stage_1 += 1; + if (kv_stage_1 == 8) { kv_stage_1 = 0; kv_phase_1 ^= 1; } + } + if ((next_tile & 1) != 0 && next_tile < num_steps_4 - 1) { + mbarrier_arrive(page_empty_addr + (page_k_stage) * 8); + page_k_stage += 1; + if (page_k_stage == 6) { page_k_stage = 0; page_k_phase ^= 1; } + page_k_stage += 1; + if (page_k_stage == 6) { page_k_stage = 0; page_k_phase ^= 1; } + } + if ((tile_2 & 1) == 0) { + mbarrier_wait(page_full_addr + (page_v_stage) * 8, page_v_phase); + } + if ((tile_2 & 1) != 0) { + mbarrier_arrive(page_empty_addr + (page_v_stage) * 8); + page_v_stage += 1; + if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } + page_v_stage += 1; + if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } + } + } + int last_tile = num_steps_4 - 1; + mbarrier_arrive(page_empty_addr + (page_k_stage) * 8); + page_k_stage += 1; + if (page_k_stage == 6) { page_k_stage = 0; page_k_phase ^= 1; } + page_k_stage += 1; + if (page_k_stage == 6) { page_k_stage = 0; page_k_phase ^= 1; } + if ((last_tile & 1) == 0) { + mbarrier_wait(page_full_addr + (page_v_stage) * 8, page_v_phase); + } + mbarrier_arrive(page_empty_addr + (page_v_stage) * 8); + page_v_stage += 1; + if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } + page_v_stage += 1; + if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } + } + mbarrier_wait(work_full_addr + (work_stage_6) * 8, _phase_work_full_6); + uint32_t _clc_valid_1 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_1) + : "r"(work_response_addr + work_stage_6 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_3 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_3) + : "r"(work_response_addr + work_stage_6 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_4 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_4) + : "r"(work_response_addr + work_stage_6 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_5 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_5) + : "r"(work_response_addr + work_stage_6 * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + work_stage_6 * 8), "r"(0) : "memory"); + work_stage_6 += 1; + if (work_stage_6 == 2) { work_stage_6 = 0; _phase_work_full_6 ^= 1; } + if (_clc_valid_1 == 0) { + break; + } + work_x_5 = _clc_ctaid_3; + work_y_5 = _clc_ctaid_4; + work_z_5 = _clc_ctaid_5; + } + } + } + + // Cleanup +} + +} // extern "C" diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_d9754db4fe394a927e67_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_d9754db4fe394a927e67_binding.cu new file mode 100644 index 00000000000..171a6c68149 --- /dev/null +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_d9754db4fe394a927e67_binding.cu @@ -0,0 +1,385 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +// Generated by the Cake source exporter — do not edit. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_d9754db4fe394a927e67'. +#include +#include +#include +#include + +#include "tvm_ffi_utils.h" + +#include +#include +#include +#include +#include +#include +#include + +extern "C" __global__ void kernel_cake_dsv4_d9754db4fe394a927e67(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks); + + +namespace cake_host_shim_01a51b3a9acbf4d0 { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t CakeDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +// 4D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + int64_t d3 = t.size(t.ndim() - 3); + TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(d3)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 64u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_swa_kv' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_swa_kv' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_kv'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_O, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckContiguous(arg_tmap_q, "tmap_q"); + CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); + CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 4, 16, 1); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_swa_indices, "swa_indices"); + CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); + CheckContiguous(arg_swa_indices, "swa_indices"); + CheckCudaTensor(arg_compressed_indices, "compressed_indices"); + CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); + CheckContiguous(arg_compressed_indices, "compressed_indices"); + CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); + CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); + CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); + CheckContiguous(arg_bmm1_scale, "bmm1_scale"); + CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); + CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); + CheckContiguous(arg_bmm2_scale, "bmm2_scale"); + TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) + << "scalar 'num_heads' value " << arg_num_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) + << "scalar 'swa_index_stride' value " << arg_swa_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) + << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) + << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) + << "scalar 'sparse_topk' value " << arg_sparse_topk + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) + << "scalar 'has_sinks' value " << arg_has_sinks + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); + CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); + CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); + CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + + cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); + __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); + int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); + int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); + int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); + float* p_sinks = static_cast(arg_sinks.data_ptr()); + float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); + float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); + int32_t v_num_heads = (int32_t)arg_num_heads; + int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; + int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; + int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; + int32_t v_sparse_topk = (int32_t)arg_sparse_topk; + int32_t v_has_sinks = (int32_t)arg_has_sinks; + void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_O, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_has_sinks}; + + static std::mutex smem_mu; + static auto* smem_status_by_device = + new std::unordered_map(); + cudaError_t smem_status = cudaSuccess; + { + std::lock_guard lock(smem_mu); + auto it = smem_status_by_device->find(dev.device_id); + if (it == smem_status_by_device->end()) { + int current_device = -1; + smem_status = cudaGetDevice(¤t_device); + if (smem_status == cudaSuccess && current_device != dev.device_id) { + smem_status = cudaErrorInvalidDevice; + } + if (smem_status == cudaSuccess) { + smem_status = cudaFuncSetAttribute( + reinterpret_cast(kernel_cake_dsv4_d9754db4fe394a927e67), + cudaFuncAttributeMaxDynamicSharedMemorySize, + 198528); + } + smem_status_by_device->emplace(dev.device_id, smem_status); + } else { + smem_status = it->second; + } + } + TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) + << "cudaFuncSetAttribute for kernel_cake_dsv4_d9754db4fe394a927e67 failed: " + << cudaGetErrorString(smem_status); + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(512u, 1u, 1u); + + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes = 198528u; + config.stream = stream; + cudaError_t launch_status = cudaLaunchKernelExC( + &config, reinterpret_cast(kernel_cake_dsv4_d9754db4fe394a927e67), kargs); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernelExC for kernel_cake_dsv4_d9754db4fe394a927e67 failed: " + << cudaGetErrorString(launch_status); +} + +} // namespace cake_host_shim_01a51b3a9acbf4d0 + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_01a51b3a9acbf4d0::Run); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_d9754db4fe394a927e67_kernel.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_d9754db4fe394a927e67_kernel.cu new file mode 100644 index 00000000000..5de2a0df80a --- /dev/null +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_d9754db4fe394a927e67_kernel.cu @@ -0,0 +1,1734 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +#define CAKE_INF CUDART_INF_F +#define TMEM_NCOLS 384 +#define TMEM_TMEM_OFFSET 0 +#define NUM_KV_PIPE_STAGES 4 +#define SMEM_SMEM_Q_OFF 1024 +#define SMEM_SMEM_Q_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_STRIDE 8192 +#define SMEM_SMEM_KV_OFF 66560 +#define SMEM_SMEM_KV_STAGE_BYTES 32768 +#define SMEM_SMEM_KV_STRIDE 32768 +#define SMEM_SMEM_V_OFF 66560 +#define SMEM_SMEM_V_STAGE_BYTES 32768 +#define SMEM_SMEM_V_STRIDE 32768 +#define SMEM_SMEM_SUM_OFF 197632 +#define SMEM_SMEM_SUM_STAGE_BYTES 256 +#define SMEM_SMEM_SUM_STRIDE 256 +#define SMEM_SMEM_INDEX_FLAGS_OFF 198400 +#define SMEM_SMEM_INDEX_FLAGS_STAGE_BYTES 16 +#define SMEM_SMEM_INDEX_FLAGS_STRIDE 16 +#define SMEM_SMEM_INDICES_OFF 197888 +#define SMEM_SMEM_INDICES_STAGE_BYTES 512 +#define SMEM_SMEM_INDICES_STRIDE 512 +#define SMEM_TOTAL 198528 +#define THREADS 512 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count) : "memory"); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_2d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3}], [%4];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr) { + // Canonical .shared::cta form for non-multicast gather4, matching + // trtllm-gen / cuda_ptx and the PTX ISA qualifier order + // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, + // .shared::cluster is reserved for the multicast variant (ctaMask). + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(512, 1) void +kernel_cake_dsv4_d9754db4fe394a927e67(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int mbar_base = smem; + #define q_full_addr (mbar_base + 0) + #define kv_full_addr (mbar_base + 8) + #define kv_empty_addr (mbar_base + 40) + #define s_full_addr (mbar_base + 72) + #define p_full_addr (mbar_base + 80) + #define sum_ready_addr (mbar_base + 88) + #define o_done_addr (mbar_base + 96) + #define tmem_dealloc_addr (mbar_base + 104) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + + // Kernel setup ops + __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_q_addr = smem + 1024; + __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_kv_addr = smem + 66560; + __nv_bfloat16* smem_v = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_v_addr = smem + 66560; + float* smem_sum = reinterpret_cast(smem_raw + 197632); + const int smem_sum_addr = smem + 197632; + int* smem_index_flags = reinterpret_cast(smem_raw + 198400); + const int smem_index_flags_addr = smem + 198400; + int* smem_indices = reinterpret_cast(smem_raw + 197888); + const int smem_indices_addr = smem + 197888; + + // Mbarrier init (8 pipeline groups, 0 ordered-sequence groups, 14 barriers) + // Mbarriers at smem_raw[0..112) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // q_full: 1 barriers, init_count=1 + mbarrier_init(smem + 0, 1); + // --- pipeline 'kv_pipe' --- + // kv_full: 4 barriers, init_count=4 + mbarrier_init(smem + 8, 4); + mbarrier_init(smem + 16, 4); + mbarrier_init(smem + 24, 4); + mbarrier_init(smem + 32, 4); + // kv_empty: 4 barriers, init_count=1 + mbarrier_init(smem + 40, 1); + mbarrier_init(smem + 48, 1); + mbarrier_init(smem + 56, 1); + mbarrier_init(smem + 64, 1); + // s_full: 1 barriers, init_count=1 + mbarrier_init(smem + 72, 1); + // p_full: 1 barriers, init_count=128 + mbarrier_init(smem + 80, 128); + // sum_ready: 1 barriers, init_count=128 + mbarrier_init(smem + 88, 128); + // o_done: 1 barriers, init_count=1 + mbarrier_init(smem + 96, 1); + // tmem_dealloc: 1 barriers, init_count=256 + mbarrier_init(smem + 104, 256); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncwarp(); + + // TMEM alloc (512 columns, 384 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 112); + if (warp == 0) { + int _tmem_hold = smem + 112; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_tmem = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 64;"); + } + + // ---- Role: softmax ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); + { // softmax_main + float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; + int query_idx = blockIdx.x >> 2; + int _max_0 = ((sparse_topk_lens[query_idx] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx] + sparse_topk_lens_offset) : (0)); + int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); + int active_topk = _min_0; + const int warp_in_compute = warp; + const int tmem_row_origin = warp_in_compute * 32; + const int logical_row_origin = warp_in_compute * 16; + const int my_row = logical_row_origin + lane % 16; + const int col_half = lane / 16; + unsigned int _phase_s_full_0 = 0; + mbarrier_wait(s_full_addr, _phase_s_full_0); + _phase_s_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int score_addr = taddr + (unsigned int)(tmem_row_origin << 16); + float _tmem_load_0[64]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_0[0])), "=r"(*reinterpret_cast(&_tmem_load_0[1])), "=r"(*reinterpret_cast(&_tmem_load_0[2])), "=r"(*reinterpret_cast(&_tmem_load_0[3])), "=r"(*reinterpret_cast(&_tmem_load_0[4])), "=r"(*reinterpret_cast(&_tmem_load_0[5])), "=r"(*reinterpret_cast(&_tmem_load_0[6])), "=r"(*reinterpret_cast(&_tmem_load_0[7])), "=r"(*reinterpret_cast(&_tmem_load_0[8])), "=r"(*reinterpret_cast(&_tmem_load_0[9])), "=r"(*reinterpret_cast(&_tmem_load_0[10])), "=r"(*reinterpret_cast(&_tmem_load_0[11])), "=r"(*reinterpret_cast(&_tmem_load_0[12])), "=r"(*reinterpret_cast(&_tmem_load_0[13])), "=r"(*reinterpret_cast(&_tmem_load_0[14])), "=r"(*reinterpret_cast(&_tmem_load_0[15])), "=r"(*reinterpret_cast(&_tmem_load_0[16])), "=r"(*reinterpret_cast(&_tmem_load_0[17])), "=r"(*reinterpret_cast(&_tmem_load_0[18])), "=r"(*reinterpret_cast(&_tmem_load_0[19])), "=r"(*reinterpret_cast(&_tmem_load_0[20])), "=r"(*reinterpret_cast(&_tmem_load_0[21])), "=r"(*reinterpret_cast(&_tmem_load_0[22])), "=r"(*reinterpret_cast(&_tmem_load_0[23])), "=r"(*reinterpret_cast(&_tmem_load_0[24])), "=r"(*reinterpret_cast(&_tmem_load_0[25])), "=r"(*reinterpret_cast(&_tmem_load_0[26])), "=r"(*reinterpret_cast(&_tmem_load_0[27])), "=r"(*reinterpret_cast(&_tmem_load_0[28])), "=r"(*reinterpret_cast(&_tmem_load_0[29])), "=r"(*reinterpret_cast(&_tmem_load_0[30])), "=r"(*reinterpret_cast(&_tmem_load_0[31])) + : "r"(score_addr)); + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_0[32])), "=r"(*reinterpret_cast(&_tmem_load_0[33])), "=r"(*reinterpret_cast(&_tmem_load_0[34])), "=r"(*reinterpret_cast(&_tmem_load_0[35])), "=r"(*reinterpret_cast(&_tmem_load_0[36])), "=r"(*reinterpret_cast(&_tmem_load_0[37])), "=r"(*reinterpret_cast(&_tmem_load_0[38])), "=r"(*reinterpret_cast(&_tmem_load_0[39])), "=r"(*reinterpret_cast(&_tmem_load_0[40])), "=r"(*reinterpret_cast(&_tmem_load_0[41])), "=r"(*reinterpret_cast(&_tmem_load_0[42])), "=r"(*reinterpret_cast(&_tmem_load_0[43])), "=r"(*reinterpret_cast(&_tmem_load_0[44])), "=r"(*reinterpret_cast(&_tmem_load_0[45])), "=r"(*reinterpret_cast(&_tmem_load_0[46])), "=r"(*reinterpret_cast(&_tmem_load_0[47])), "=r"(*reinterpret_cast(&_tmem_load_0[48])), "=r"(*reinterpret_cast(&_tmem_load_0[49])), "=r"(*reinterpret_cast(&_tmem_load_0[50])), "=r"(*reinterpret_cast(&_tmem_load_0[51])), "=r"(*reinterpret_cast(&_tmem_load_0[52])), "=r"(*reinterpret_cast(&_tmem_load_0[53])), "=r"(*reinterpret_cast(&_tmem_load_0[54])), "=r"(*reinterpret_cast(&_tmem_load_0[55])), "=r"(*reinterpret_cast(&_tmem_load_0[56])), "=r"(*reinterpret_cast(&_tmem_load_0[57])), "=r"(*reinterpret_cast(&_tmem_load_0[58])), "=r"(*reinterpret_cast(&_tmem_load_0[59])), "=r"(*reinterpret_cast(&_tmem_load_0[60])), "=r"(*reinterpret_cast(&_tmem_load_0[61])), "=r"(*reinterpret_cast(&_tmem_load_0[62])), "=r"(*reinterpret_cast(&_tmem_load_0[63])) + : "r"(score_addr + 32)); + int valid_cols = active_topk - col_half * 64; + int causal_cols = smem_index_flags[col_half * 2 + 1]; + valid_cols = ((valid_cols < causal_cols) ? valid_cols : causal_cols); + uint32_t _slice_lo_mask_0; + { + int _lim_0 = valid_cols; + if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_0 & (1u << 0))) _tmem_load_0[0] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 1))) _tmem_load_0[1] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 2))) _tmem_load_0[2] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 3))) _tmem_load_0[3] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 4))) _tmem_load_0[4] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 5))) _tmem_load_0[5] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 6))) _tmem_load_0[6] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 7))) _tmem_load_0[7] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 8))) _tmem_load_0[8] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 9))) _tmem_load_0[9] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 10))) _tmem_load_0[10] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 11))) _tmem_load_0[11] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 12))) _tmem_load_0[12] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 13))) _tmem_load_0[13] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 14))) _tmem_load_0[14] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 15))) _tmem_load_0[15] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 16))) _tmem_load_0[16] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 17))) _tmem_load_0[17] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 18))) _tmem_load_0[18] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 19))) _tmem_load_0[19] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 20))) _tmem_load_0[20] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 21))) _tmem_load_0[21] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 22))) _tmem_load_0[22] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 23))) _tmem_load_0[23] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 24))) _tmem_load_0[24] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 25))) _tmem_load_0[25] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 26))) _tmem_load_0[26] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 27))) _tmem_load_0[27] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 28))) _tmem_load_0[28] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 29))) _tmem_load_0[29] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 30))) _tmem_load_0[30] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 31))) _tmem_load_0[31] = -CAKE_INF; + uint32_t _slice_lo_mask_1; + { + int _lim_1 = valid_cols - 32; + if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) _tmem_load_0[32] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) _tmem_load_0[33] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) _tmem_load_0[34] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) _tmem_load_0[35] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) _tmem_load_0[36] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) _tmem_load_0[37] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) _tmem_load_0[38] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) _tmem_load_0[39] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) _tmem_load_0[40] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) _tmem_load_0[41] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) _tmem_load_0[42] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) _tmem_load_0[43] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) _tmem_load_0[44] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) _tmem_load_0[45] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) _tmem_load_0[46] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) _tmem_load_0[47] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) _tmem_load_0[48] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) _tmem_load_0[49] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) _tmem_load_0[50] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) _tmem_load_0[51] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) _tmem_load_0[52] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) _tmem_load_0[53] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) _tmem_load_0[54] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) _tmem_load_0[55] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) _tmem_load_0[56] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) _tmem_load_0[57] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) _tmem_load_0[58] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) _tmem_load_0[59] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) _tmem_load_0[60] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) _tmem_load_0[61] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) _tmem_load_0[62] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) _tmem_load_0[63] = -CAKE_INF; + float2 _reg_reduce_max2_2 = {-CAKE_INF, -CAKE_INF}; + row_max_x32_accum(&_tmem_load_0[0], _reg_reduce_max2_2); + row_max_x32_accum(&_tmem_load_0[32], _reg_reduce_max2_2); + float _tmem_load_0_max = row_max_reduce(_reg_reduce_max2_2); + float row_max = _tmem_load_0_max; + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, row_max, 16); + float _max_1 = max_noftz(row_max, _shfl_xor_0); + row_max = _max_1; + if (has_sinks != 0 && my_row < num_heads) { + float sink_unscaled = sinks[my_row] * 1.4426950408889634f / softmax_scale_log2; + float _max_2 = max_noftz(row_max, sink_unscaled); + row_max = _max_2; + } + bool empty_row = row_max == -CAKE_INF; + float safe_max = ((empty_row) ? 0.0f : row_max); + float max_scaled = safe_max * softmax_scale_log2; + float2 _f2_0 = make_float2(softmax_scale_log2, softmax_scale_log2); + float2 _f2_1 = make_float2(-max_scaled, -max_scaled); + float2 _f2_2 = make_float2(0.0f, 0.0f); + float2 sum0 = _f2_2; + float2 _f2_3 = make_float2(0.0f, 0.0f); + float2 sum1 = _f2_3; + float2 _f2_4 = make_float2(0.0f, 0.0f); + float2 sum2 = _f2_4; + float2 _f2_5 = make_float2(0.0f, 0.0f); + float2 sum3 = _f2_5; + #pragma unroll + for (int i = 0; i < 64; i += 8) { + float2 _f2_6 = make_float2(_tmem_load_0[i], _tmem_load_0[i + 1]); + float2 _f2_7 = make_float2(_tmem_load_0[i + 2], _tmem_load_0[i + 3]); + float2 _f2_8 = make_float2(_tmem_load_0[i + 4], _tmem_load_0[i + 5]); + float2 _f2_9 = make_float2(_tmem_load_0[i + 6], _tmem_load_0[i + 7]); + float2 affine0 = fma_f32x2_rn_ftz(_f2_6, _f2_0, _f2_1); + float2 affine1 = fma_f32x2_rn_ftz(_f2_7, _f2_0, _f2_1); + float2 affine2 = fma_f32x2_rn_ftz(_f2_8, _f2_0, _f2_1); + float2 affine3 = fma_f32x2_rn_ftz(_f2_9, _f2_0, _f2_1); + float _exp2_0 = approx_exp2(affine0.x); + float exp0 = _exp2_0; + float _exp2_1 = approx_exp2(affine0.y); + float exp1 = _exp2_1; + float _exp2_2 = approx_exp2(affine1.x); + float exp2 = _exp2_2; + float _exp2_3 = approx_exp2(affine1.y); + float exp3 = _exp2_3; + float _exp2_4 = approx_exp2(affine2.x); + float exp4 = _exp2_4; + float _exp2_5 = approx_exp2(affine2.y); + float exp5 = _exp2_5; + float _exp2_6 = approx_exp2(affine3.x); + float exp6 = _exp2_6; + float _exp2_7 = approx_exp2(affine3.y); + float exp7 = _exp2_7; + _tmem_load_0[i] = exp0; + _tmem_load_0[i + 1] = exp1; + _tmem_load_0[i + 2] = exp2; + _tmem_load_0[i + 3] = exp3; + _tmem_load_0[i + 4] = exp4; + _tmem_load_0[i + 5] = exp5; + _tmem_load_0[i + 6] = exp6; + _tmem_load_0[i + 7] = exp7; + float2 _f2_10 = make_float2(exp0, exp1); + sum0 = add_f32x2(sum0, _f2_10); + float2 _f2_11 = make_float2(exp2, exp3); + sum1 = add_f32x2(sum1, _f2_11); + float2 _f2_12 = make_float2(exp4, exp5); + sum2 = add_f32x2(sum2, _f2_12); + float2 _f2_13 = make_float2(exp6, exp7); + sum3 = add_f32x2(sum3, _f2_13); + } + float2 sum01 = add_f32x2(sum0, sum1); + float2 sum23 = add_f32x2(sum2, sum3); + float2 row_sum_pair = add_f32x2(sum01, sum23); + float row_sum = row_sum_pair.x + row_sum_pair.y; + float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, row_sum, 16); + row_sum = row_sum + _shfl_xor_1; + if (has_sinks != 0 && my_row < num_heads) { + float _exp2_8 = approx_exp2(sinks[my_row] * 1.4426950408889634f - max_scaled); + row_sum = row_sum + _exp2_8; + } + unsigned int packed_p[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_0[_lp*2 + 0], _tmem_load_0[_lp*2+1 + 0])); + packed_p[_lp] = *(uint32_t*)&_bf2; + } + int p_addr = taddr + 128 + (unsigned int)(tmem_row_origin << 16); + asm volatile( + "tcgen05.st.sync.aligned.16x32bx2.x16.b32" + " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(p_addr), "r"(*reinterpret_cast(&packed_p[0])), "r"(*reinterpret_cast(&packed_p[1])), "r"(*reinterpret_cast(&packed_p[2])), "r"(*reinterpret_cast(&packed_p[3])), "r"(*reinterpret_cast(&packed_p[4])), "r"(*reinterpret_cast(&packed_p[5])), "r"(*reinterpret_cast(&packed_p[6])), "r"(*reinterpret_cast(&packed_p[7])), "r"(*reinterpret_cast(&packed_p[8])), "r"(*reinterpret_cast(&packed_p[9])), "r"(*reinterpret_cast(&packed_p[10])), "r"(*reinterpret_cast(&packed_p[11])), "r"(*reinterpret_cast(&packed_p[12])), "r"(*reinterpret_cast(&packed_p[13])), "r"(*reinterpret_cast(&packed_p[14])), "r"(*reinterpret_cast(&packed_p[15]))); + asm volatile( + "tcgen05.st.sync.aligned.16x32bx2.x16.b32" + " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(p_addr + 16), "r"(*reinterpret_cast(&(packed_p + 16)[0])), "r"(*reinterpret_cast(&(packed_p + 16)[1])), "r"(*reinterpret_cast(&(packed_p + 16)[2])), "r"(*reinterpret_cast(&(packed_p + 16)[3])), "r"(*reinterpret_cast(&(packed_p + 16)[4])), "r"(*reinterpret_cast(&(packed_p + 16)[5])), "r"(*reinterpret_cast(&(packed_p + 16)[6])), "r"(*reinterpret_cast(&(packed_p + 16)[7])), "r"(*reinterpret_cast(&(packed_p + 16)[8])), "r"(*reinterpret_cast(&(packed_p + 16)[9])), "r"(*reinterpret_cast(&(packed_p + 16)[10])), "r"(*reinterpret_cast(&(packed_p + 16)[11])), "r"(*reinterpret_cast(&(packed_p + 16)[12])), "r"(*reinterpret_cast(&(packed_p + 16)[13])), "r"(*reinterpret_cast(&(packed_p + 16)[14])), "r"(*reinterpret_cast(&(packed_p + 16)[15]))); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(p_full_addr); + if (col_half == 0) { + float _max_3 = max_noftz(row_sum, 1.1754943508222875e-38f); + float published_sum = _max_3; + smem_sum[my_row] = published_sum; + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(sum_ready_addr); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_dealloc_addr); + } + } + // ---- Role: epilogue ---- + if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 128;"); + { // epilogue_main + float output_scale = bmm2_scale[0]; + int work_idx = blockIdx.x; + int query_idx_1 = work_idx >> 2; + int v_chunk = work_idx & 3; + const int warp_in_role = warp - 4; + const int tmem_row_origin_1 = warp_in_role * 32; + const int logical_row_origin_1 = warp_in_role * 16; + const int my_row_1 = logical_row_origin_1 + lane % 16; + const int col_half_1 = lane / 16; + const int row_addr = tmem_row_origin_1 << 16; + unsigned int _phase_o_done_0 = 0; + mbarrier_wait(o_done_addr, _phase_o_done_0); + _phase_o_done_0 ^= 1; + unsigned int _phase_sum_ready_0 = 0; + mbarrier_wait(sum_ready_addr, _phase_sum_ready_0); + _phase_sum_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + float row_sum_1 = smem_sum[my_row_1]; + float _rcp_0 = approx_rcp(row_sum_1); + float inv_sum = _rcp_0; + int output_base = (query_idx_1 * num_heads + my_row_1) * 512 + v_chunk * 128; + float _tmem_load_1[64]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_1[0])), "=r"(*reinterpret_cast(&_tmem_load_1[1])), "=r"(*reinterpret_cast(&_tmem_load_1[2])), "=r"(*reinterpret_cast(&_tmem_load_1[3])), "=r"(*reinterpret_cast(&_tmem_load_1[4])), "=r"(*reinterpret_cast(&_tmem_load_1[5])), "=r"(*reinterpret_cast(&_tmem_load_1[6])), "=r"(*reinterpret_cast(&_tmem_load_1[7])), "=r"(*reinterpret_cast(&_tmem_load_1[8])), "=r"(*reinterpret_cast(&_tmem_load_1[9])), "=r"(*reinterpret_cast(&_tmem_load_1[10])), "=r"(*reinterpret_cast(&_tmem_load_1[11])), "=r"(*reinterpret_cast(&_tmem_load_1[12])), "=r"(*reinterpret_cast(&_tmem_load_1[13])), "=r"(*reinterpret_cast(&_tmem_load_1[14])), "=r"(*reinterpret_cast(&_tmem_load_1[15])), "=r"(*reinterpret_cast(&_tmem_load_1[16])), "=r"(*reinterpret_cast(&_tmem_load_1[17])), "=r"(*reinterpret_cast(&_tmem_load_1[18])), "=r"(*reinterpret_cast(&_tmem_load_1[19])), "=r"(*reinterpret_cast(&_tmem_load_1[20])), "=r"(*reinterpret_cast(&_tmem_load_1[21])), "=r"(*reinterpret_cast(&_tmem_load_1[22])), "=r"(*reinterpret_cast(&_tmem_load_1[23])), "=r"(*reinterpret_cast(&_tmem_load_1[24])), "=r"(*reinterpret_cast(&_tmem_load_1[25])), "=r"(*reinterpret_cast(&_tmem_load_1[26])), "=r"(*reinterpret_cast(&_tmem_load_1[27])), "=r"(*reinterpret_cast(&_tmem_load_1[28])), "=r"(*reinterpret_cast(&_tmem_load_1[29])), "=r"(*reinterpret_cast(&_tmem_load_1[30])), "=r"(*reinterpret_cast(&_tmem_load_1[31])) + : "r"(taddr + 256 + (unsigned int)row_addr)); + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_1[32])), "=r"(*reinterpret_cast(&_tmem_load_1[33])), "=r"(*reinterpret_cast(&_tmem_load_1[34])), "=r"(*reinterpret_cast(&_tmem_load_1[35])), "=r"(*reinterpret_cast(&_tmem_load_1[36])), "=r"(*reinterpret_cast(&_tmem_load_1[37])), "=r"(*reinterpret_cast(&_tmem_load_1[38])), "=r"(*reinterpret_cast(&_tmem_load_1[39])), "=r"(*reinterpret_cast(&_tmem_load_1[40])), "=r"(*reinterpret_cast(&_tmem_load_1[41])), "=r"(*reinterpret_cast(&_tmem_load_1[42])), "=r"(*reinterpret_cast(&_tmem_load_1[43])), "=r"(*reinterpret_cast(&_tmem_load_1[44])), "=r"(*reinterpret_cast(&_tmem_load_1[45])), "=r"(*reinterpret_cast(&_tmem_load_1[46])), "=r"(*reinterpret_cast(&_tmem_load_1[47])), "=r"(*reinterpret_cast(&_tmem_load_1[48])), "=r"(*reinterpret_cast(&_tmem_load_1[49])), "=r"(*reinterpret_cast(&_tmem_load_1[50])), "=r"(*reinterpret_cast(&_tmem_load_1[51])), "=r"(*reinterpret_cast(&_tmem_load_1[52])), "=r"(*reinterpret_cast(&_tmem_load_1[53])), "=r"(*reinterpret_cast(&_tmem_load_1[54])), "=r"(*reinterpret_cast(&_tmem_load_1[55])), "=r"(*reinterpret_cast(&_tmem_load_1[56])), "=r"(*reinterpret_cast(&_tmem_load_1[57])), "=r"(*reinterpret_cast(&_tmem_load_1[58])), "=r"(*reinterpret_cast(&_tmem_load_1[59])), "=r"(*reinterpret_cast(&_tmem_load_1[60])), "=r"(*reinterpret_cast(&_tmem_load_1[61])), "=r"(*reinterpret_cast(&_tmem_load_1[62])), "=r"(*reinterpret_cast(&_tmem_load_1[63])) + : "r"(taddr + 256 + (unsigned int)row_addr + 32)); + float2 _f2_14 = make_float2(inv_sum * output_scale, inv_sum * output_scale); + #pragma unroll + for (int i_1 = 0; i_1 < 64; i_1 += 2) { + float2 _f2_15 = make_float2(_tmem_load_1[i_1], _tmem_load_1[i_1 + 1]); + float2 _mul_f32x2_0; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_0) : "l"(*(const unsigned long long*)&_f2_15), "l"(*(const unsigned long long*)&_f2_14)); + _tmem_load_1[i_1] = _mul_f32x2_0.x; + _tmem_load_1[i_1 + 1] = _mul_f32x2_0.y; + } + unsigned int packed_o[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_1[_lp*2 + 0], _tmem_load_1[_lp*2+1 + 0])); + packed_o[_lp] = *(uint32_t*)&_bf2; + } + bool stage_hi_1 = (lane & 1) != 0; + bool stage_hi_2 = (lane & 2) != 0; + bool stage_hi_4 = (lane & 4) != 0; + unsigned int _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[0] : packed_o[4]), 1); + packed_o[0] = ((stage_hi_1) ? _shfl_xor_2 : packed_o[0]); + packed_o[4] = ((stage_hi_1) ? packed_o[4] : _shfl_xor_2); + unsigned int _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[1] : packed_o[5]), 1); + packed_o[1] = ((stage_hi_1) ? _shfl_xor_3 : packed_o[1]); + packed_o[5] = ((stage_hi_1) ? packed_o[5] : _shfl_xor_3); + unsigned int _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[2] : packed_o[6]), 1); + packed_o[2] = ((stage_hi_1) ? _shfl_xor_4 : packed_o[2]); + packed_o[6] = ((stage_hi_1) ? packed_o[6] : _shfl_xor_4); + unsigned int _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[3] : packed_o[7]), 1); + packed_o[3] = ((stage_hi_1) ? _shfl_xor_5 : packed_o[3]); + packed_o[7] = ((stage_hi_1) ? packed_o[7] : _shfl_xor_5); + unsigned int _shfl_xor_6 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[8] : packed_o[12]), 1); + packed_o[8] = ((stage_hi_1) ? _shfl_xor_6 : packed_o[8]); + packed_o[12] = ((stage_hi_1) ? packed_o[12] : _shfl_xor_6); + unsigned int _shfl_xor_7 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[9] : packed_o[13]), 1); + packed_o[9] = ((stage_hi_1) ? _shfl_xor_7 : packed_o[9]); + packed_o[13] = ((stage_hi_1) ? packed_o[13] : _shfl_xor_7); + unsigned int _shfl_xor_8 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[10] : packed_o[14]), 1); + packed_o[10] = ((stage_hi_1) ? _shfl_xor_8 : packed_o[10]); + packed_o[14] = ((stage_hi_1) ? packed_o[14] : _shfl_xor_8); + unsigned int _shfl_xor_9 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[11] : packed_o[15]), 1); + packed_o[11] = ((stage_hi_1) ? _shfl_xor_9 : packed_o[11]); + packed_o[15] = ((stage_hi_1) ? packed_o[15] : _shfl_xor_9); + unsigned int _shfl_xor_10 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[16] : packed_o[20]), 1); + packed_o[16] = ((stage_hi_1) ? _shfl_xor_10 : packed_o[16]); + packed_o[20] = ((stage_hi_1) ? packed_o[20] : _shfl_xor_10); + unsigned int _shfl_xor_11 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[17] : packed_o[21]), 1); + packed_o[17] = ((stage_hi_1) ? _shfl_xor_11 : packed_o[17]); + packed_o[21] = ((stage_hi_1) ? packed_o[21] : _shfl_xor_11); + unsigned int _shfl_xor_12 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[18] : packed_o[22]), 1); + packed_o[18] = ((stage_hi_1) ? _shfl_xor_12 : packed_o[18]); + packed_o[22] = ((stage_hi_1) ? packed_o[22] : _shfl_xor_12); + unsigned int _shfl_xor_13 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[19] : packed_o[23]), 1); + packed_o[19] = ((stage_hi_1) ? _shfl_xor_13 : packed_o[19]); + packed_o[23] = ((stage_hi_1) ? packed_o[23] : _shfl_xor_13); + unsigned int _shfl_xor_14 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[24] : packed_o[28]), 1); + packed_o[24] = ((stage_hi_1) ? _shfl_xor_14 : packed_o[24]); + packed_o[28] = ((stage_hi_1) ? packed_o[28] : _shfl_xor_14); + unsigned int _shfl_xor_15 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[25] : packed_o[29]), 1); + packed_o[25] = ((stage_hi_1) ? _shfl_xor_15 : packed_o[25]); + packed_o[29] = ((stage_hi_1) ? packed_o[29] : _shfl_xor_15); + unsigned int _shfl_xor_16 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[26] : packed_o[30]), 1); + packed_o[26] = ((stage_hi_1) ? _shfl_xor_16 : packed_o[26]); + packed_o[30] = ((stage_hi_1) ? packed_o[30] : _shfl_xor_16); + unsigned int _shfl_xor_17 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[27] : packed_o[31]), 1); + packed_o[27] = ((stage_hi_1) ? _shfl_xor_17 : packed_o[27]); + packed_o[31] = ((stage_hi_1) ? packed_o[31] : _shfl_xor_17); + int lane_t = lane & 1; + unsigned int chunk_words[4]; + chunk_words[0] = packed_o[0]; + chunk_words[1] = packed_o[1]; + chunk_words[2] = packed_o[2]; + chunk_words[3] = packed_o[3]; + if ((my_row_1 & -2 | 0) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + (my_row_1 & -2 | 0)) * 512 + v_chunk * 128 + col_half_1 * 64 + (0 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[4]; + chunk_words[1] = packed_o[5]; + chunk_words[2] = packed_o[6]; + chunk_words[3] = packed_o[7]; + if ((my_row_1 & -2 | 1) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + (my_row_1 & -2 | 1)) * 512 + v_chunk * 128 + col_half_1 * 64 + (0 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[8]; + chunk_words[1] = packed_o[9]; + chunk_words[2] = packed_o[10]; + chunk_words[3] = packed_o[11]; + if ((my_row_1 & -2 | 0) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + (my_row_1 & -2 | 0)) * 512 + v_chunk * 128 + col_half_1 * 64 + (2 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[12]; + chunk_words[1] = packed_o[13]; + chunk_words[2] = packed_o[14]; + chunk_words[3] = packed_o[15]; + if ((my_row_1 & -2 | 1) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + (my_row_1 & -2 | 1)) * 512 + v_chunk * 128 + col_half_1 * 64 + (2 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[16]; + chunk_words[1] = packed_o[17]; + chunk_words[2] = packed_o[18]; + chunk_words[3] = packed_o[19]; + if ((my_row_1 & -2 | 0) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + (my_row_1 & -2 | 0)) * 512 + v_chunk * 128 + col_half_1 * 64 + (4 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[20]; + chunk_words[1] = packed_o[21]; + chunk_words[2] = packed_o[22]; + chunk_words[3] = packed_o[23]; + if ((my_row_1 & -2 | 1) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + (my_row_1 & -2 | 1)) * 512 + v_chunk * 128 + col_half_1 * 64 + (4 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[24]; + chunk_words[1] = packed_o[25]; + chunk_words[2] = packed_o[26]; + chunk_words[3] = packed_o[27]; + if ((my_row_1 & -2 | 0) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + (my_row_1 & -2 | 0)) * 512 + v_chunk * 128 + col_half_1 * 64 + (6 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[28]; + chunk_words[1] = packed_o[29]; + chunk_words[2] = packed_o[30]; + chunk_words[3] = packed_o[31]; + if ((my_row_1 & -2 | 1) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + (my_row_1 & -2 | 1)) * 512 + v_chunk * 128 + col_half_1 * 64 + (6 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_dealloc_addr); + } + } + // ---- Role: mma_warp ---- + if (warp == 8) { + { // mma_warp_main + unsigned int _phase_q_full_0 = 0; + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + unsigned int _phase_kv_full = 0; + #pragma unroll + for (int k_stage = 0; k_stage < 4; k_stage++) { + mbarrier_wait(kv_full_addr + (k_stage) * 8, _phase_kv_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_0 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (k_stage * 2) * 512); + int _mma_b_lo_0 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (k_stage) * 2048); + { + uint64_t _mma_ss_a_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_0); + uint64_t _mma_ss_b_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_0); + if (elect_sync()) { + tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, ((k_stage == 0) ? 0 : 1)); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); + } + } + int _mma_a_lo_1 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (k_stage * 2 + 1) * 512); + int _mma_b_lo_1 = make_warp_uniform((((smem_kv_addr + 16384) >> 4) & 0x3FFF) + (k_stage) * 2048); + { + uint64_t _mma_ss_a_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_1); + uint64_t _mma_ss_b_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_1); + if (elect_sync()) { + tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); + } + } + if (k_stage == 3) { + elect_commit(s_full_addr); + } + elect_commit(kv_empty_addr + (k_stage) * 8); + } + int _mma_prepared_b_lo_2 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (0) * 2048); + unsigned int _phase_p_full_0 = 0; + mbarrier_wait(p_full_addr, _phase_p_full_0); + _phase_p_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + mbarrier_wait(kv_full_addr, 1); + asm volatile("tcgen05.fence::after_thread_sync;"); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 dhi, blo, ta, id;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 69272720;\n\t" + "mov.b32 ta, %2;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p0;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "}\n" + :: "r"((tmem_tmem + (256))), "r"(_mma_prepared_b_lo_2), "r"(tmem_tmem + 128), "r"(0)); + elect_commit(o_done_addr); + elect_commit(kv_empty_addr); + unsigned int _phase_tmem_dealloc_0 = 0; + mbarrier_wait(tmem_dealloc_addr, _phase_tmem_dealloc_0); + _phase_tmem_dealloc_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + } + // ---- Role: empty ---- + if (warp >= 9 && warp <= 11) { + // idle — no tasks assigned + } + // ---- Role: load_warp ---- + if (warp >= 12 && warp <= 15) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 64;"); + { // load_warp_main + const int load_warp_rank = warp - 12; + int work_idx_1 = blockIdx.x; + int query_idx_2 = work_idx_1 >> 2; + int v_chunk_1 = work_idx_1 & 3; + int* row_ptr = swa_indices + (query_idx_2 * swa_index_stride); + int index_offset = load_warp_rank * 32 + lane; + int sparse_row = -1; + if (index_offset < sparse_topk) { + sparse_row = row_ptr[index_offset]; + } + if (load_warp_rank == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(q_full_addr, 65536); + #pragma unroll + for (int q_stage = 0; q_stage < 8; q_stage++) { + tma_4d_gmem2smem(smem_q_addr + (unsigned int)(q_stage * 8192), (&tmap_q), 0, 0, q_stage, query_idx_2, q_full_addr); + } + } + } + smem_indices[index_offset] = sparse_row; + unsigned int _vote_0 = __ballot_sync(0xFFFFFFFF, sparse_row < 0); + unsigned int invalid_indices = _vote_0; + int valid_prefix = 32; + if (invalid_indices != 0) { + int _ffs_0 = __ffs(invalid_indices); + valid_prefix = _ffs_0 - 1; + } + if (elect_sync()) { + smem_index_flags[load_warp_rank] = valid_prefix; + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("barrier.sync 8, 128;" ::: "memory"); + if ((load_warp_rank & 1) != 0) { + if (elect_sync()) { + int lower_prefix = smem_index_flags[load_warp_rank - 1]; + int upper_prefix = smem_index_flags[load_warp_rank]; + smem_index_flags[load_warp_rank] = ((lower_prefix < 32) ? lower_prefix : 32 + upper_prefix); + } + } + unsigned int _phase_kv_empty = 1; + #pragma unroll + for (int k_stage_1 = 0; k_stage_1 < 4; k_stage_1++) { + mbarrier_wait(kv_empty_addr + (k_stage_1) * 8, _phase_kv_empty); + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (k_stage_1) * 8, 8192); + } + #pragma unroll + for (int local_group = 0; local_group < 8; local_group++) { + int dst_k = smem_kv_addr + (unsigned int)(k_stage_1 * 32768); + int group = load_warp_rank * 8 + local_group; + int group_offset = group * 4; + int raw_rows[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows[0])), "=r"(*reinterpret_cast(&raw_rows[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows[(0) + 3])) + : "r"(smem_indices_addr + (unsigned int)(group_offset * 4))); + int raw0 = raw_rows[0]; + int raw1 = raw_rows[1]; + int raw2 = raw_rows[2]; + int raw3 = raw_rows[3]; + int row0 = ((raw0 >= 0) ? raw0 : 0); + int row1 = ((raw1 >= 0) ? raw1 : 0); + int row2 = ((raw2 >= 0) ? raw2 : 0); + int row3 = ((raw3 >= 0) ? raw3 : 0); + if (elect_sync()) { + tma_gather4_gmem2smem(dst_k + group * 512, (&tmap_swa_kv), k_stage_1 * 128, row0, row1, row2, row3, kv_full_addr + (k_stage_1) * 8); + tma_gather4_gmem2smem(dst_k + 16384 + group * 512, (&tmap_swa_kv), k_stage_1 * 128 + 64, row0, row1, row2, row3, kv_full_addr + (k_stage_1) * 8); + } + } + } + #pragma unroll + for (int v_stage = 0; v_stage < 1; v_stage++) { + mbarrier_wait(kv_empty_addr + (v_stage) * 8, 0); + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (v_stage) * 8, 8192); + } + #pragma unroll + for (int local_group_1 = 0; local_group_1 < 8; local_group_1++) { + int dst_v = smem_v_addr + (unsigned int)(v_stage * 32768); + int v_col = v_chunk_1 * 128 + v_stage * 64; + int group_1 = load_warp_rank * 8 + local_group_1; + int group_offset_1 = group_1 * 4; + int raw_rows_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_1[0])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 3])) + : "r"(smem_indices_addr + (unsigned int)(group_offset_1 * 4))); + int raw0_1 = raw_rows_1[0]; + int raw1_1 = raw_rows_1[1]; + int raw2_1 = raw_rows_1[2]; + int raw3_1 = raw_rows_1[3]; + int row0_1 = ((raw0_1 >= 0) ? raw0_1 : 0); + int row1_1 = ((raw1_1 >= 0) ? raw1_1 : 0); + int row2_1 = ((raw2_1 >= 0) ? raw2_1 : 0); + int row3_1 = ((raw3_1 >= 0) ? raw3_1 : 0); + if (elect_sync()) { + tma_gather4_gmem2smem(dst_v + group_1 * 512, (&tmap_swa_kv), v_col, row0_1, row1_1, row2_1, row3_1, kv_full_addr + (v_stage) * 8); + tma_gather4_gmem2smem(dst_v + 16384 + group_1 * 512, (&tmap_swa_kv), v_col + 64, row0_1, row1_1, row2_1, row3_1, kv_full_addr + (v_stage) * 8); + } + } + } + } + } + + // Cleanup +} + +} // extern "C" diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_e48c5e0a89a41eb1e57a_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_e48c5e0a89a41eb1e57a_binding.cu new file mode 100644 index 00000000000..0b839f4f0ba --- /dev/null +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_e48c5e0a89a41eb1e57a_binding.cu @@ -0,0 +1,526 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +// Generated by the Cake source exporter — do not edit. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_e48c5e0a89a41eb1e57a'. +#include +#include +#include +#include + +#include "tvm_ffi_utils.h" + +#include +#include +#include +#include +#include +#include +#include + +struct CakeTensorMap; +extern "C" __global__ void kernel_cake_dsv4_e48c5e0a89a41eb1e57a(CakeTensorMap const* tmap_q, CakeTensorMap const* tmap_swa_kv, CakeTensorMap const* tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int max_q_len, int batch_size, int has_sinks); + + +namespace cake_host_shim_b67714d8806118f0 { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t CakeDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +// Explicit immutable caller storage: initialize once, outside graph capture, +// exactly as the production TmaDeviceSlot path. The owner must retain the +// allocation and must never expose it as mutable scratch or rewrite its bytes. +template +static inline void PrepareImmutableCallerTmaWorkspace( + const tvm::ffi::Tensor& retained_workspace, + const CUtensorMap (&maps)[N], cudaStream_t stream) { + void* workspace = retained_workspace.data_ptr(); + CUcontext context = nullptr; + CUresult result = cuCtxGetCurrent(&context); + TVM_FFI_CHECK(result == CUDA_SUCCESS && context != nullptr, RuntimeError) + << "immutable TMA workspace requires an active CUDA context"; + CUcontext allocation_context = nullptr; + result = cuPointerGetAttribute( + &allocation_context, CU_POINTER_ATTRIBUTE_CONTEXT, + reinterpret_cast(workspace)); + TVM_FFI_CHECK(result == CUDA_SUCCESS && allocation_context == context, RuntimeError) + << "immutable TMA workspace must belong to the active CUDA context"; + unsigned long long buffer_id = 0; + result = cuPointerGetAttribute( + &buffer_id, CU_POINTER_ATTRIBUTE_BUFFER_ID, + reinterpret_cast(workspace)); + TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) + << "querying immutable TMA allocation identity failed"; + std::string key = std::to_string(reinterpret_cast(context)); + key += ":" + std::to_string(reinterpret_cast(workspace)); + key += ":" + std::to_string(buffer_id); + std::string descriptor_bytes(reinterpret_cast(maps), sizeof(maps)); + struct InitializedWorkspace { + tvm::ffi::Tensor owner; + std::string descriptor_bytes; + }; + static std::mutex mu; + // Keep the actual Tensor allocation alive: a caching allocator may reuse a + // freed suballocation without changing its enclosing CUDA buffer identity. + static auto* initialized = new std::unordered_map(); + std::lock_guard lock(mu); + auto it = initialized->find(key); + if (it != initialized->end()) { + TVM_FFI_CHECK(it->second.descriptor_bytes == descriptor_bytes, ValueError) + << "immutable TMA descriptors changed; provide new private descriptor storage"; + return; + } + CUstreamCaptureStatus capture_status = CU_STREAM_CAPTURE_STATUS_NONE; + result = cuStreamIsCapturing(reinterpret_cast(stream), &capture_status); + TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) + << "querying immutable TMA initialization capture state failed"; + TVM_FFI_CHECK(capture_status == CU_STREAM_CAPTURE_STATUS_NONE, RuntimeError) + << "immutable TMA storage must be initialized before CUDA Graph capture"; + // Synchronous host-to-device initialization completes before publishing the + // slot to any thread/stream, matching TmaDeviceSlot. The consumer retains its + // generated tensor-map acquire; the mutable upload's release fence is unchanged. + result = cuMemcpyHtoD(reinterpret_cast(workspace), maps, sizeof(maps)); + TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) + << "initializing immutable caller TMA descriptors failed"; + initialized->emplace( + std::move(key), InitializedWorkspace{retained_workspace, std::move(descriptor_bytes)}); +} + +static inline void* CallerTmaWorkspaceSlot( + const tvm::ffi::TensorView& workspace, size_t slot) { + return static_cast(workspace.data_ptr()) + slot * sizeof(CUtensorMap); +} + +// 4D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + int64_t d3 = t.size(t.ndim() - 3); + TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(d3)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 2u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 8, 2, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 8u, 2u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_swa_kv' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_swa_kv' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_kv'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_compressed_kv' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_compressed_kv(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_compressed_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_compressed_kv' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_compressed_kv' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_compressed_kv"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_compressed_kv' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_kv'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_kv') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_compressed_kv, TensorView arg_O, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_seq_lens, TensorView arg_cum_seq_lens_q, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_max_q_len, int64_t arg_batch_size, int64_t arg_has_sinks, tvm::ffi::Tensor arg_tma_descriptor_workspace, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckContiguous(arg_tmap_q, "tmap_q"); + CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); + CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 4, 16, 1); + CheckCudaTensor(arg_tmap_compressed_kv, "tmap_compressed_kv"); + CheckDtype(arg_tmap_compressed_kv, "tmap_compressed_kv", 4, 16, 1); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_swa_indices, "swa_indices"); + CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); + CheckContiguous(arg_swa_indices, "swa_indices"); + CheckCudaTensor(arg_compressed_indices, "compressed_indices"); + CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); + CheckContiguous(arg_compressed_indices, "compressed_indices"); + CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); + CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckCudaTensor(arg_seq_lens, "seq_lens"); + CheckDtype(arg_seq_lens, "seq_lens", 0, 32, 1); + CheckContiguous(arg_seq_lens, "seq_lens"); + CheckCudaTensor(arg_cum_seq_lens_q, "cum_seq_lens_q"); + CheckDtype(arg_cum_seq_lens_q, "cum_seq_lens_q", 0, 32, 1); + CheckContiguous(arg_cum_seq_lens_q, "cum_seq_lens_q"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); + CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); + CheckContiguous(arg_bmm1_scale, "bmm1_scale"); + CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); + CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); + CheckContiguous(arg_bmm2_scale, "bmm2_scale"); + TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) + << "scalar 'swa_index_stride' value " << arg_swa_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) + << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) + << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) + << "scalar 'sparse_topk' value " << arg_sparse_topk + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_q_len >= -2147483648LL && arg_max_q_len <= 2147483647LL, ValueError) + << "scalar 'max_q_len' value " << arg_max_q_len + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_batch_size >= -2147483648LL && arg_batch_size <= 2147483647LL, ValueError) + << "scalar 'batch_size' value " << arg_batch_size + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) + << "scalar 'has_sinks' value " << arg_has_sinks + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); + CheckSameCudaDevice(arg_tmap_compressed_kv, arg_tmap_q, "tmap_compressed_kv", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); + CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); + CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); + CheckSameCudaDevice(arg_seq_lens, arg_tmap_q, "seq_lens", "tmap_q"); + CheckSameCudaDevice(arg_cum_seq_lens_q, arg_tmap_q, "cum_seq_lens_q", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); + CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); + CheckCudaTensor(arg_tma_descriptor_workspace, "tma_descriptor_workspace"); + CheckSameCudaDevice(arg_tma_descriptor_workspace, arg_tmap_q, "tma_descriptor_workspace", "tmap_q"); + CheckContiguous(arg_tma_descriptor_workspace, "tma_descriptor_workspace"); + CheckDtype(arg_tma_descriptor_workspace, "tma_descriptor_workspace", 1, 8, 1); + TVM_FFI_CHECK(arg_tma_descriptor_workspace.numel() >= 384, ValueError) + << "tma_descriptor_workspace requires at least 384 bytes"; + TVM_FFI_CHECK( + reinterpret_cast(arg_tma_descriptor_workspace.data_ptr()) % 128 == 0, + ValueError) + << "tma_descriptor_workspace must be 128-byte aligned"; + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + + cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); + CUtensorMap h_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CakeTensorMap const* p_tmap_q = static_cast(CallerTmaWorkspaceSlot(arg_tma_descriptor_workspace, 0u)); + CUtensorMap h_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); + CakeTensorMap const* p_tmap_swa_kv = static_cast(CallerTmaWorkspaceSlot(arg_tma_descriptor_workspace, 1u)); + CUtensorMap h_tmap_compressed_kv = EncodeTma_tmap_compressed_kv(arg_tmap_compressed_kv); + CakeTensorMap const* p_tmap_compressed_kv = static_cast(CallerTmaWorkspaceSlot(arg_tma_descriptor_workspace, 2u)); + __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); + int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); + int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); + int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); + int* p_seq_lens = static_cast(arg_seq_lens.data_ptr()); + int* p_cum_seq_lens_q = static_cast(arg_cum_seq_lens_q.data_ptr()); + float* p_sinks = static_cast(arg_sinks.data_ptr()); + float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); + float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); + int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; + int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; + int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; + int32_t v_sparse_topk = (int32_t)arg_sparse_topk; + int32_t v_max_q_len = (int32_t)arg_max_q_len; + int32_t v_batch_size = (int32_t)arg_batch_size; + int32_t v_has_sinks = (int32_t)arg_has_sinks; + void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_tmap_compressed_kv, &p_O, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_seq_lens, &p_cum_seq_lens_q, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_max_q_len, &v_batch_size, &v_has_sinks}; + + static std::mutex smem_mu; + static auto* smem_status_by_device = + new std::unordered_map(); + cudaError_t smem_status = cudaSuccess; + { + std::lock_guard lock(smem_mu); + auto it = smem_status_by_device->find(dev.device_id); + if (it == smem_status_by_device->end()) { + int current_device = -1; + smem_status = cudaGetDevice(¤t_device); + if (smem_status == cudaSuccess && current_device != dev.device_id) { + smem_status = cudaErrorInvalidDevice; + } + if (smem_status == cudaSuccess) { + smem_status = cudaFuncSetAttribute( + reinterpret_cast(kernel_cake_dsv4_e48c5e0a89a41eb1e57a), + cudaFuncAttributeMaxDynamicSharedMemorySize, + 218752); + } + smem_status_by_device->emplace(dev.device_id, smem_status); + } else { + smem_status = it->second; + } + } + TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) + << "cudaFuncSetAttribute for kernel_cake_dsv4_e48c5e0a89a41eb1e57a failed: " + << cudaGetErrorString(smem_status); + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(640u, 1u, 1u); + const CUtensorMap workspace_maps[] = {h_tmap_q, h_tmap_swa_kv, h_tmap_compressed_kv}; + PrepareImmutableCallerTmaWorkspace(arg_tma_descriptor_workspace, workspace_maps, stream); + + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes = 218752u; + config.stream = stream; + cudaError_t launch_status = cudaLaunchKernelExC( + &config, reinterpret_cast(kernel_cake_dsv4_e48c5e0a89a41eb1e57a), kargs); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernelExC for kernel_cake_dsv4_e48c5e0a89a41eb1e57a failed: " + << cudaGetErrorString(launch_status); +} + +} // namespace cake_host_shim_b67714d8806118f0 + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_b67714d8806118f0::Run); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_e48c5e0a89a41eb1e57a_kernel.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_e48c5e0a89a41eb1e57a_kernel.cu new file mode 100644 index 00000000000..f841e132e77 --- /dev/null +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_e48c5e0a89a41eb1e57a_kernel.cu @@ -0,0 +1,2734 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); +template +struct __align__(128) CakeTensorMapPack { CakeTensorMap maps[N]; }; + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +#define CAKE_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SOURCE_EXACT_SCRATCH_OFFSET 0 +#define NUM_Q_PIPE_STAGES 1 +#define NUM_KV_PIPE_STAGES 6 +#define NUM_PAGE_PIPE_STAGES 6 +#define NUM_SCORE_PIPE_STAGES 2 +#define NUM_STATS_PIPE_STAGES 2 +#define NUM_P_PIPE_STAGES 2 +#define NUM_O_PIPE_STAGES 1 +#define NUM_WORK_PIPE_STAGES 2 +#define NUM_THROTTLE_PIPE_STAGES 2 +#define SMEM_SMEM_Q_OFF 1024 +#define SMEM_SMEM_Q_STAGE_BYTES 2048 +#define SMEM_SMEM_Q_STRIDE 2048 +#define SMEM_SMEM_KV_OFF 9216 +#define SMEM_SMEM_KV_STAGE_BYTES 32768 +#define SMEM_SMEM_KV_STRIDE 32768 +#define SMEM_SMEM_V_OFF 9216 +#define SMEM_SMEM_V_STAGE_BYTES 32768 +#define SMEM_SMEM_V_STRIDE 32768 +#define SMEM_SMEM_P_OFF 205824 +#define SMEM_SMEM_P_STAGE_BYTES 2048 +#define SMEM_SMEM_P_STRIDE 2048 +#define SMEM_SMEM_PAGE_OFFSETS_OFF 209920 +#define SMEM_SMEM_PAGE_OFFSETS_STAGE_BYTES 1024 +#define SMEM_SMEM_PAGE_OFFSETS_STRIDE 1024 +#define SMEM_SMEM_O_OFF 216064 +#define SMEM_SMEM_O_STAGE_BYTES 2048 +#define SMEM_SMEM_O_STRIDE 2048 +#define SMEM_SMEM_SOFTMAX_OFF 218112 +#define SMEM_SMEM_SOFTMAX_STAGE_BYTES 128 +#define SMEM_SMEM_SOFTMAX_STRIDE 128 +#define SMEM_SMEM_SOFTMAX_U32_OFF 218112 +#define SMEM_SMEM_SOFTMAX_U32_STAGE_BYTES 128 +#define SMEM_SMEM_SOFTMAX_U32_STRIDE 128 +#define SMEM_SMEM_SOFTMAX_UNUSED_OFF 218240 +#define SMEM_SMEM_SOFTMAX_UNUSED_STAGE_BYTES 128 +#define SMEM_SMEM_SOFTMAX_UNUSED_STRIDE 128 +#define SMEM_SMEM_CORR_OFF 218368 +#define SMEM_SMEM_CORR_STAGE_BYTES 128 +#define SMEM_SMEM_CORR_STRIDE 128 +#define SMEM_SMEM_STATS_OFF 218496 +#define SMEM_SMEM_STATS_STAGE_BYTES 96 +#define SMEM_SMEM_STATS_STRIDE 96 +#define SMEM_WORK_RESPONSE_OFF 218688 +#define SMEM_WORK_RESPONSE_STAGE_BYTES 16 +#define SMEM_WORK_RESPONSE_STRIDE 16 +#define SMEM_TOTAL 218752 +#define THREADS 640 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count) : "memory"); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x4(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x4.b32" + " {%0, %1, %2, %3}, [%4];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_2d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3}], [%4];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr) { + // Canonical .shared::cta form for non-multicast gather4, matching + // trtllm-gen / cuda_ptx and the PTX ISA qualifier order + // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, + // .shared::cluster is reserved for the multicast variant (ctaMask). + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x4_wait(float* dst, int addr) { + tmem_ld_x4(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(640, 1) void +kernel_cake_dsv4_e48c5e0a89a41eb1e57a(CakeTensorMap const* tmap_q, CakeTensorMap const* tmap_swa_kv, CakeTensorMap const* tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int max_q_len, int batch_size, int has_sinks) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int mbar_base = smem; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define kv_full_addr (mbar_base + 16) + #define kv_empty_addr (mbar_base + 64) + #define page_full_addr (mbar_base + 112) + #define page_empty_addr (mbar_base + 160) + #define score_full_addr (mbar_base + 208) + #define score_empty_addr (mbar_base + 224) + #define stats_full_addr (mbar_base + 240) + #define stats_empty_addr (mbar_base + 256) + #define p_full_addr (mbar_base + 272) + #define p_empty_addr (mbar_base + 288) + #define o_full_addr (mbar_base + 304) + #define o_empty_addr (mbar_base + 312) + #define work_full_addr (mbar_base + 320) + #define work_empty_addr (mbar_base + 336) + #define throttle_full_addr (mbar_base + 352) + #define throttle_empty_addr (mbar_base + 368) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + if (tid == 0) { + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_q)) : "memory"); + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_swa_kv)) : "memory"); + asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_compressed_kv)) : "memory"); + } + __syncthreads(); + + + // Kernel setup ops + __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_q_addr = smem + 1024; + __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 9216); + const int smem_kv_addr = smem + 9216; + __nv_bfloat16* smem_v = reinterpret_cast<__nv_bfloat16*>(smem_raw + 9216); + const int smem_v_addr = smem + 9216; + __nv_bfloat16* smem_p = reinterpret_cast<__nv_bfloat16*>(smem_raw + 205824); + const int smem_p_addr = smem + 205824; + int* smem_page_offsets = reinterpret_cast(smem_raw + 209920); + const int smem_page_offsets_addr = smem + 209920; + __nv_bfloat16* smem_o = reinterpret_cast<__nv_bfloat16*>(smem_raw + 216064); + const int smem_o_addr = smem + 216064; + float* smem_softmax = reinterpret_cast(smem_raw + 218112); + const int smem_softmax_addr = smem + 218112; + unsigned int* smem_softmax_u32 = reinterpret_cast(smem_raw + 218112); + const int smem_softmax_u32_addr = smem + 218112; + float* smem_softmax_unused = reinterpret_cast(smem_raw + 218240); + const int smem_softmax_unused_addr = smem + 218240; + float* smem_corr = reinterpret_cast(smem_raw + 218368); + const int smem_corr_addr = smem + 218368; + float* smem_stats = reinterpret_cast(smem_raw + 218496); + const int smem_stats_addr = smem + 218496; + unsigned int* work_response = reinterpret_cast(smem_raw + 218688); + const int work_response_addr = smem + 218688; + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)(tmap_q)) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)(tmap_swa_kv)) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)(tmap_compressed_kv)) : "memory"); + + // Mbarrier init (18 pipeline groups, 0 ordered-sequence groups, 48 barriers) + // Mbarriers at smem_raw[0..384) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // --- pipeline 'q_pipe' --- + // q_full: 1 barriers, init_count=1 + mbarrier_init(smem + 0, 1); + // q_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 8, 1); + // --- pipeline 'kv_pipe' --- + // kv_full: 6 barriers, init_count=1 + mbarrier_init(smem + 16, 1); + mbarrier_init(smem + 24, 1); + mbarrier_init(smem + 32, 1); + mbarrier_init(smem + 40, 1); + mbarrier_init(smem + 48, 1); + mbarrier_init(smem + 56, 1); + // kv_empty: 6 barriers, init_count=1 + mbarrier_init(smem + 64, 1); + mbarrier_init(smem + 72, 1); + mbarrier_init(smem + 80, 1); + mbarrier_init(smem + 88, 1); + mbarrier_init(smem + 96, 1); + mbarrier_init(smem + 104, 1); + // --- pipeline 'page_pipe' --- + // page_full: 6 barriers, init_count=32 + mbarrier_init(smem + 112, 32); + mbarrier_init(smem + 120, 32); + mbarrier_init(smem + 128, 32); + mbarrier_init(smem + 136, 32); + mbarrier_init(smem + 144, 32); + mbarrier_init(smem + 152, 32); + // page_empty: 6 barriers, init_count=256 + mbarrier_init(smem + 160, 256); + mbarrier_init(smem + 168, 256); + mbarrier_init(smem + 176, 256); + mbarrier_init(smem + 184, 256); + mbarrier_init(smem + 192, 256); + mbarrier_init(smem + 200, 256); + // --- pipeline 'score_pipe' --- + // score_full: 2 barriers, init_count=1 + mbarrier_init(smem + 208, 1); + mbarrier_init(smem + 216, 1); + // score_empty: 2 barriers, init_count=128 + mbarrier_init(smem + 224, 128); + mbarrier_init(smem + 232, 128); + // --- pipeline 'stats_pipe' --- + // stats_full: 2 barriers, init_count=128 + mbarrier_init(smem + 240, 128); + mbarrier_init(smem + 248, 128); + // stats_empty: 2 barriers, init_count=128 + mbarrier_init(smem + 256, 128); + mbarrier_init(smem + 264, 128); + // --- pipeline 'p_pipe' --- + // p_full: 2 barriers, init_count=128 + mbarrier_init(smem + 272, 128); + mbarrier_init(smem + 280, 128); + // p_empty: 2 barriers, init_count=1 + mbarrier_init(smem + 288, 1); + mbarrier_init(smem + 296, 1); + // --- pipeline 'o_pipe' --- + // o_full: 1 barriers, init_count=1 + mbarrier_init(smem + 304, 1); + // o_empty: 1 barriers, init_count=128 + mbarrier_init(smem + 312, 128); + // --- pipeline 'work_pipe' --- + // work_full: 2 barriers, init_count=1 + mbarrier_init(smem + 320, 1); + mbarrier_init(smem + 328, 1); + // work_empty: 2 barriers, init_count=640 + mbarrier_init(smem + 336, 640); + mbarrier_init(smem + 344, 640); + // --- pipeline 'throttle_pipe' --- + // throttle_full: 2 barriers, init_count=256 + mbarrier_init(smem + 352, 256); + mbarrier_init(smem + 360, 256); + // throttle_empty: 2 barriers, init_count=32 + mbarrier_init(smem + 368, 32); + mbarrier_init(smem + 376, 32); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncwarp(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 384); + if (warp == 0) { + int _tmem_hold = smem + 384; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_source_exact_scratch = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 48;"); + } + + // ---- Role: softmax ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 152;"); + { // softmax_main + unsigned int work_stage = 0; + unsigned int score_stage = 0; + unsigned int score_phase = 0; + unsigned int stats_stage = 0; + unsigned int stats_phase = 1; + unsigned int p_stage = 0; + unsigned int p_phase = 1; + unsigned int work_x = blockIdx.x; + unsigned int work_y = blockIdx.y; + unsigned int work_z = blockIdx.z; + const int warp_in_wg = warp; + const int wg_tid = warp_in_wg * 32 + lane; + const int col_pair_base = wg_tid % 4 * 2; + const int ldtm_row_base = warp_in_wg * 32 + lane / 4; + float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; + unsigned int _phase_work_full = 0; + #pragma unroll 1 + for (unsigned int _work = 0; _work < 120; _work++) { + int q_begin = cum_seq_lens_q[work_z]; + int q_end = cum_seq_lens_q[work_z + 1]; + int q_len = q_end - q_begin; + int query_idx = (unsigned int)q_begin + work_x; + int active_topk = 0; + if (work_x < (unsigned int)q_len) { + int _max_4 = ((sparse_topk_lens[query_idx] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx] + sparse_topk_lens_offset) : (0)); + int _min_7 = ((_max_4) < (sparse_topk) ? (_max_4) : (sparse_topk)); + active_topk = _min_7; + } + int num_steps = (active_topk + 128 - 1) / 128; + float row_max_pair[2]; + float row_sum_pair[2]; + row_max_pair[0] = -CAKE_INF; + row_max_pair[1] = -CAKE_INF; + row_sum_pair[0] = 0.0f; + row_sum_pair[1] = 0.0f; + asm volatile("barrier.sync 1, 128;" ::: "memory"); + uint32_t _amf_u_0 = __float_as_uint(-3.4028235e+38f); + uint32_t _amf_mask_0 = -int32_t(_amf_u_0 >> 31) | 0x80000000u; + unsigned int _amf_enc_0 = _amf_u_0 ^ _amf_mask_0; + if (wg_tid < 8) { + smem_softmax_u32[wg_tid] = _amf_enc_0; + } + asm volatile("barrier.sync 1, 128;" ::: "memory"); + #pragma unroll 1 + for (int tile = 0; tile < num_steps; tile++) { + mbarrier_wait(score_full_addr + (score_stage) * 8, score_phase); + int score_col = ((score_stage == 0) ? 0 : 8); + float scores[8]; + float score_lo[4]; + float score_hi[4]; + int tile_end = (tile + 1) * 128; + int is_full_tile = ((active_topk % 128 == 0) ? 1 : 0); + is_full_tile = is_full_tile | ((tile_end < active_topk) ? 1 : 0); + if (is_full_tile != 0) { + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x1.b32" + " {%0, %1, %2, %3}, [%4];" + : "=r"(*reinterpret_cast(&score_lo[0])), "=r"(*reinterpret_cast(&score_lo[1])), "=r"(*reinterpret_cast(&score_lo[2])), "=r"(*reinterpret_cast(&score_lo[3])) + : "r"(taddr + (unsigned int)score_col)); + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x1.b32" + " {%0, %1, %2, %3}, [%4];" + : "=r"(*reinterpret_cast(&score_hi[0])), "=r"(*reinterpret_cast(&score_hi[1])), "=r"(*reinterpret_cast(&score_hi[2])), "=r"(*reinterpret_cast(&score_hi[3])) + : "r"(taddr + (unsigned int)score_col + 1048576)); + #pragma unroll + for (int c = 0; c < 4; c++) { + scores[c] = score_lo[c]; + scores[c + 4] = score_hi[c]; + } + } else { + asm volatile("" ::: "memory"); + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x1.b32" + " {%0, %1, %2, %3}, [%4];" + : "=r"(*reinterpret_cast(&score_lo[0])), "=r"(*reinterpret_cast(&score_lo[1])), "=r"(*reinterpret_cast(&score_lo[2])), "=r"(*reinterpret_cast(&score_lo[3])) + : "r"(taddr + (unsigned int)score_col)); + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x1.b32" + " {%0, %1, %2, %3}, [%4];" + : "=r"(*reinterpret_cast(&score_hi[0])), "=r"(*reinterpret_cast(&score_hi[1])), "=r"(*reinterpret_cast(&score_hi[2])), "=r"(*reinterpret_cast(&score_hi[3])) + : "r"(taddr + (unsigned int)score_col + 1048576)); + #pragma unroll + for (int c_1 = 0; c_1 < 4; c_1++) { + scores[c_1] = score_lo[c_1]; + scores[c_1 + 4] = score_hi[c_1]; + } + if (active_topk <= tile * 128 + ldtm_row_base) { + scores[0] = -CAKE_INF; + scores[1] = -CAKE_INF; + } + if (active_topk <= tile * 128 + ldtm_row_base + 8) { + scores[2] = -CAKE_INF; + scores[3] = -CAKE_INF; + } + if (active_topk <= tile * 128 + ldtm_row_base + 16) { + scores[4] = -CAKE_INF; + scores[5] = -CAKE_INF; + } + if (active_topk <= tile * 128 + ldtm_row_base + 24) { + scores[6] = -CAKE_INF; + scores[7] = -CAKE_INF; + } + asm volatile("" ::: "memory"); + } + float pair_max[2]; + float _max_5 = max_noftz(scores[0], scores[2]); + float _max_6 = max_noftz(scores[4], scores[6]); + float _max_7 = max_noftz(_max_5, _max_6); + pair_max[0] = _max_7; + float _max_8 = max_noftz(scores[1], scores[3]); + float _max_9 = max_noftz(scores[5], scores[7]); + float _max_10 = max_noftz(_max_8, _max_9); + pair_max[1] = _max_10; + #pragma unroll + for (int c_2 = 0; c_2 < 2; c_2++) { + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, pair_max[c_2], 16); + float _max_11 = max_noftz(pair_max[c_2], _shfl_xor_0); + pair_max[c_2] = _max_11; + float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, pair_max[c_2], 8); + float _max_12 = max_noftz(pair_max[c_2], _shfl_xor_1); + pair_max[c_2] = _max_12; + } + float old_max_pair[2]; + float new_max_pair[2]; + #pragma unroll + for (int c_3 = 0; c_3 < 2; c_3++) { + old_max_pair[c_3] = row_max_pair[c_3]; + float _max_13 = max_noftz(row_max_pair[c_3], pair_max[c_3]); + new_max_pair[c_3] = _max_13; + } + if (lane < 8) { + uint32_t _amf_u_1 = __float_as_uint(new_max_pair[0]); + uint32_t _amf_mask_1 = -int32_t(_amf_u_1 >> 31) | 0x80000000u; + unsigned int _amf_enc_1 = _amf_u_1 ^ _amf_mask_1; + uint32_t _amf_u_2 = __float_as_uint(new_max_pair[1]); + uint32_t _amf_mask_2 = -int32_t(_amf_u_2 >> 31) | 0x80000000u; + unsigned int _amf_enc_2 = _amf_u_2 ^ _amf_mask_2; + atomicMax(&smem_softmax_u32[col_pair_base], _amf_enc_1); + atomicMax(&smem_softmax_u32[col_pair_base + 1], _amf_enc_2); + } + asm volatile("barrier.sync 1, 128;" ::: "memory"); + uint32_t _amf_u_3 = smem_softmax_u32[col_pair_base]; + uint32_t _amf_mask_3 = ((_amf_u_3 >> 31) - 1u) | 0x80000000u; + float _amf_dec_0 = __uint_as_float(_amf_u_3 ^ _amf_mask_3); + new_max_pair[0] = _amf_dec_0; + uint32_t _amf_u_4 = smem_softmax_u32[col_pair_base + 1]; + uint32_t _amf_mask_4 = ((_amf_u_4 >> 31) - 1u) | 0x80000000u; + float _amf_dec_1 = __uint_as_float(_amf_u_4 ^ _amf_mask_4); + new_max_pair[1] = _amf_dec_1; + mbarrier_wait(stats_empty_addr + (stats_stage) * 8, stats_phase); + float stats_pair[4]; + stats_pair[0] = old_max_pair[0]; + stats_pair[1] = old_max_pair[1]; + stats_pair[2] = new_max_pair[0]; + stats_pair[3] = new_max_pair[1]; + int stats_col = ((stats_stage == 0) ? 16 : 48); + tmem_st_x4_f32(taddr + (unsigned int)stats_col, stats_pair); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(stats_full_addr + (stats_stage) * 8); + stats_stage += 1; + if (stats_stage == 2) { stats_stage = 0; stats_phase ^= 1; } + float acc_scale_pair[2]; + #pragma unroll + for (int c_4 = 0; c_4 < 2; c_4++) { + float delta = softmax_scale_log2 * (old_max_pair[c_4] - new_max_pair[c_4]); + float _exp2_0 = approx_exp2(delta); + acc_scale_pair[c_4] = ((old_max_pair[c_4] > -CAKE_INF) ? _exp2_0 : 1.0f); + row_max_pair[c_4] = new_max_pair[c_4]; + } + asm volatile("" ::: "memory"); + float exp_values[8]; + #pragma unroll + for (int c_5 = 0; c_5 < 8; c_5++) { + float safe_max = ((new_max_pair[c_5 % 2] == -CAKE_INF) ? 0.0f : new_max_pair[c_5 % 2]); + float _exp2_1 = approx_exp2(scores[c_5] * softmax_scale_log2 - safe_max * softmax_scale_log2); + exp_values[c_5] = _exp2_1; + } + mbarrier_wait(p_empty_addr + (p_stage) * 8, p_phase); + unsigned int regs_p[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(exp_values[_lp*2 + 0], exp_values[_lp*2+1 + 0])); + regs_p[_lp] = *(uint32_t*)&_bf2; + } + int slice_idx = warp_in_wg / 2; + int warp_idx_in_slice = warp_in_wg % 2; + int mtx_idx = lane / 8; + int thr_row_idx = lane % 8; + int seg_col_idx = warp_idx_in_slice * 4 + mtx_idx ^ thr_row_idx; + int stsm_offset = slice_idx * 8 * 128 + thr_row_idx * 128 + seg_col_idx * 16; + asm volatile("stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 [%0], {%1, %2, %3, %4};\n" + :: "l"(__cvta_generic_to_shared(reinterpret_cast(smem_p) + (p_stage * 2048 + (unsigned int)stsm_offset))), "r"(*reinterpret_cast(®s_p[0])), "r"(*reinterpret_cast(®s_p[1])), "r"(*reinterpret_cast(®s_p[2])), "r"(*reinterpret_cast(®s_p[3])) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(p_full_addr + (p_stage) * 8); + p_stage += 1; + if (p_stage == 2) { p_stage = 0; p_phase ^= 1; } + float2 _f2_0 = make_float2(exp_values[0], exp_values[1]); + float2 _f2_1 = make_float2(exp_values[2], exp_values[3]); + float2 _f2_2 = make_float2(exp_values[4], exp_values[5]); + float2 _f2_3 = make_float2(exp_values[6], exp_values[7]); + float2 _f2_4 = make_float2(row_sum_pair[0], row_sum_pair[1]); + float2 _f2_5 = make_float2(acc_scale_pair[0], acc_scale_pair[1]); + float2 sum_p0_pair = fma_f32x2_rn_ftz(_f2_4, _f2_5, _f2_0); + float2 sum_p1_pair = add_f32x2(sum_p0_pair, _f2_1); + float2 sum_p2_pair = add_f32x2(sum_p1_pair, _f2_2); + float2 next_sum_pair = add_f32x2(sum_p2_pair, _f2_3); + row_sum_pair[0] = next_sum_pair.x; + row_sum_pair[1] = next_sum_pair.y; + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + mbarrier_arrive(score_empty_addr + (score_stage) * 8); + score_stage += 1; + if (score_stage == 2) { score_stage = 0; score_phase ^= 1; } + } + if (num_steps > 0) { + mbarrier_wait(stats_empty_addr + (stats_stage) * 8, stats_phase); + float final_stats_pair[4]; + final_stats_pair[0] = row_sum_pair[0]; + final_stats_pair[1] = row_sum_pair[1]; + final_stats_pair[2] = row_max_pair[0]; + final_stats_pair[3] = row_max_pair[1]; + int final_stats_col = ((stats_stage == 0) ? 16 : 48); + tmem_st_x4_f32(taddr + (unsigned int)final_stats_col, final_stats_pair); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(stats_full_addr + (stats_stage) * 8); + stats_stage += 1; + if (stats_stage == 2) { stats_stage = 0; stats_phase ^= 1; } + } + mbarrier_wait(work_full_addr + (work_stage) * 8, _phase_work_full); + uint32_t _clc_valid_3 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_3) + : "r"(work_response_addr + work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_9 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_9) + : "r"(work_response_addr + work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_10 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_10) + : "r"(work_response_addr + work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_11 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_11) + : "r"(work_response_addr + work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + work_stage * 8), "r"(0) : "memory"); + work_stage += 1; + if (work_stage == 2) { work_stage = 0; _phase_work_full ^= 1; } + if (_clc_valid_3 == 0) { + break; + } + work_x = _clc_ctaid_9; + work_y = _clc_ctaid_10; + work_z = _clc_ctaid_11; + } + } + } + // ---- Role: correction ---- + if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 136;"); + { // correction_main + unsigned int work_stage_1 = 0; + unsigned int stats_stage_1 = 0; + unsigned int stats_phase_1 = 0; + unsigned int work_x_1 = blockIdx.x; + unsigned int work_y_1 = blockIdx.y; + unsigned int work_z_1 = blockIdx.z; + const int corr_rank = warp - 4; + const int corr_tid = corr_rank * 32 + lane; + const int col_pair_base_1 = corr_tid % 4 * 2; + float softmax_scale_log2_1 = bmm1_scale[0] * 1.4426950408889634f; + float output_scale = bmm2_scale[0]; + unsigned int _phase_o_full_0 = 0; + unsigned int _phase_work_full_1 = 0; + #pragma unroll 1 + for (unsigned int _work_1 = 0; _work_1 < 120; _work_1++) { + int q_begin_1 = cum_seq_lens_q[work_z_1]; + int q_end_1 = cum_seq_lens_q[work_z_1 + 1]; + int q_len_1 = q_end_1 - q_begin_1; + int query_idx_1 = (unsigned int)q_begin_1 + work_x_1; + int active_topk_1 = 0; + if (work_x_1 < (unsigned int)q_len_1) { + int _max_14 = ((sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) : (0)); + int _min_8 = ((_max_14) < (sparse_topk) ? (_max_14) : (sparse_topk)); + active_topk_1 = _min_8; + } + int num_steps_1 = (active_topk_1 + 128 - 1) / 128; + int head_group = work_y_1 / 4; + int value_quarter = work_y_1 % 4; + if (num_steps_1 > 0) { + mbarrier_wait(stats_full_addr + (stats_stage_1) * 8, stats_phase_1); + mbarrier_arrive(stats_empty_addr + (stats_stage_1) * 8); + stats_stage_1 += 1; + if (stats_stage_1 == 2) { stats_stage_1 = 0; stats_phase_1 ^= 1; } + } + #pragma unroll 1 + for (int tile_1 = 1; tile_1 < num_steps_1; tile_1++) { + mbarrier_wait(stats_full_addr + (stats_stage_1) * 8, stats_phase_1); + int stats_col_1 = ((stats_stage_1 == 0) ? 16 : 48); + float _tmem_load_0[4]; + tmem_ld_x4(&_tmem_load_0[0], taddr + (unsigned int)stats_col_1); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + float scale_pair[2]; + #pragma unroll + for (int c_6 = 0; c_6 < 2; c_6++) { + float max_diff = _tmem_load_0[c_6] - _tmem_load_0[c_6 + 2]; + float _exp2_2 = approx_exp2(softmax_scale_log2_1 * max_diff); + scale_pair[c_6] = ((max_diff != 0.0f) ? _exp2_2 : 1.0f); + } + mbarrier_arrive(stats_empty_addr + (stats_stage_1) * 8); + stats_stage_1 += 1; + if (stats_stage_1 == 2) { stats_stage_1 = 0; stats_phase_1 ^= 1; } + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int needs_scale = ((scale_pair[0] != 1.0f) ? 1 : 0); + needs_scale = needs_scale | ((scale_pair[1] != 1.0f) ? 1 : 0); + int _vote_0 = __any_sync(0xFFFFFFFF, needs_scale != 0); + needs_scale = _vote_0; + if (needs_scale != 0) { + float values_lo[4]; + float values_hi[4]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x1.b32" + " {%0, %1, %2, %3}, [%4];" + : "=r"(*reinterpret_cast(&values_lo[0])), "=r"(*reinterpret_cast(&values_lo[1])), "=r"(*reinterpret_cast(&values_lo[2])), "=r"(*reinterpret_cast(&values_lo[3])) + : "r"(taddr + 80)); + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x1.b32" + " {%0, %1, %2, %3}, [%4];" + : "=r"(*reinterpret_cast(&values_hi[0])), "=r"(*reinterpret_cast(&values_hi[1])), "=r"(*reinterpret_cast(&values_hi[2])), "=r"(*reinterpret_cast(&values_hi[3])) + : "r"(taddr + 80 + 1048576)); + #pragma unroll + for (int pair = 0; pair < 4; pair++) { + asm volatile("" ::: "memory"); + values_lo[pair] = values_lo[pair] * scale_pair[pair % 2]; + values_hi[pair] = values_hi[pair] * scale_pair[pair % 2]; + } + asm volatile( + "tcgen05.st.sync.aligned.16x256b.x1.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(taddr + 80), "r"(*reinterpret_cast(&values_lo[0])), "r"(*reinterpret_cast(&values_lo[1])), "r"(*reinterpret_cast(&values_lo[2])), "r"(*reinterpret_cast(&values_lo[3]))); + asm volatile( + "tcgen05.st.sync.aligned.16x256b.x1.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(taddr + 80 + 1048576), "r"(*reinterpret_cast(&values_hi[0])), "r"(*reinterpret_cast(&values_hi[1])), "r"(*reinterpret_cast(&values_hi[2])), "r"(*reinterpret_cast(&values_hi[3]))); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + } + mbarrier_arrive(o_empty_addr); + } + if (num_steps_1 > 0) { + mbarrier_wait(stats_full_addr + (stats_stage_1) * 8, stats_phase_1); + int final_stats_col_1 = ((stats_stage_1 == 0) ? 16 : 48); + float _tmem_load_1[4]; + tmem_ld_x4(&_tmem_load_1[0], taddr + (unsigned int)final_stats_col_1); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + mbarrier_arrive(stats_empty_addr + (stats_stage_1) * 8); + stats_stage_1 += 1; + if (stats_stage_1 == 2) { stats_stage_1 = 0; stats_phase_1 ^= 1; } + float pair_sum[2]; + pair_sum[0] = _tmem_load_1[0]; + pair_sum[1] = _tmem_load_1[1]; + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, pair_sum[0], 16); + pair_sum[0] = pair_sum[0] + _shfl_xor_2; + float _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, pair_sum[1], 16); + pair_sum[1] = pair_sum[1] + _shfl_xor_3; + float _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, pair_sum[0], 8); + pair_sum[0] = pair_sum[0] + _shfl_xor_4; + float _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, pair_sum[1], 8); + pair_sum[1] = pair_sum[1] + _shfl_xor_5; + float _shfl_xor_6 = __shfl_xor_sync(0xFFFFFFFF, pair_sum[0], 4); + pair_sum[0] = pair_sum[0] + _shfl_xor_6; + float _shfl_xor_7 = __shfl_xor_sync(0xFFFFFFFF, pair_sum[1], 4); + pair_sum[1] = pair_sum[1] + _shfl_xor_7; + if (lane < 4) { + smem_corr[corr_rank * 8 + col_pair_base_1] = pair_sum[0]; + smem_corr[corr_rank * 8 + col_pair_base_1 + 1] = pair_sum[1]; + } + asm volatile("barrier.sync 2, 128;" ::: "memory"); + float sum0 = smem_corr[col_pair_base_1] + smem_corr[8 + col_pair_base_1] + smem_corr[16 + col_pair_base_1] + smem_corr[24 + col_pair_base_1]; + float sum1 = smem_corr[col_pair_base_1 + 1] + smem_corr[8 + col_pair_base_1 + 1] + smem_corr[16 + col_pair_base_1 + 1] + smem_corr[24 + col_pair_base_1 + 1]; + if (has_sinks != 0) { + float _exp2_3 = approx_exp2(sinks[head_group * 8 + col_pair_base_1] * 1.4426950408889634f - _tmem_load_1[2] * softmax_scale_log2_1); + sum0 = sum0 + _exp2_3; + float _exp2_4 = approx_exp2(sinks[head_group * 8 + col_pair_base_1 + 1] * 1.4426950408889634f - _tmem_load_1[3] * softmax_scale_log2_1); + sum1 = sum1 + _exp2_4; + } + float _rcp_0 = approx_rcp(sum0); + float scale0 = ((sum0 > 0.0f) ? output_scale * _rcp_0 : 0.0f); + float _rcp_1 = approx_rcp(sum1); + float scale1 = ((sum1 > 0.0f) ? output_scale * _rcp_1 : 0.0f); + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + float output_lo[4]; + float output_hi[4]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x1.b32" + " {%0, %1, %2, %3}, [%4];" + : "=r"(*reinterpret_cast(&output_lo[0])), "=r"(*reinterpret_cast(&output_lo[1])), "=r"(*reinterpret_cast(&output_lo[2])), "=r"(*reinterpret_cast(&output_lo[3])) + : "r"(taddr + 80)); + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x1.b32" + " {%0, %1, %2, %3}, [%4];" + : "=r"(*reinterpret_cast(&output_hi[0])), "=r"(*reinterpret_cast(&output_hi[1])), "=r"(*reinterpret_cast(&output_hi[2])), "=r"(*reinterpret_cast(&output_hi[3])) + : "r"(taddr + 80 + 1048576)); + unsigned int packed_o[4]; + #pragma unroll + for (int pair_1 = 0; pair_1 < 4; pair_1++) { + asm volatile("" ::: "memory"); + output_lo[pair_1] = output_lo[pair_1] * ((pair_1 % 2 == 0) ? scale0 : scale1); + output_hi[pair_1] = output_hi[pair_1] * ((pair_1 % 2 == 0) ? scale0 : scale1); + } + float output_values[8]; + #pragma unroll + for (int c_7 = 0; c_7 < 4; c_7++) { + output_values[c_7] = output_lo[c_7]; + output_values[c_7 + 4] = output_hi[c_7]; + } + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(output_values[_lp*2 + 0], output_values[_lp*2+1 + 0])); + packed_o[_lp] = *(uint32_t*)&_bf2; + } + int slice_idx_1 = corr_rank / 2; + int warp_idx_in_slice_1 = corr_rank % 2; + int mtx_idx_1 = lane / 8; + int thr_row_idx_1 = lane % 8; + int seg_col_idx_1 = warp_idx_in_slice_1 * 4 + mtx_idx_1 ^ thr_row_idx_1; + int stsm_offset_1 = slice_idx_1 * 8 * 128 + thr_row_idx_1 * 128 + seg_col_idx_1 * 16; + asm volatile("stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 [%0], {%1, %2, %3, %4};\n" + :: "l"(__cvta_generic_to_shared(reinterpret_cast(smem_o) + stsm_offset_1)), "r"(*reinterpret_cast(&packed_o[0])), "r"(*reinterpret_cast(&packed_o[1])), "r"(*reinterpret_cast(&packed_o[2])), "r"(*reinterpret_cast(&packed_o[3])) + : "memory"); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + int copy_base = corr_tid * 16; + int smem_row = copy_base / 128; + int copy_smem_offset = copy_base ^ smem_row % 8 * 16; + int copy_row = smem_row % 8; + int copy_col = (smem_row / 8 * 128 + copy_base % 128) / 2; + unsigned int copy_vec[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(©_vec[0])), "=r"(*reinterpret_cast(©_vec[(0) + 1])), "=r"(*reinterpret_cast(©_vec[(0) + 2])), "=r"(*reinterpret_cast(©_vec[(0) + 3])) + : "r"(smem_o_addr + (unsigned int)copy_smem_offset)); + int head_groups = gridDim.y / 4; + int output_base = query_idx_1 * head_groups * 8 * 512 + head_group * 8 * 512 + value_quarter * 128; + reinterpret_cast(O + (output_base + copy_row * 512 + copy_col))[0] = reinterpret_cast(copy_vec)[0]; + mbarrier_arrive(o_empty_addr); + } + mbarrier_wait(work_full_addr + (work_stage_1) * 8, _phase_work_full_1); + uint32_t _clc_valid_4 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_4) + : "r"(work_response_addr + work_stage_1 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_12 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_12) + : "r"(work_response_addr + work_stage_1 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_13 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_13) + : "r"(work_response_addr + work_stage_1 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_14 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_14) + : "r"(work_response_addr + work_stage_1 * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + work_stage_1 * 8), "r"(0) : "memory"); + work_stage_1 += 1; + if (work_stage_1 == 2) { work_stage_1 = 0; _phase_work_full_1 ^= 1; } + if (_clc_valid_4 == 0) { + break; + } + work_x_1 = _clc_ctaid_12; + work_y_1 = _clc_ctaid_13; + work_z_1 = _clc_ctaid_14; + } + asm volatile("barrier.sync 11, 128;" ::: "memory"); + if (corr_rank == 0) { + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + } + } + // ---- Role: mma ---- + if (warp == 8) { + { // mma_main + unsigned int work_stage_2 = 0; + unsigned int q_stage = 0; + unsigned int q_phase = 0; + unsigned int kv_stage = 0; + unsigned int kv_phase = 0; + unsigned int score_stage_1 = 0; + unsigned int score_phase_1 = 1; + unsigned int p_stage_1 = 0; + unsigned int p_phase_1 = 0; + unsigned int work_x_2 = blockIdx.x; + unsigned int work_y_2 = blockIdx.y; + unsigned int work_z_2 = blockIdx.z; + unsigned int _phase_o_empty_0 = 1; + unsigned int _phase_work_full_2 = 0; + #pragma unroll 1 + for (unsigned int _work_2 = 0; _work_2 < 120; _work_2++) { + int q_begin_2 = cum_seq_lens_q[work_z_2]; + int q_end_2 = cum_seq_lens_q[work_z_2 + 1]; + int q_len_2 = q_end_2 - q_begin_2; + int query_idx_2 = (unsigned int)q_begin_2 + work_x_2; + int active_topk_2 = 0; + if (work_x_2 < (unsigned int)q_len_2) { + int _max_3 = ((sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) : (0)); + int _min_6 = ((_max_3) < (sparse_topk) ? (_max_3) : (sparse_topk)); + active_topk_2 = _min_6; + } + int num_steps_2 = (active_topk_2 + 128 - 1) / 128; + int value_quarter_1 = work_y_2 % 4; + if (num_steps_2 > 0) { + mbarrier_wait(q_full_addr + (q_stage) * 8, q_phase); + int first_output = 1; + mbarrier_wait(score_empty_addr + (score_stage_1) * 8, score_phase_1); + int score_col_1 = ((score_stage_1 == 0) ? 0 : 8); + #pragma unroll + for (int hs = 0; hs < 3; hs++) { + int hs_sel = ((value_quarter_1 > hs) ? hs : hs + 1); + mbarrier_wait(kv_full_addr + (kv_stage) * 8, kv_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_0 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage) * 2048); + int _mma_b_lo_0 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (hs_sel) * 128); + { + uint64_t _mma_ss_a_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_0); + uint64_t _mma_ss_b_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_0); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134349968, ((hs == 0) ? 0 : 1)); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134349968, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134349968, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134349968, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 1018U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 58U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134349968, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134349968, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134349968, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134349968, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage) * 8); + kv_stage += 1; + if (kv_stage == 6) { kv_stage = 0; kv_phase ^= 1; } + } + mbarrier_wait(kv_full_addr + (kv_stage) * 8, kv_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_1 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage) * 2048); + int _mma_b_lo_1 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (value_quarter_1) * 128); + { + uint64_t _mma_ss_a_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_1); + uint64_t _mma_ss_b_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_1); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134349968, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134349968, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134349968, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134349968, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 1018U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 58U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134349968, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134349968, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134349968, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134349968, 1); + } + } + kv_stage += 1; + if (kv_stage == 6) { kv_stage = 0; kv_phase ^= 1; } + elect_commit(score_full_addr + (score_stage_1) * 8); + score_stage_1 += 1; + if (score_stage_1 == 2) { score_stage_1 = 0; score_phase_1 ^= 1; } + #pragma unroll 1 + for (int _tile = 0; _tile < num_steps_2 - 1; _tile++) { + mbarrier_wait(score_empty_addr + (score_stage_1) * 8, score_phase_1); + int score_col_0 = ((score_stage_1 == 0) ? 0 : 8); + #pragma unroll + for (int hs_1 = 0; hs_1 < 3; hs_1++) { + int hs_sel_1 = ((value_quarter_1 > hs_1) ? hs_1 : hs_1 + 1); + mbarrier_wait(kv_full_addr + (kv_stage) * 8, kv_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_2 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage) * 2048); + int _mma_b_lo_2 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (hs_sel_1) * 128); + { + uint64_t _mma_ss_a_desc_2 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_2); + uint64_t _mma_ss_b_desc_2 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_2); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134349968, ((hs_1 == 0) ? 0 : 1)); + } + incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134349968, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134349968, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134349968, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_2, 1018U); + incr_smem_desc_lo(_mma_ss_b_desc_2, 58U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134349968, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134349968, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134349968, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134349968, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage) * 8); + kv_stage += 1; + if (kv_stage == 6) { kv_stage = 0; kv_phase ^= 1; } + } + mbarrier_wait(kv_full_addr + (kv_stage) * 8, kv_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_3 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage) * 2048); + int _mma_b_lo_3 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (value_quarter_1) * 128); + { + uint64_t _mma_ss_a_desc_3 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_3); + uint64_t _mma_ss_b_desc_3 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_3); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134349968, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134349968, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134349968, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134349968, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_3, 1018U); + incr_smem_desc_lo(_mma_ss_b_desc_3, 58U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134349968, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134349968, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134349968, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134349968, 1); + } + } + kv_stage += 1; + if (kv_stage == 6) { kv_stage = 0; kv_phase ^= 1; } + elect_commit(score_full_addr + (score_stage_1) * 8); + score_stage_1 += 1; + if (score_stage_1 == 2) { score_stage_1 = 0; score_phase_1 ^= 1; } + unsigned int v_cur = (kv_stage + 1) % 6; + mbarrier_wait(p_full_addr + (p_stage_1) * 8, p_phase_1); + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_4 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (v_cur) * 2048); + int _mma_b_lo_4 = make_warp_uniform(((((smem_p_addr) >> 4) & 0x3FFF) | 0x400000) + (p_stage_1) * 128); + { + uint64_t _mma_ss_a_desc_4 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_4); + uint64_t _mma_ss_b_desc_4 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_4); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134382736, ((first_output) ? 0 : 1)); + } + incr_smem_desc_lo(_mma_ss_a_desc_4, 128U); + incr_smem_desc_lo(_mma_ss_b_desc_4, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134382736, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_4, 128U); + incr_smem_desc_lo(_mma_ss_b_desc_4, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134382736, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_4, 128U); + incr_smem_desc_lo(_mma_ss_b_desc_4, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134382736, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_4, 128U); + incr_smem_desc_lo(_mma_ss_b_desc_4, 58U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134382736, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_4, 128U); + incr_smem_desc_lo(_mma_ss_b_desc_4, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134382736, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_4, 128U); + incr_smem_desc_lo(_mma_ss_b_desc_4, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134382736, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_4, 128U); + incr_smem_desc_lo(_mma_ss_b_desc_4, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134382736, 1); + } + } + elect_commit(o_full_addr); + elect_commit(kv_empty_addr + (v_cur) * 8); + elect_commit(p_empty_addr + (p_stage_1) * 8); + p_stage_1 += 1; + if (p_stage_1 == 2) { p_stage_1 = 0; p_phase_1 ^= 1; } + first_output = 0; + } + unsigned int v_tail = (kv_stage + 5) % 6; + mbarrier_wait(p_full_addr + (p_stage_1) * 8, p_phase_1); + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_5 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (v_tail) * 2048); + int _mma_b_lo_5 = make_warp_uniform(((((smem_p_addr) >> 4) & 0x3FFF) | 0x400000) + (p_stage_1) * 128); + { + uint64_t _mma_ss_a_desc_5 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_5); + uint64_t _mma_ss_b_desc_5 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_5); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134382736, ((first_output) ? 0 : 1)); + } + incr_smem_desc_lo(_mma_ss_a_desc_5, 128U); + incr_smem_desc_lo(_mma_ss_b_desc_5, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134382736, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_5, 128U); + incr_smem_desc_lo(_mma_ss_b_desc_5, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134382736, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_5, 128U); + incr_smem_desc_lo(_mma_ss_b_desc_5, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134382736, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_5, 128U); + incr_smem_desc_lo(_mma_ss_b_desc_5, 58U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134382736, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_5, 128U); + incr_smem_desc_lo(_mma_ss_b_desc_5, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134382736, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_5, 128U); + incr_smem_desc_lo(_mma_ss_b_desc_5, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134382736, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_5, 128U); + incr_smem_desc_lo(_mma_ss_b_desc_5, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134382736, 1); + } + } + elect_commit(o_full_addr); + elect_commit(kv_empty_addr + (v_tail) * 8); + elect_commit(p_empty_addr + (p_stage_1) * 8); + p_stage_1 += 1; + if (p_stage_1 == 2) { p_stage_1 = 0; p_phase_1 ^= 1; } + elect_commit(q_empty_addr + (q_stage) * 8); + q_phase ^= 1; + } + mbarrier_wait(work_full_addr + (work_stage_2) * 8, _phase_work_full_2); + uint32_t _clc_valid_2 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_2) + : "r"(work_response_addr + work_stage_2 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_6 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_6) + : "r"(work_response_addr + work_stage_2 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_7 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_7) + : "r"(work_response_addr + work_stage_2 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_8 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_8) + : "r"(work_response_addr + work_stage_2 * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + work_stage_2 * 8), "r"(0) : "memory"); + work_stage_2 += 1; + if (work_stage_2 == 2) { work_stage_2 = 0; _phase_work_full_2 ^= 1; } + if (_clc_valid_2 == 0) { + break; + } + work_x_2 = _clc_ctaid_6; + work_y_2 = _clc_ctaid_7; + work_z_2 = _clc_ctaid_8; + } + #pragma unroll + for (int tail = 0; tail < 2; tail++) { + elect_commit(score_full_addr + (tail) * 8); + } + } + } + // ---- Role: page_loader ---- + if (warp == 9) { + { // page_loader_main + unsigned int work_stage_3 = 0; + unsigned int page_stage = 0; + unsigned int page_phase = 1; + unsigned int work_x_3 = blockIdx.x; + unsigned int work_y_3 = blockIdx.y; + unsigned int work_z_3 = blockIdx.z; + unsigned int _phase_work_full_3 = 0; + #pragma unroll 1 + for (unsigned int _work_3 = 0; _work_3 < 120; _work_3++) { + int q_begin_3 = cum_seq_lens_q[work_z_3]; + int q_end_3 = cum_seq_lens_q[work_z_3 + 1]; + int q_len_3 = q_end_3 - q_begin_3; + int query_idx_3 = (unsigned int)q_begin_3 + work_x_3; + int active_topk_3 = 0; + if (work_x_3 < (unsigned int)q_len_3) { + int _max_0 = ((sparse_topk_lens[query_idx_3] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_3] + sparse_topk_lens_offset) : (0)); + int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); + active_topk_3 = _min_0; + } + int num_steps_3 = (active_topk_3 + 128 - 1) / 128; + int num_pairs = (num_steps_3 + 1) / 2; + int _max_1 = ((active_topk_3 - 1 & -4) > (0) ? (active_topk_3 - 1 & -4) : (0)); + int last_aligned = _max_1; + int* swa_row = swa_indices + (query_idx_3 * swa_index_stride); + int* compressed_row = compressed_indices + (query_idx_3 * compressed_index_stride); + #pragma unroll 1 + for (int pair_2 = 0; pair_2 < num_pairs; pair_2++) { + int lane_base = lane * 4; + #pragma unroll + for (int duplicate = 0; duplicate < 2; duplicate++) { + mbarrier_wait(page_empty_addr + (page_stage) * 8, page_phase); + int page0 = pair_2 * 2 * 128 + lane_base; + int page1 = page0 + 128; + int safe0 = ((page0 <= last_aligned) ? page0 : last_aligned); + int safe1 = ((page1 <= last_aligned) ? page1 : last_aligned); + int* src0 = ((safe0 < 128) ? swa_row + safe0 : compressed_row + (safe0 - 128)); + int* src1 = ((safe1 < 128) ? swa_row + safe1 : compressed_row + (safe1 - 128)); + asm volatile("cp.async.ca.shared::cta.global [%0], [%1], 16;" + :: "r"(smem_page_offsets_addr + page_stage * 1024 + (unsigned int)(lane_base * 4)), "l"(src0)); + asm volatile("cp.async.ca.shared::cta.global [%0], [%1], 16;" + :: "r"(smem_page_offsets_addr + page_stage * 1024 + (unsigned int)((128 + lane_base) * 4)), "l"(src1)); + asm volatile( + "{\n\t" + "cp.async.mbarrier.arrive.shared::cta.b64 [%0];\n\t" + "}" + :: "r"(page_full_addr + (page_stage) * 8) : "memory"); + mbarrier_arrive(page_full_addr + (page_stage) * 8); + page_stage += 1; + if (page_stage == 6) { page_stage = 0; page_phase ^= 1; } + } + } + mbarrier_wait(work_full_addr + (work_stage_3) * 8, _phase_work_full_3); + uint32_t _clc_valid_0 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_0) + : "r"(work_response_addr + work_stage_3 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_0 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_0) + : "r"(work_response_addr + work_stage_3 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_1 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_1) + : "r"(work_response_addr + work_stage_3 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_2 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_2) + : "r"(work_response_addr + work_stage_3 * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + work_stage_3 * 8), "r"(0) : "memory"); + work_stage_3 += 1; + if (work_stage_3 == 2) { work_stage_3 = 0; _phase_work_full_3 ^= 1; } + if (_clc_valid_0 == 0) { + break; + } + work_x_3 = _clc_ctaid_0; + work_y_3 = _clc_ctaid_1; + work_z_3 = _clc_ctaid_2; + } + } + } + // ---- Role: scheduler ---- + if (warp == 10) { + { // scheduler_main + asm volatile("griddepcontrol.wait;" ::: "memory"); + unsigned int work_stage_4 = 0; + unsigned int throttle_stage = 0; + unsigned int throttle_phase = 0; + unsigned int _phase_work_empty = 1; + unsigned int _phase_work_full_4 = 0; + #pragma unroll 1 + for (unsigned int _work_4 = 0; _work_4 < 120; _work_4++) { + mbarrier_wait(throttle_full_addr + (throttle_stage) * 8, throttle_phase); + mbarrier_arrive(throttle_empty_addr + (throttle_stage) * 8); + throttle_stage += 1; + if (throttle_stage == 2) { throttle_stage = 0; throttle_phase ^= 1; } + mbarrier_wait(work_empty_addr + (work_stage_4) * 8, _phase_work_empty); + if (elect_sync()) { + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [remAddr32], %2;\n\t" + "}" + :: "r"(work_full_addr + work_stage_4 * 8), "r"(0), "r"((uint32_t)(16)) : "memory"); + asm volatile( + "fence.proxy.async.shared::cta;\n\t" + "clusterlaunchcontrol.try_cancel.async.shared::cta" + ".mbarrier::complete_tx::bytes.b128" + " [%0], [%1];" + :: "r"(work_response_addr + work_stage_4 * 16 + 0 * 16), "r"(work_full_addr + work_stage_4 * 8) + : "memory"); + } + mbarrier_wait(work_full_addr + (work_stage_4) * 8, _phase_work_full_4); + uint32_t _clc_valid_5 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_5) + : "r"(work_response_addr + work_stage_4 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_15 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_15) + : "r"(work_response_addr + work_stage_4 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_16 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_16) + : "r"(work_response_addr + work_stage_4 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_17 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_17) + : "r"(work_response_addr + work_stage_4 * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + work_stage_4 * 8), "r"(0) : "memory"); + work_stage_4 += 1; + if (work_stage_4 == 2) { work_stage_4 = 0; _phase_work_empty ^= 1; _phase_work_full_4 ^= 1; } + if (_clc_valid_5 == 0) { + break; + } + } + } + } + // ---- Role: padding ---- + if (warp == 11) { + { // padding_main + unsigned int work_stage_5 = 0; + unsigned int work_x_4 = blockIdx.x; + unsigned int work_y_4 = blockIdx.y; + unsigned int work_z_4 = blockIdx.z; + unsigned int _phase_work_full_5 = 0; + #pragma unroll 1 + for (unsigned int _work_5 = 0; _work_5 < 120; _work_5++) { + mbarrier_wait(work_full_addr + (work_stage_5) * 8, _phase_work_full_5); + uint32_t _clc_valid_6 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_6) + : "r"(work_response_addr + work_stage_5 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_18 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_18) + : "r"(work_response_addr + work_stage_5 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_19 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_19) + : "r"(work_response_addr + work_stage_5 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_20 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_20) + : "r"(work_response_addr + work_stage_5 * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + work_stage_5 * 8), "r"(0) : "memory"); + work_stage_5 += 1; + if (work_stage_5 == 2) { work_stage_5 = 0; _phase_work_full_5 ^= 1; } + if (_clc_valid_6 == 0) { + break; + } + work_x_4 = _clc_ctaid_18; + work_y_4 = _clc_ctaid_19; + work_z_4 = _clc_ctaid_20; + } + } + } + // ---- Role: loader ---- + if (warp >= 12 && warp <= 19) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 72;"); + { // loader_main + asm volatile("griddepcontrol.wait;" ::: "memory"); + const int load_rank = warp - 12; + unsigned int work_stage_6 = 0; + unsigned int throttle_stage_1 = 0; + unsigned int throttle_phase_1 = 1; + unsigned int q_stage_1 = 0; + unsigned int q_phase_1 = 1; + unsigned int kv_stage_1 = 0; + unsigned int kv_phase_1 = 1; + unsigned int page_k_stage = 0; + unsigned int page_k_phase = 0; + unsigned int page_v_stage = 1; + unsigned int page_v_phase = 0; + unsigned int work_x_5 = blockIdx.x; + unsigned int work_y_5 = blockIdx.y; + unsigned int work_z_5 = blockIdx.z; + unsigned int _phase_work_full_6 = 0; + #pragma unroll 1 + for (unsigned int _work_6 = 0; _work_6 < 120; _work_6++) { + int q_begin_4 = cum_seq_lens_q[work_z_5]; + int q_end_4 = cum_seq_lens_q[work_z_5 + 1]; + int q_len_4 = q_end_4 - q_begin_4; + int query_idx_4 = (unsigned int)q_begin_4 + work_x_5; + int active_topk_4 = 0; + if (work_x_5 < (unsigned int)q_len_4) { + int _max_2 = ((sparse_topk_lens[query_idx_4] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_4] + sparse_topk_lens_offset) : (0)); + int _min_1 = ((_max_2) < (sparse_topk) ? (_max_2) : (sparse_topk)); + active_topk_4 = _min_1; + } + int num_steps_4 = (active_topk_4 + 128 - 1) / 128; + int head_group_1 = work_y_5 / 4; + int value_quarter_2 = work_y_5 % 4; + mbarrier_wait(throttle_empty_addr + (throttle_stage_1) * 8, throttle_phase_1); + mbarrier_arrive(throttle_full_addr + (throttle_stage_1) * 8); + throttle_stage_1 += 1; + if (throttle_stage_1 == 2) { throttle_stage_1 = 0; throttle_phase_1 ^= 1; } + if (num_steps_4 > 0) { + mbarrier_wait(q_empty_addr + (q_stage_1) * 8, q_phase_1); + if (load_rank == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(q_full_addr + (q_stage_1) * 8, 8192); + #pragma unroll + for (int hs_2 = 0; hs_2 < 4; hs_2++) { + tma_4d_gmem2smem(smem_q_addr + (unsigned int)(hs_2 * 2048), tmap_q, 0, head_group_1 * 8, hs_2 * 2, query_idx_4, q_full_addr + (q_stage_1) * 8); + } + } + } + mbarrier_wait(page_full_addr + (page_k_stage) * 8, page_k_phase); + int page_base = 0; + int warps_tile_idx = 0; + int _min_2 = ((active_topk_4) < (128) ? (active_topk_4) : (128)); + int warps_valid_rows = _min_2; + int _min_3 = (((warps_valid_rows + 3) / 4) < (8) ? ((warps_valid_rows + 3) / 4) : (8)); + int warps_tile = _min_3; + int warps_ge1 = 8; + int warps_ge2 = 8; + if (warps_tile_idx >= 1) { + warps_ge1 = warps_tile; + } + if (warps_tile_idx >= 2) { + warps_ge2 = warps_tile; + } + #pragma unroll + for (int hs_3 = 0; hs_3 < 4; hs_3++) { + int hs_sel_2 = ((hs_3 == 3) ? value_quarter_2 : ((value_quarter_2 > hs_3) ? hs_3 : hs_3 + 1)); + mbarrier_wait(kv_empty_addr + (kv_stage_1) * 8, kv_phase_1); + if (load_rank == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_1) * 8, ((hs_3 >= 2) ? warps_ge1 : warps_ge2) * 4096); + } + } + asm volatile("barrier.sync 9, 256;" ::: "memory"); + if (load_rank < ((hs_3 >= 2) ? warps_ge1 : warps_ge2)) { + #pragma unroll + for (int group = 0; group < 4; group++) { + const int row_base = load_rank * 4 + group * 32; + int rows[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&rows[0])), "=r"(*reinterpret_cast(&rows[(0) + 1])), "=r"(*reinterpret_cast(&rows[(0) + 2])), "=r"(*reinterpret_cast(&rows[(0) + 3])) + : "r"(smem_page_offsets_addr + page_k_stage * 1024 + (unsigned int)((page_base + row_base) * 4))); + int dst_k = smem_kv_addr + kv_stage_1 * 32768 + (unsigned int)(row_base * 64 * 2); + if (elect_sync()) { + { + tma_gather4_gmem2smem(dst_k, tmap_swa_kv, hs_sel_2 * 128, rows[0], rows[1], rows[2], rows[3], kv_full_addr + (kv_stage_1) * 8); + tma_gather4_gmem2smem(dst_k + 16384, tmap_swa_kv, hs_sel_2 * 128 + 64, rows[0], rows[1], rows[2], rows[3], kv_full_addr + (kv_stage_1) * 8); + } + } + } + } + kv_stage_1 += 1; + if (kv_stage_1 == 6) { kv_stage_1 = 0; kv_phase_1 ^= 1; } + } + #pragma unroll 1 + for (int tile_2 = 0; tile_2 < num_steps_4 - 1; tile_2++) { + int next_tile = tile_2 + 1; + if ((next_tile & 1) == 0) { + mbarrier_wait(page_full_addr + (page_k_stage) * 8, page_k_phase); + } + int page_base_0 = (next_tile & 1) * 128; + int warps_tile_idx_1 = next_tile; + int _min_4 = ((active_topk_4 - next_tile * 128) < (128) ? (active_topk_4 - next_tile * 128) : (128)); + int warps_valid_rows_2 = _min_4; + int _min_5 = (((warps_valid_rows_2 + 3) / 4) < (8) ? ((warps_valid_rows_2 + 3) / 4) : (8)); + int warps_tile_3 = _min_5; + int warps_ge1_4 = 8; + int warps_ge2_5 = 8; + if (warps_tile_idx_1 >= 1) { + warps_ge1_4 = warps_tile_3; + } + if (warps_tile_idx_1 >= 2) { + warps_ge2_5 = warps_tile_3; + } + #pragma unroll + for (int hs_4 = 0; hs_4 < 4; hs_4++) { + int hs_sel_3 = ((hs_4 == 3) ? value_quarter_2 : ((value_quarter_2 > hs_4) ? hs_4 : hs_4 + 1)); + mbarrier_wait(kv_empty_addr + (kv_stage_1) * 8, kv_phase_1); + if (load_rank == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_1) * 8, ((hs_4 >= 2) ? warps_ge1_4 : warps_ge2_5) * 4096); + } + } + asm volatile("barrier.sync 9, 256;" ::: "memory"); + if (load_rank < ((hs_4 >= 2) ? warps_ge1_4 : warps_ge2_5)) { + #pragma unroll + for (int group_1 = 0; group_1 < 4; group_1++) { + const int row_base_1 = load_rank * 4 + group_1 * 32; + int rows_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&rows_1[0])), "=r"(*reinterpret_cast(&rows_1[(0) + 1])), "=r"(*reinterpret_cast(&rows_1[(0) + 2])), "=r"(*reinterpret_cast(&rows_1[(0) + 3])) + : "r"(smem_page_offsets_addr + page_k_stage * 1024 + (unsigned int)((page_base_0 + row_base_1) * 4))); + int dst_k_1 = smem_kv_addr + kv_stage_1 * 32768 + (unsigned int)(row_base_1 * 64 * 2); + if (elect_sync()) { + if (next_tile == 0) { + tma_gather4_gmem2smem(dst_k_1, tmap_swa_kv, hs_sel_3 * 128, rows_1[0], rows_1[1], rows_1[2], rows_1[3], kv_full_addr + (kv_stage_1) * 8); + tma_gather4_gmem2smem(dst_k_1 + 16384, tmap_swa_kv, hs_sel_3 * 128 + 64, rows_1[0], rows_1[1], rows_1[2], rows_1[3], kv_full_addr + (kv_stage_1) * 8); + } else { + tma_gather4_gmem2smem(dst_k_1, tmap_compressed_kv, hs_sel_3 * 128, rows_1[0], rows_1[1], rows_1[2], rows_1[3], kv_full_addr + (kv_stage_1) * 8); + tma_gather4_gmem2smem(dst_k_1 + 16384, tmap_compressed_kv, hs_sel_3 * 128 + 64, rows_1[0], rows_1[1], rows_1[2], rows_1[3], kv_full_addr + (kv_stage_1) * 8); + } + } + } + } + kv_stage_1 += 1; + if (kv_stage_1 == 6) { kv_stage_1 = 0; kv_phase_1 ^= 1; } + } + if ((next_tile & 1) != 0 && next_tile < num_steps_4 - 1) { + mbarrier_arrive(page_empty_addr + (page_k_stage) * 8); + page_k_stage += 1; + if (page_k_stage == 6) { page_k_stage = 0; page_k_phase ^= 1; } + page_k_stage += 1; + if (page_k_stage == 6) { page_k_stage = 0; page_k_phase ^= 1; } + } + if ((tile_2 & 1) == 0) { + mbarrier_wait(page_full_addr + (page_v_stage) * 8, page_v_phase); + } + if ((tile_2 & 1) != 0) { + mbarrier_arrive(page_empty_addr + (page_v_stage) * 8); + page_v_stage += 1; + if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } + page_v_stage += 1; + if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } + } + } + int last_tile = num_steps_4 - 1; + mbarrier_arrive(page_empty_addr + (page_k_stage) * 8); + page_k_stage += 1; + if (page_k_stage == 6) { page_k_stage = 0; page_k_phase ^= 1; } + page_k_stage += 1; + if (page_k_stage == 6) { page_k_stage = 0; page_k_phase ^= 1; } + if ((last_tile & 1) == 0) { + mbarrier_wait(page_full_addr + (page_v_stage) * 8, page_v_phase); + } + mbarrier_arrive(page_empty_addr + (page_v_stage) * 8); + page_v_stage += 1; + if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } + page_v_stage += 1; + if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } + } + mbarrier_wait(work_full_addr + (work_stage_6) * 8, _phase_work_full_6); + uint32_t _clc_valid_1 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_1) + : "r"(work_response_addr + work_stage_6 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_3 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_3) + : "r"(work_response_addr + work_stage_6 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_4 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_4) + : "r"(work_response_addr + work_stage_6 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_5 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_5) + : "r"(work_response_addr + work_stage_6 * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + work_stage_6 * 8), "r"(0) : "memory"); + work_stage_6 += 1; + if (work_stage_6 == 2) { work_stage_6 = 0; _phase_work_full_6 ^= 1; } + if (_clc_valid_1 == 0) { + break; + } + work_x_5 = _clc_ctaid_3; + work_y_5 = _clc_ctaid_4; + work_z_5 = _clc_ctaid_5; + } + } + } + + // Cleanup +} + +} // extern "C" diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_family_0d09e8b242706c4e05a9_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_family_0d09e8b242706c4e05a9_binding.cu new file mode 100644 index 00000000000..bc529abed8b --- /dev/null +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_family_0d09e8b242706c4e05a9_binding.cu @@ -0,0 +1,703 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +// Generated by cake.runtime.ffi_dispatch — do not edit. +// tvm-ffi dispatch module for kernel family 'cake_dsv4_family_0d09e8b242706c4e05a9'. +// dispatch-policy-revision: unpinned +#include +#include +#include +#include + +#include "tvm_ffi_utils.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +extern "C" __global__ void kernel_cake_dsv4_230143c89d42c6c89387(__nv_bfloat16* __restrict__ partial_O, float* __restrict__ partial_lse, __nv_bfloat16* __restrict__ O, int num_q_heads, int num_split); +extern "C" __global__ void kernel_cake_dsv4_4ab0487d5fde70d4569e(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, const __grid_constant__ CUtensorMap O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); + +namespace cake_host_shim_4bb3ed291da4a466 { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t CakeDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +namespace variant_cake_dsv4_4ab0487d5fde70d4569e_sm_103a_7e73b6214907 { + +// 4D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + int64_t d3 = t.size(t.ndim() - 3); + TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(d3)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 64u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_swa_k' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_swa_k(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_swa_k' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_swa_k' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_swa_k' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_swa_k"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_swa_k' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_swa_k' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_k'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_swa_k') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_compressed_k' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_compressed_k(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_compressed_k' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_compressed_k' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_compressed_k' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_compressed_k"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_compressed_k' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_compressed_k' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_k'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_k') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_swa_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_swa_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_swa_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_swa_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_swa_v' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_swa_v"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_swa_v' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_swa_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_v'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_swa_v') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_compressed_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_compressed_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_compressed_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_compressed_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_compressed_v' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_compressed_v"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_compressed_v' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_compressed_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_v'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_v') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'O' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_O(const TensorView& t, int64_t tp_num_heads, int64_t tp_num_query_tokens) { + TVM_FFI_CHECK(t.ndim() >= 1, ValueError) + << "TMA source 'O' must have at least 1 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'O' must have unit innermost stride, got " << t.stride(-1); + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(8), (uint64_t)(5), (uint64_t)((tp_num_query_tokens * tp_num_heads))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'O' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1] && 1u <= global_dim[2] && 32u <= global_dim[3], ValueError) + << "TMA box (64, 1, 1, 32) exceeds resolved global dims for 'O'"; + int64_t carrier_stride_0 = 64; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 512; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = 2560; + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 1u, 1u, 32u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'O') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_compressed_k, TensorView arg_tmap_swa_v, TensorView arg_tmap_compressed_v, TensorView arg_O, TensorView arg_partial_lse, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_num_query_tokens, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckContiguous(arg_tmap_q, "tmap_q"); + CheckCudaTensor(arg_tmap_swa_k, "tmap_swa_k"); + CheckDtype(arg_tmap_swa_k, "tmap_swa_k", 4, 16, 1); + CheckCudaTensor(arg_tmap_compressed_k, "tmap_compressed_k"); + CheckDtype(arg_tmap_compressed_k, "tmap_compressed_k", 4, 16, 1); + CheckCudaTensor(arg_tmap_swa_v, "tmap_swa_v"); + CheckDtype(arg_tmap_swa_v, "tmap_swa_v", 4, 16, 1); + CheckCudaTensor(arg_tmap_compressed_v, "tmap_compressed_v"); + CheckDtype(arg_tmap_compressed_v, "tmap_compressed_v", 4, 16, 1); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_partial_lse, "partial_lse"); + CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); + CheckContiguous(arg_partial_lse, "partial_lse"); + CheckCudaTensor(arg_swa_indices, "swa_indices"); + CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); + CheckContiguous(arg_swa_indices, "swa_indices"); + CheckCudaTensor(arg_compressed_indices, "compressed_indices"); + CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); + CheckContiguous(arg_compressed_indices, "compressed_indices"); + CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); + CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); + CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); + CheckContiguous(arg_bmm1_scale, "bmm1_scale"); + CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); + CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); + CheckContiguous(arg_bmm2_scale, "bmm2_scale"); + TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) + << "scalar 'num_heads' value " << arg_num_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_query_tokens >= -2147483648LL && arg_num_query_tokens <= 2147483647LL, ValueError) + << "scalar 'num_query_tokens' value " << arg_num_query_tokens + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) + << "scalar 'swa_index_stride' value " << arg_swa_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) + << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) + << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) + << "scalar 'sparse_topk' value " << arg_sparse_topk + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) + << "scalar 'has_sinks' value " << arg_has_sinks + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_swa_k, arg_tmap_q, "tmap_swa_k", "tmap_q"); + CheckSameCudaDevice(arg_tmap_compressed_k, arg_tmap_q, "tmap_compressed_k", "tmap_q"); + CheckSameCudaDevice(arg_tmap_swa_v, arg_tmap_q, "tmap_swa_v", "tmap_q"); + CheckSameCudaDevice(arg_tmap_compressed_v, arg_tmap_q, "tmap_compressed_v", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_partial_lse, arg_tmap_q, "partial_lse", "tmap_q"); + CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); + CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); + CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); + CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") must be divisible by cluster dims (2, 1, 1)"; + + + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_swa_k = EncodeTma_tmap_swa_k(arg_tmap_swa_k); + CUtensorMap p_tmap_compressed_k = EncodeTma_tmap_compressed_k(arg_tmap_compressed_k); + CUtensorMap p_tmap_swa_v = EncodeTma_tmap_swa_v(arg_tmap_swa_v); + CUtensorMap p_tmap_compressed_v = EncodeTma_tmap_compressed_v(arg_tmap_compressed_v); + CUtensorMap p_O = EncodeTma_O(arg_O, (int64_t)arg_num_heads, (int64_t)arg_num_query_tokens); + float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); + int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); + int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); + int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); + float* p_sinks = static_cast(arg_sinks.data_ptr()); + float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); + float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); + int32_t v_num_heads = (int32_t)arg_num_heads; + int32_t v_num_query_tokens = (int32_t)arg_num_query_tokens; + int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; + int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; + int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; + int32_t v_sparse_topk = (int32_t)arg_sparse_topk; + int32_t v_has_sinks = (int32_t)arg_has_sinks; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + void* kargs[] = {&p_tmap_q, &p_tmap_swa_k, &p_tmap_compressed_k, &p_tmap_swa_v, &p_tmap_compressed_v, &p_O, &p_partial_lse, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_num_query_tokens, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_has_sinks, &v_total_work_items}; + + static std::mutex smem_mu; + static auto* smem_status_by_device = + new std::unordered_map(); + cudaError_t smem_status = cudaSuccess; + { + std::lock_guard lock(smem_mu); + auto it = smem_status_by_device->find(dev.device_id); + if (it == smem_status_by_device->end()) { + int current_device = -1; + smem_status = cudaGetDevice(¤t_device); + if (smem_status == cudaSuccess && current_device != dev.device_id) { + smem_status = cudaErrorInvalidDevice; + } + if (smem_status == cudaSuccess) { + smem_status = cudaFuncSetAttribute( + reinterpret_cast(kernel_cake_dsv4_4ab0487d5fde70d4569e), + cudaFuncAttributeMaxDynamicSharedMemorySize, + 231936); + } + smem_status_by_device->emplace(dev.device_id, smem_status); + } else { + smem_status = it->second; + } + } + TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) + << "cudaFuncSetAttribute for kernel_cake_dsv4_4ab0487d5fde70d4569e failed: " + << cudaGetErrorString(smem_status); + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(512u, 1u, 1u); + + // Extended direct-source launch for a linked CUDA kernel symbol. + cudaLaunchAttribute attrs[1]{}; + int n = 0; + attrs[n].id = cudaLaunchAttributeClusterDimension; + attrs[n].val.clusterDim.x = 2u; + attrs[n].val.clusterDim.y = 1u; + attrs[n].val.clusterDim.z = 1u; + ++n; + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes = 231936u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; + cudaError_t launch_status = cudaLaunchKernelExC( + &config, reinterpret_cast(kernel_cake_dsv4_4ab0487d5fde70d4569e), kargs); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernelExC for kernel_cake_dsv4_4ab0487d5fde70d4569e failed: " + << cudaGetErrorString(launch_status); +} +} // namespace variant_cake_dsv4_4ab0487d5fde70d4569e_sm_103a_7e73b6214907 + +namespace variant_cake_dsv4_230143c89d42c6c89387_sm_103a_c6d3bf2db9b6 { + +void Run(TensorView arg_partial_O, TensorView arg_partial_lse, TensorView arg_O, int64_t arg_num_q_heads, int64_t arg_num_split, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_partial_O.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_partial_O, "partial_O"); + CheckDtype(arg_partial_O, "partial_O", 4, 16, 1); + CheckContiguous(arg_partial_O, "partial_O"); + CheckCudaTensor(arg_partial_lse, "partial_lse"); + CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); + CheckContiguous(arg_partial_lse, "partial_lse"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + TVM_FFI_CHECK(arg_num_q_heads >= -2147483648LL && arg_num_q_heads <= 2147483647LL, ValueError) + << "scalar 'num_q_heads' value " << arg_num_q_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_split >= -2147483648LL && arg_num_split <= 2147483647LL, ValueError) + << "scalar 'num_split' value " << arg_num_split + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_partial_lse, arg_partial_O, "partial_lse", "partial_O"); + CheckSameCudaDevice(arg_O, arg_partial_O, "O", "partial_O"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + + + __nv_bfloat16* p_partial_O = static_cast<__nv_bfloat16*>(arg_partial_O.data_ptr()); + float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); + __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); + int32_t v_num_q_heads = (int32_t)arg_num_q_heads; + int32_t v_num_split = (int32_t)arg_num_split; + void* kargs[] = {&p_partial_O, &p_partial_lse, &p_O, &v_num_q_heads, &v_num_split}; + + + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(128u, 1u, 1u); + + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes = 0u; + config.stream = stream; + cudaError_t launch_status = cudaLaunchKernelExC( + &config, reinterpret_cast(kernel_cake_dsv4_230143c89d42c6c89387), kargs); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernelExC for kernel_cake_dsv4_230143c89d42c6c89387 failed: " + << cudaGetErrorString(launch_status); +} +} // namespace variant_cake_dsv4_230143c89d42c6c89387_sm_103a_c6d3bf2db9b6 + +int SelectVariant(int32_t num_heads, int32_t num_query_tokens, int32_t swa_index_stride, int32_t compressed_index_stride, int32_t sparse_topk_lens_offset, int32_t sparse_topk, int32_t has_sinks, int32_t total_work_items, int32_t num_split) { + if ((num_split == 5)) return 0; + return -1; +} + +void Dispatch(tvm::ffi::TensorView arg_Q, tvm::ffi::TensorView arg_SWA_cache, tvm::ffi::TensorView arg_compressed_KV_cache, tvm::ffi::TensorView arg_partial_O, tvm::ffi::TensorView arg_partial_lse, tvm::ffi::TensorView arg_swa_indices, tvm::ffi::TensorView arg_compressed_indices, tvm::ffi::TensorView arg_sparse_topk_lens, tvm::ffi::TensorView arg_sinks, tvm::ffi::TensorView arg_bmm1_scale, tvm::ffi::TensorView arg_bmm2_scale, tvm::ffi::TensorView arg_O, int32_t num_heads, int32_t num_query_tokens, int32_t swa_index_stride, int32_t compressed_index_stride, int32_t sparse_topk_lens_offset, int32_t sparse_topk, int32_t has_sinks, int32_t total_work_items, int32_t num_split, int64_t cuda_stream_ptr) { + CheckCudaTensor(arg_Q, "Q"); + CheckCudaTensor(arg_SWA_cache, "SWA_cache"); + CheckCudaTensor(arg_compressed_KV_cache, "compressed_KV_cache"); + CheckCudaTensor(arg_partial_O, "partial_O"); + CheckCudaTensor(arg_partial_lse, "partial_lse"); + CheckCudaTensor(arg_swa_indices, "swa_indices"); + CheckCudaTensor(arg_compressed_indices, "compressed_indices"); + CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); + CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); + CheckCudaTensor(arg_O, "O"); + cudaStream_t stream = reinterpret_cast(static_cast(cuda_stream_ptr)); + int v = SelectVariant(num_heads, num_query_tokens, swa_index_stride, compressed_index_stride, sparse_topk_lens_offset, sparse_topk, has_sinks, total_work_items, num_split); + switch (v) { + case 0: { + { + int64_t gx = (total_work_items * 2); + int64_t gy = 1; + int64_t gz = 1; + cake_host_shim_4bb3ed291da4a466::variant_cake_dsv4_4ab0487d5fde70d4569e_sm_103a_7e73b6214907::Run(arg_Q, arg_SWA_cache, arg_compressed_KV_cache, arg_SWA_cache, arg_compressed_KV_cache, arg_partial_O, arg_partial_lse, arg_swa_indices, arg_compressed_indices, arg_sparse_topk_lens, arg_sinks, arg_bmm1_scale, arg_bmm2_scale, num_heads, num_query_tokens, swa_index_stride, compressed_index_stride, sparse_topk_lens_offset, sparse_topk, has_sinks, total_work_items, gx, gy, gz, stream); + } + { + int64_t gx = num_query_tokens; + int64_t gy = (num_heads / 4); + int64_t gz = 1; + cake_host_shim_4bb3ed291da4a466::variant_cake_dsv4_230143c89d42c6c89387_sm_103a_c6d3bf2db9b6::Run(arg_partial_O, arg_partial_lse, arg_O, num_heads, num_split, gx, gy, gz, stream); + } + return; + } + default: + TVM_FFI_CHECK(false, ValueError) << "no dispatch route for cake_dsv4_family_0d09e8b242706c4e05a9 shape (num_heads=" << num_heads << ", num_query_tokens=" << num_query_tokens << ", swa_index_stride=" << swa_index_stride << ", compressed_index_stride=" << compressed_index_stride << ", sparse_topk_lens_offset=" << sparse_topk_lens_offset << ", sparse_topk=" << sparse_topk << ", has_sinks=" << has_sinks << ", total_work_items=" << total_work_items << ", num_split=" << num_split << ")"; + } +} + +int64_t RouteOf(int32_t num_heads, int32_t num_query_tokens, int32_t swa_index_stride, int32_t compressed_index_stride, int32_t sparse_topk_lens_offset, int32_t sparse_topk, int32_t has_sinks, int32_t total_work_items, int32_t num_split) { return SelectVariant(num_heads, num_query_tokens, swa_index_stride, compressed_index_stride, sparse_topk_lens_offset, sparse_topk, has_sinks, total_work_items, num_split); } + +} // namespace cake_host_shim_4bb3ed291da4a466 + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_4bb3ed291da4a466::Dispatch); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(route_of, cake_host_shim_4bb3ed291da4a466::RouteOf); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_family_ebb85adb872529ca9e83_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_family_ebb85adb872529ca9e83_binding.cu new file mode 100644 index 00000000000..aa2fd81aeb1 --- /dev/null +++ b/csrc/cake_dsv4/sm_103a/cake_dsv4_family_ebb85adb872529ca9e83_binding.cu @@ -0,0 +1,703 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +// Generated by cake.runtime.ffi_dispatch — do not edit. +// tvm-ffi dispatch module for kernel family 'cake_dsv4_family_ebb85adb872529ca9e83'. +// dispatch-policy-revision: unpinned +#include +#include +#include +#include + +#include "tvm_ffi_utils.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +extern "C" __global__ void kernel_cake_dsv4_397f817e2931a6cc7c92(__nv_bfloat16* __restrict__ partial_O, float* __restrict__ partial_lse, __nv_bfloat16* __restrict__ O, int num_q_heads, int num_split); +extern "C" __global__ void kernel_cake_dsv4_064850af5858dd03c104(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, const __grid_constant__ CUtensorMap O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); + +namespace cake_host_shim_71d3307de7b4df2d { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t CakeDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +namespace variant_cake_dsv4_064850af5858dd03c104_sm_103a_6276393afd19 { + +// 4D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + int64_t d3 = t.size(t.ndim() - 3); + TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(d3)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 64u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_swa_k' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_swa_k(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_swa_k' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_swa_k' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_swa_k' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_swa_k"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_swa_k' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_swa_k' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_k'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_swa_k') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_compressed_k' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_compressed_k(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_compressed_k' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_compressed_k' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_compressed_k' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_compressed_k"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_compressed_k' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_compressed_k' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_k'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_k') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_swa_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_swa_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_swa_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_swa_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_swa_v' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_swa_v"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_swa_v' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_swa_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_v'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_swa_v') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_compressed_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_compressed_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_compressed_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_compressed_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_compressed_v' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_compressed_v"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_compressed_v' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_compressed_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_v'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_v') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'O' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_O(const TensorView& t, int64_t tp_num_heads, int64_t tp_num_query_tokens) { + TVM_FFI_CHECK(t.ndim() >= 1, ValueError) + << "TMA source 'O' must have at least 1 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'O' must have unit innermost stride, got " << t.stride(-1); + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(8), (uint64_t)(4), (uint64_t)((tp_num_query_tokens * tp_num_heads))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'O' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1] && 1u <= global_dim[2] && 32u <= global_dim[3], ValueError) + << "TMA box (64, 1, 1, 32) exceeds resolved global dims for 'O'"; + int64_t carrier_stride_0 = 64; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 512; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = 2048; + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 1u, 1u, 32u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'O') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_compressed_k, TensorView arg_tmap_swa_v, TensorView arg_tmap_compressed_v, TensorView arg_O, TensorView arg_partial_lse, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_num_query_tokens, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckContiguous(arg_tmap_q, "tmap_q"); + CheckCudaTensor(arg_tmap_swa_k, "tmap_swa_k"); + CheckDtype(arg_tmap_swa_k, "tmap_swa_k", 4, 16, 1); + CheckCudaTensor(arg_tmap_compressed_k, "tmap_compressed_k"); + CheckDtype(arg_tmap_compressed_k, "tmap_compressed_k", 4, 16, 1); + CheckCudaTensor(arg_tmap_swa_v, "tmap_swa_v"); + CheckDtype(arg_tmap_swa_v, "tmap_swa_v", 4, 16, 1); + CheckCudaTensor(arg_tmap_compressed_v, "tmap_compressed_v"); + CheckDtype(arg_tmap_compressed_v, "tmap_compressed_v", 4, 16, 1); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_partial_lse, "partial_lse"); + CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); + CheckContiguous(arg_partial_lse, "partial_lse"); + CheckCudaTensor(arg_swa_indices, "swa_indices"); + CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); + CheckContiguous(arg_swa_indices, "swa_indices"); + CheckCudaTensor(arg_compressed_indices, "compressed_indices"); + CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); + CheckContiguous(arg_compressed_indices, "compressed_indices"); + CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); + CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); + CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); + CheckContiguous(arg_bmm1_scale, "bmm1_scale"); + CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); + CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); + CheckContiguous(arg_bmm2_scale, "bmm2_scale"); + TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) + << "scalar 'num_heads' value " << arg_num_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_query_tokens >= -2147483648LL && arg_num_query_tokens <= 2147483647LL, ValueError) + << "scalar 'num_query_tokens' value " << arg_num_query_tokens + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) + << "scalar 'swa_index_stride' value " << arg_swa_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) + << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) + << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) + << "scalar 'sparse_topk' value " << arg_sparse_topk + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) + << "scalar 'has_sinks' value " << arg_has_sinks + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_swa_k, arg_tmap_q, "tmap_swa_k", "tmap_q"); + CheckSameCudaDevice(arg_tmap_compressed_k, arg_tmap_q, "tmap_compressed_k", "tmap_q"); + CheckSameCudaDevice(arg_tmap_swa_v, arg_tmap_q, "tmap_swa_v", "tmap_q"); + CheckSameCudaDevice(arg_tmap_compressed_v, arg_tmap_q, "tmap_compressed_v", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_partial_lse, arg_tmap_q, "partial_lse", "tmap_q"); + CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); + CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); + CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); + CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") must be divisible by cluster dims (2, 1, 1)"; + + + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_swa_k = EncodeTma_tmap_swa_k(arg_tmap_swa_k); + CUtensorMap p_tmap_compressed_k = EncodeTma_tmap_compressed_k(arg_tmap_compressed_k); + CUtensorMap p_tmap_swa_v = EncodeTma_tmap_swa_v(arg_tmap_swa_v); + CUtensorMap p_tmap_compressed_v = EncodeTma_tmap_compressed_v(arg_tmap_compressed_v); + CUtensorMap p_O = EncodeTma_O(arg_O, (int64_t)arg_num_heads, (int64_t)arg_num_query_tokens); + float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); + int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); + int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); + int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); + float* p_sinks = static_cast(arg_sinks.data_ptr()); + float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); + float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); + int32_t v_num_heads = (int32_t)arg_num_heads; + int32_t v_num_query_tokens = (int32_t)arg_num_query_tokens; + int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; + int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; + int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; + int32_t v_sparse_topk = (int32_t)arg_sparse_topk; + int32_t v_has_sinks = (int32_t)arg_has_sinks; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + void* kargs[] = {&p_tmap_q, &p_tmap_swa_k, &p_tmap_compressed_k, &p_tmap_swa_v, &p_tmap_compressed_v, &p_O, &p_partial_lse, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_num_query_tokens, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_has_sinks, &v_total_work_items}; + + static std::mutex smem_mu; + static auto* smem_status_by_device = + new std::unordered_map(); + cudaError_t smem_status = cudaSuccess; + { + std::lock_guard lock(smem_mu); + auto it = smem_status_by_device->find(dev.device_id); + if (it == smem_status_by_device->end()) { + int current_device = -1; + smem_status = cudaGetDevice(¤t_device); + if (smem_status == cudaSuccess && current_device != dev.device_id) { + smem_status = cudaErrorInvalidDevice; + } + if (smem_status == cudaSuccess) { + smem_status = cudaFuncSetAttribute( + reinterpret_cast(kernel_cake_dsv4_064850af5858dd03c104), + cudaFuncAttributeMaxDynamicSharedMemorySize, + 220032); + } + smem_status_by_device->emplace(dev.device_id, smem_status); + } else { + smem_status = it->second; + } + } + TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) + << "cudaFuncSetAttribute for kernel_cake_dsv4_064850af5858dd03c104 failed: " + << cudaGetErrorString(smem_status); + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(512u, 1u, 1u); + + // Extended direct-source launch for a linked CUDA kernel symbol. + cudaLaunchAttribute attrs[1]{}; + int n = 0; + attrs[n].id = cudaLaunchAttributeClusterDimension; + attrs[n].val.clusterDim.x = 2u; + attrs[n].val.clusterDim.y = 1u; + attrs[n].val.clusterDim.z = 1u; + ++n; + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes = 220032u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; + cudaError_t launch_status = cudaLaunchKernelExC( + &config, reinterpret_cast(kernel_cake_dsv4_064850af5858dd03c104), kargs); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernelExC for kernel_cake_dsv4_064850af5858dd03c104 failed: " + << cudaGetErrorString(launch_status); +} +} // namespace variant_cake_dsv4_064850af5858dd03c104_sm_103a_6276393afd19 + +namespace variant_cake_dsv4_397f817e2931a6cc7c92_sm_103a_aa3105f3b9de { + +void Run(TensorView arg_partial_O, TensorView arg_partial_lse, TensorView arg_O, int64_t arg_num_q_heads, int64_t arg_num_split, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_partial_O.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_partial_O, "partial_O"); + CheckDtype(arg_partial_O, "partial_O", 4, 16, 1); + CheckContiguous(arg_partial_O, "partial_O"); + CheckCudaTensor(arg_partial_lse, "partial_lse"); + CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); + CheckContiguous(arg_partial_lse, "partial_lse"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + TVM_FFI_CHECK(arg_num_q_heads >= -2147483648LL && arg_num_q_heads <= 2147483647LL, ValueError) + << "scalar 'num_q_heads' value " << arg_num_q_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_split >= -2147483648LL && arg_num_split <= 2147483647LL, ValueError) + << "scalar 'num_split' value " << arg_num_split + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_partial_lse, arg_partial_O, "partial_lse", "partial_O"); + CheckSameCudaDevice(arg_O, arg_partial_O, "O", "partial_O"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + + + __nv_bfloat16* p_partial_O = static_cast<__nv_bfloat16*>(arg_partial_O.data_ptr()); + float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); + __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); + int32_t v_num_q_heads = (int32_t)arg_num_q_heads; + int32_t v_num_split = (int32_t)arg_num_split; + void* kargs[] = {&p_partial_O, &p_partial_lse, &p_O, &v_num_q_heads, &v_num_split}; + + + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(128u, 1u, 1u); + + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes = 0u; + config.stream = stream; + cudaError_t launch_status = cudaLaunchKernelExC( + &config, reinterpret_cast(kernel_cake_dsv4_397f817e2931a6cc7c92), kargs); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernelExC for kernel_cake_dsv4_397f817e2931a6cc7c92 failed: " + << cudaGetErrorString(launch_status); +} +} // namespace variant_cake_dsv4_397f817e2931a6cc7c92_sm_103a_aa3105f3b9de + +int SelectVariant(int32_t num_heads, int32_t num_query_tokens, int32_t swa_index_stride, int32_t compressed_index_stride, int32_t sparse_topk_lens_offset, int32_t sparse_topk, int32_t has_sinks, int32_t total_work_items, int32_t num_split) { + if ((num_split == 4)) return 0; + return -1; +} + +void Dispatch(tvm::ffi::TensorView arg_Q, tvm::ffi::TensorView arg_SWA_cache, tvm::ffi::TensorView arg_compressed_KV_cache, tvm::ffi::TensorView arg_partial_O, tvm::ffi::TensorView arg_partial_lse, tvm::ffi::TensorView arg_swa_indices, tvm::ffi::TensorView arg_compressed_indices, tvm::ffi::TensorView arg_sparse_topk_lens, tvm::ffi::TensorView arg_sinks, tvm::ffi::TensorView arg_bmm1_scale, tvm::ffi::TensorView arg_bmm2_scale, tvm::ffi::TensorView arg_O, int32_t num_heads, int32_t num_query_tokens, int32_t swa_index_stride, int32_t compressed_index_stride, int32_t sparse_topk_lens_offset, int32_t sparse_topk, int32_t has_sinks, int32_t total_work_items, int32_t num_split, int64_t cuda_stream_ptr) { + CheckCudaTensor(arg_Q, "Q"); + CheckCudaTensor(arg_SWA_cache, "SWA_cache"); + CheckCudaTensor(arg_compressed_KV_cache, "compressed_KV_cache"); + CheckCudaTensor(arg_partial_O, "partial_O"); + CheckCudaTensor(arg_partial_lse, "partial_lse"); + CheckCudaTensor(arg_swa_indices, "swa_indices"); + CheckCudaTensor(arg_compressed_indices, "compressed_indices"); + CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); + CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); + CheckCudaTensor(arg_O, "O"); + cudaStream_t stream = reinterpret_cast(static_cast(cuda_stream_ptr)); + int v = SelectVariant(num_heads, num_query_tokens, swa_index_stride, compressed_index_stride, sparse_topk_lens_offset, sparse_topk, has_sinks, total_work_items, num_split); + switch (v) { + case 0: { + { + int64_t gx = (total_work_items * 2); + int64_t gy = 1; + int64_t gz = 1; + cake_host_shim_71d3307de7b4df2d::variant_cake_dsv4_064850af5858dd03c104_sm_103a_6276393afd19::Run(arg_Q, arg_SWA_cache, arg_compressed_KV_cache, arg_SWA_cache, arg_compressed_KV_cache, arg_partial_O, arg_partial_lse, arg_swa_indices, arg_compressed_indices, arg_sparse_topk_lens, arg_sinks, arg_bmm1_scale, arg_bmm2_scale, num_heads, num_query_tokens, swa_index_stride, compressed_index_stride, sparse_topk_lens_offset, sparse_topk, has_sinks, total_work_items, gx, gy, gz, stream); + } + { + int64_t gx = num_query_tokens; + int64_t gy = num_heads; + int64_t gz = 1; + cake_host_shim_71d3307de7b4df2d::variant_cake_dsv4_397f817e2931a6cc7c92_sm_103a_aa3105f3b9de::Run(arg_partial_O, arg_partial_lse, arg_O, num_heads, num_split, gx, gy, gz, stream); + } + return; + } + default: + TVM_FFI_CHECK(false, ValueError) << "no dispatch route for cake_dsv4_family_ebb85adb872529ca9e83 shape (num_heads=" << num_heads << ", num_query_tokens=" << num_query_tokens << ", swa_index_stride=" << swa_index_stride << ", compressed_index_stride=" << compressed_index_stride << ", sparse_topk_lens_offset=" << sparse_topk_lens_offset << ", sparse_topk=" << sparse_topk << ", has_sinks=" << has_sinks << ", total_work_items=" << total_work_items << ", num_split=" << num_split << ")"; + } +} + +int64_t RouteOf(int32_t num_heads, int32_t num_query_tokens, int32_t swa_index_stride, int32_t compressed_index_stride, int32_t sparse_topk_lens_offset, int32_t sparse_topk, int32_t has_sinks, int32_t total_work_items, int32_t num_split) { return SelectVariant(num_heads, num_query_tokens, swa_index_stride, compressed_index_stride, sparse_topk_lens_offset, sparse_topk, has_sinks, total_work_items, num_split); } + +} // namespace cake_host_shim_71d3307de7b4df2d + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_71d3307de7b4df2d::Dispatch); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(route_of, cake_host_shim_71d3307de7b4df2d::RouteOf); diff --git a/flashinfer/jit/cake_dsv4.py b/flashinfer/jit/cake_dsv4.py index 4381bdac79c..6c1525f8151 100644 --- a/flashinfer/jit/cake_dsv4.py +++ b/flashinfer/jit/cake_dsv4.py @@ -1237,21 +1237,33 @@ }, "sm_103a": { "programs": { - "cake_dsv4_family_1b5ee6b39dec3c3538c6": { + "cake_dsv4_family_0d09e8b242706c4e05a9": { "arch": "sm_103a", "compile_flags": ["-std=c++17"], - "dependencies": ["bf16_h128_prefill_v42"], + "dependencies": [ + "bf16_h128_topk4x_v52", + "split_reduce_cake_dsv4_230143c89d42c6c89387", + ], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "860901dbab47cf43c94ed40b20ba0847a3137cf7deb3f0a88e699c67963ca890", + "identity": "897382fde021af3afe3cbf39d8771c6f45a849efd11f723e57ca4fe800f91f6a", "route_entry": "route_of", - "routes": [{"index": 0, "stage_variants": ["bf16_h128_prefill_v42"]}], + "routes": [ + { + "index": 0, + "stage_variants": [ + "bf16_h128_topk4x_v52", + "split_reduce_cake_dsv4_230143c89d42c6c89387", + ], + } + ], "signature": { "scalar_dtypes": { "compressed_index_stride": "int32", "has_sinks": "int32", "num_heads": "int32", "num_query_tokens": "int32", + "num_split": "int32", "sparse_topk": "int32", "sparse_topk_lens_offset": "int32", "swa_index_stride": "int32", @@ -1266,6 +1278,7 @@ "sparse_topk", "has_sinks", "total_work_items", + "num_split", ], "tensor_dtypes": { "O": "bfloat16", @@ -1275,6 +1288,7 @@ "bmm2_scale": "float32", "compressed_KV_cache": "bfloat16", "compressed_indices": "int32", + "partial_O": "bfloat16", "partial_lse": "float32", "sinks": "float32", "sparse_topk_lens": "int32", @@ -1284,6 +1298,7 @@ "Q", "SWA_cache", "compressed_KV_cache", + "partial_O", "partial_lse", "swa_indices", "compressed_indices", @@ -1296,19 +1311,17 @@ "workspace_dtypes": {}, "workspace_keys": [], }, - "sources": ["cake_dsv4_family_1b5ee6b39dec3c3538c6_binding.cu"], + "sources": ["cake_dsv4_family_0d09e8b242706c4e05a9_binding.cu"], }, - "cake_dsv4_family_42ae3e93a01e7a3bf20d": { + "cake_dsv4_family_1b5ee6b39dec3c3538c6": { "arch": "sm_103a", "compile_flags": ["-std=c++17"], - "dependencies": ["bf16_h128_prefill_v42_snake"], + "dependencies": ["bf16_h128_prefill_v42"], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "285cf074aae8f3086ca239ade301c7eb313c22c15286c2f26f0386be98ebb4f5", + "identity": "860901dbab47cf43c94ed40b20ba0847a3137cf7deb3f0a88e699c67963ca890", "route_entry": "route_of", - "routes": [ - {"index": 0, "stage_variants": ["bf16_h128_prefill_v42_snake"]} - ], + "routes": [{"index": 0, "stage_variants": ["bf16_h128_prefill_v42"]}], "signature": { "scalar_dtypes": { "compressed_index_stride": "int32", @@ -1359,18 +1372,18 @@ "workspace_dtypes": {}, "workspace_keys": [], }, - "sources": ["cake_dsv4_family_42ae3e93a01e7a3bf20d_binding.cu"], + "sources": ["cake_dsv4_family_1b5ee6b39dec3c3538c6_binding.cu"], }, - "cake_dsv4_family_585100536b80a2cec8ba": { + "cake_dsv4_family_42ae3e93a01e7a3bf20d": { "arch": "sm_103a", "compile_flags": ["-std=c++17"], - "dependencies": ["bf16_h128_topk128x_row_first"], + "dependencies": ["bf16_h128_prefill_v42_snake"], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "72b48d77a88ee5d7e37313209035e21d2037d2b74d4793d7660367a6c9003750", + "identity": "285cf074aae8f3086ca239ade301c7eb313c22c15286c2f26f0386be98ebb4f5", "route_entry": "route_of", "routes": [ - {"index": 0, "stage_variants": ["bf16_h128_topk128x_row_first"]} + {"index": 0, "stage_variants": ["bf16_h128_prefill_v42_snake"]} ], "signature": { "scalar_dtypes": { @@ -1422,27 +1435,18 @@ "workspace_dtypes": {}, "workspace_keys": [], }, - "sources": ["cake_dsv4_family_585100536b80a2cec8ba_binding.cu"], + "sources": ["cake_dsv4_family_42ae3e93a01e7a3bf20d_binding.cu"], }, - "cake_dsv4_family_956aa0f74f4a82042f2c": { + "cake_dsv4_family_585100536b80a2cec8ba": { "arch": "sm_103a", "compile_flags": ["-std=c++17"], - "dependencies": [ - "bf16_h128_topk4x_v52", - "split_reduce_cake_dsv4_230143c89d42c6c89387", - ], + "dependencies": ["bf16_h128_topk128x_row_first"], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "53306debf3bddc3348c65222e0ed7b69fc316a69af2ce6d7314061f98104b233", + "identity": "72b48d77a88ee5d7e37313209035e21d2037d2b74d4793d7660367a6c9003750", "route_entry": "route_of", "routes": [ - { - "index": 0, - "stage_variants": [ - "bf16_h128_topk4x_v52", - "split_reduce_cake_dsv4_230143c89d42c6c89387", - ], - } + {"index": 0, "stage_variants": ["bf16_h128_topk128x_row_first"]} ], "signature": { "scalar_dtypes": { @@ -1450,7 +1454,6 @@ "has_sinks": "int32", "num_heads": "int32", "num_query_tokens": "int32", - "num_split": "int32", "sparse_topk": "int32", "sparse_topk_lens_offset": "int32", "swa_index_stride": "int32", @@ -1465,7 +1468,6 @@ "sparse_topk", "has_sinks", "total_work_items", - "num_split", ], "tensor_dtypes": { "O": "bfloat16", @@ -1475,7 +1477,6 @@ "bmm2_scale": "float32", "compressed_KV_cache": "bfloat16", "compressed_indices": "int32", - "partial_O": "bfloat16", "partial_lse": "float32", "sinks": "float32", "sparse_topk_lens": "int32", @@ -1485,7 +1486,6 @@ "Q", "SWA_cache", "compressed_KV_cache", - "partial_O", "partial_lse", "swa_indices", "compressed_indices", @@ -1498,9 +1498,9 @@ "workspace_dtypes": {}, "workspace_keys": [], }, - "sources": ["cake_dsv4_family_956aa0f74f4a82042f2c_binding.cu"], + "sources": ["cake_dsv4_family_585100536b80a2cec8ba_binding.cu"], }, - "cake_dsv4_family_ec44ddd94e3be88cfc83": { + "cake_dsv4_family_ebb85adb872529ca9e83": { "arch": "sm_103a", "compile_flags": ["-std=c++17"], "dependencies": [ @@ -1509,7 +1509,7 @@ ], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "7c7b9440d3cf7f9ad7a4726319dfcba580a58bf74449e03f048096044be3bee5", + "identity": "42d62a52aae085402e0ffb43549c81c591c5c7bd90b08b4d71ecdf5e0c44fb1f", "route_entry": "route_of", "routes": [ { @@ -1574,7 +1574,7 @@ "workspace_dtypes": {}, "workspace_keys": [], }, - "sources": ["cake_dsv4_family_ec44ddd94e3be88cfc83_binding.cu"], + "sources": ["cake_dsv4_family_ebb85adb872529ca9e83_binding.cu"], }, }, "variants": { @@ -1609,7 +1609,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "3c18335a89ae686e86a9d8f7e02bedfceda722cf8eca63b648d0bf1871fe4cf5", + "identity": "89495d4d9f2fe8ee1b03873aa440cff30e4d413047e7aae57b372c68b430df7f", "sources": [ "cake_dsv4_a40bd49fee9655d8b1c6_kernel.cu", "cake_dsv4_a40bd49fee9655d8b1c6_binding.cu", @@ -1647,7 +1647,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "4ca772248ca31846c587fafbac94c2da2efdf58a844e78832268403caca66abb", + "identity": "525777264ca547d1e7d9e732ff609890eda28ca396da31184dce78c7c23253bb", "sources": [ "cake_dsv4_d338c32a69477808a1ef_kernel.cu", "cake_dsv4_d338c32a69477808a1ef_binding.cu", @@ -1680,10 +1680,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "55e52ad1bad7c5f6d9f0287db61de94bd0c95041a6cb2933c495784d302f9d68", + "identity": "669f6c1916289193a683ac74b7d3c5967c0a8a5ea1e2e652f0f99a9bdcbb40cd", "sources": [ - "cake_dsv4_f89bc6561b38462e9538_kernel.cu", - "cake_dsv4_f89bc6561b38462e9538_binding.cu", + "cake_dsv4_1bad879f597b2e2adf3e_kernel.cu", + "cake_dsv4_1bad879f597b2e2adf3e_binding.cu", ], "tma_workspace_bytes": 0, }, @@ -1718,7 +1718,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "8fdcc0ef01be0f9758a29d5556b2f7bae6864aa2ed8a79e56200ea722bb60c78", + "identity": "f4a7ad606f7e1cbc294f4021c2f443da82e6735e6118ec12af544bb948f798bb", "sources": [ "cake_dsv4_af6f9d240794a942a124_kernel.cu", "cake_dsv4_af6f9d240794a942a124_binding.cu", @@ -1733,7 +1733,7 @@ ["tma_buffer", "tmap_compressed_k"], ["tma_buffer", "tmap_swa_v"], ["tma_buffer", "tmap_compressed_v"], - ["buffer", "O"], + ["tma_buffer", "O"], ["buffer", "partial_lse"], ["buffer", "swa_indices"], ["buffer", "compressed_indices"], @@ -1756,10 +1756,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "5bed0c801c9f03c5a4157722eca0768a8457156f1aab8d80ee1d002b00029ca5", + "identity": "6008a89e4a7f9470bbd0b258bad5d335dc84a0027763f5c00745d93787c5efdf", "sources": [ - "cake_dsv4_4c41fb8347f137d4d46f_kernel.cu", - "cake_dsv4_4c41fb8347f137d4d46f_binding.cu", + "cake_dsv4_064850af5858dd03c104_kernel.cu", + "cake_dsv4_064850af5858dd03c104_binding.cu", ], "tma_workspace_bytes": 0, }, @@ -1771,7 +1771,7 @@ ["tma_buffer", "tmap_compressed_k"], ["tma_buffer", "tmap_swa_v"], ["tma_buffer", "tmap_compressed_v"], - ["buffer", "O"], + ["tma_buffer", "O"], ["buffer", "partial_lse"], ["buffer", "swa_indices"], ["buffer", "compressed_indices"], @@ -1794,10 +1794,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "1ca21a0f7b4c6edfddb5340b989d3cc2964c772e7ecd16fc74bcea8088fd5eec", + "identity": "2adf6df71253d0032819574f14607f44a8c86dae4c489e3d77dd50806316dcf6", "sources": [ - "cake_dsv4_125332c2066d89bd7610_kernel.cu", - "cake_dsv4_125332c2066d89bd7610_binding.cu", + "cake_dsv4_4ab0487d5fde70d4569e_kernel.cu", + "cake_dsv4_4ab0487d5fde70d4569e_binding.cu", ], "tma_workspace_bytes": 0, }, @@ -1826,7 +1826,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "181de2841cecdd15e69e910e6afec211743131dcf5fb7fed2023027d0570b314", + "identity": "19423e1c286750ba3f6842d6809313687520c8ae382c1e5b3f14297c75155a96", "sources": [ "cake_dsv4_ae27071dabfc85fb47c1_kernel.cu", "cake_dsv4_ae27071dabfc85fb47c1_binding.cu", @@ -1864,7 +1864,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "456fc2e0220f173e148d85b2fd7b3d73c245f83a64a29847de39a0a29b8a183b", + "identity": "2780b283e3889a568e1f4bce3009a3a3d71b4c076930bc8481381308d5379d67", "sources": [ "cake_dsv4_6b7225f9e37a48354992_kernel.cu", "cake_dsv4_6b7225f9e37a48354992_binding.cu", @@ -1899,7 +1899,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "9f675bcf369d90feaf54659043953dd0734c40b14e65b5cc2e8db3ad756b6b3f", + "identity": "3c9bff552cbbe14c33c1d07469dad72b836d2c3816b91dcf6150b1affc5f9aed", "sources": [ "cake_dsv4_fb5ca25f5af355ef89b5_kernel.cu", "cake_dsv4_fb5ca25f5af355ef89b5_binding.cu", @@ -1921,7 +1921,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "e88d9e3ab9ee0f46c4d50a03e328e447396b8ab121c01f134d447065b10394f0", + "identity": "8b08b4fd9e1a83550104302232433821bb3a83d4ce9762fcfee3fd540bc2af59", "sources": [ "cake_dsv4_0590fb03afd3391febef_kernel.cu", "cake_dsv4_0590fb03afd3391febef_binding.cu", @@ -1960,10 +1960,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "6991e3e7c4210570953f9cdc15c6293f15bf9a0e1798b739abee8df34c6d3551", + "identity": "7cbab296d40363c89a8116b277938459a7a15cbd3ea2bad99ac36d24f0f92166", "sources": [ - "cake_dsv4_c967e7efe220160e28d3_kernel.cu", - "cake_dsv4_c967e7efe220160e28d3_binding.cu", + "cake_dsv4_1ff52fece183b851b441_kernel.cu", + "cake_dsv4_1ff52fece183b851b441_binding.cu", ], "tma_workspace_bytes": 384, }, @@ -1987,7 +1987,6 @@ ["parameter", "sparse_topk"], ["parameter", "num_query_tokens"], ["parameter", "has_sinks"], - ["workspace", "tma_descriptor_workspace"], ["grid", "grid_x"], ["grid", "grid_y"], ["grid", "grid_z"], @@ -1995,12 +1994,12 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "aae71071ae4aa66b4bdd00a852c061efda61475dc4a69cdf4b2bd2705f4a4661", + "identity": "6f16d1265dc27464472cc2464c8faa84a788e83345d6c5fa698c57897dcad362", "sources": [ - "cake_dsv4_3e5a644d61420c9c7ea1_kernel.cu", - "cake_dsv4_3e5a644d61420c9c7ea1_binding.cu", + "cake_dsv4_1e3cf4361de31c58b69b_kernel.cu", + "cake_dsv4_1e3cf4361de31c58b69b_binding.cu", ], - "tma_workspace_bytes": 384, + "tma_workspace_bytes": 0, }, "bf16_h8_h16_source_exact": { "arch": "sm_103a", @@ -2032,10 +2031,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "ad819a642ec063a95836b08029677d93d1c6516c4b096e3106f9032257a12bd0", + "identity": "a08e408c6cdc5bbb63728eae56cb3c556c8b15513eb1a23d68d25712751ba903", "sources": [ - "cake_dsv4_d23652aae3370c30d978_kernel.cu", - "cake_dsv4_d23652aae3370c30d978_binding.cu", + "cake_dsv4_e48c5e0a89a41eb1e57a_kernel.cu", + "cake_dsv4_e48c5e0a89a41eb1e57a_binding.cu", ], "tma_workspace_bytes": 384, }, @@ -2064,7 +2063,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "12c3d8e97fc54d26fe53f365da8b548e3cfe388c84a0101316c88a55a840ab0e", + "identity": "a837022a12132a14e91375f54c3ff7fcabea99a3cae8d5583d13070284df984f", "sources": [ "cake_dsv4_fdbe4c4737d7f6dd1d9e_kernel.cu", "cake_dsv4_fdbe4c4737d7f6dd1d9e_binding.cu", @@ -2096,10 +2095,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "5b806976ca4dcb0b14105e2f1b06df0630452af2283a61256c8df96b0ec8b7e4", + "identity": "9503eaa1245660fa320fbbd82dfec0253c49d57bdd1b2208f3dcc687d8cb1587", "sources": [ - "cake_dsv4_9cf30aacb43b71034660_kernel.cu", - "cake_dsv4_9cf30aacb43b71034660_binding.cu", + "cake_dsv4_d9754db4fe394a927e67_kernel.cu", + "cake_dsv4_d9754db4fe394a927e67_binding.cu", ], "tma_workspace_bytes": 0, }, @@ -2140,7 +2139,7 @@ ], "entry": "run", "host_linkage_flags": [], - "identity": "72c980c50469dd2c8c22342b2384ccd6d6770535e77343c9aea15c9f1f99860f", + "identity": "8a292d1b18d0e15c944d598de09148bc80b5d5bd09a855539482791253d4ff15", "sources": [ "cake_dsv4_85ef6cd9a950ca02c484_kernel.cu", "cake_dsv4_85ef6cd9a950ca02c484_binding.cu", @@ -2181,7 +2180,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "4cbb2c362b859ca355d32bb1c802aabf58d88fd8f9ee36140273165754e8820c", + "identity": "cfef2a7086ba2562eeebc50d196c8486ab4522e3b059ce3a0952915c03d8dad0", "sources": [ "cake_dsv4_32e01d0e8961dfa26bf3_kernel.cu", "cake_dsv4_32e01d0e8961dfa26bf3_binding.cu", @@ -2221,7 +2220,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "9499459d9e38d660258217d59ae733e72c9421c68b93e53534379a4353cf4145", + "identity": "587da4e34f457b41b1b540e71a2a841fe03c0a5e0f056e846da82eeafe874185", "sources": [ "cake_dsv4_bd3f29f66e3196dc9451_kernel.cu", "cake_dsv4_bd3f29f66e3196dc9451_binding.cu", @@ -2256,10 +2255,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "a0c09ef4273aa62ee82bc9f0f221af205a14fbbc44509bbdf674ff321f5f66b6", + "identity": "910f0fea7bd4f6be5f5b62ddb6ca05b4365799b0e8fe122b4a02018ee0da8ae4", "sources": [ - "cake_dsv4_f55f69afe4ac3ec96817_kernel.cu", - "cake_dsv4_f55f69afe4ac3ec96817_binding.cu", + "cake_dsv4_0a925b9bcb3044777c76_kernel.cu", + "cake_dsv4_0a925b9bcb3044777c76_binding.cu", ], "tma_workspace_bytes": 0, }, @@ -2293,10 +2292,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "d29ebd3f5584b9189fb165159f8117f93dcf919fd9b35bf6ddfccdbbb3f05593", + "identity": "30120c1636cb83664be716d2b2df67d2f2dd5f8faaa9b7b2e05237fa9a402557", "sources": [ - "cake_dsv4_7a4b9556b3c84ebff6ce_kernel.cu", - "cake_dsv4_7a4b9556b3c84ebff6ce_binding.cu", + "cake_dsv4_c1d444bba4cdb8273c7c_kernel.cu", + "cake_dsv4_c1d444bba4cdb8273c7c_binding.cu", ], "tma_workspace_bytes": 384, }, @@ -2329,10 +2328,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "7ff0f4e03616e8a841b5726c9bae67e2ec29b56d21c7da962980cf1d0fba8d57", + "identity": "22eff1cf99b136db0c474ddb719cbca1a10a199b8d7a1639c83182dad1c6ac15", "sources": [ - "cake_dsv4_9e1188b192c9c806d7c6_kernel.cu", - "cake_dsv4_9e1188b192c9c806d7c6_binding.cu", + "cake_dsv4_983db4c490ceed3c455a_kernel.cu", + "cake_dsv4_983db4c490ceed3c455a_binding.cu", ], "tma_workspace_bytes": 0, }, @@ -2365,10 +2364,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "936e66b0b3c4a12668dcca14881ec70097823c247ca505c455fd0db5da4c99b5", + "identity": "4e9cdeed8ed71d70dc7b1d02f7a417e006a6ff14789f8af8779658b85fa60c5d", "sources": [ - "cake_dsv4_c20607c57cc11784015e_kernel.cu", - "cake_dsv4_c20607c57cc11784015e_binding.cu", + "cake_dsv4_98db1f55c8c129e16005_kernel.cu", + "cake_dsv4_98db1f55c8c129e16005_binding.cu", ], "tma_workspace_bytes": 0, }, @@ -2402,10 +2401,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "61781095b5e9377262e277a85c730dcd0bb3588b3633cbe248765b64980e8e23", + "identity": "325bbfbb67496e600ba75e025966b45f9df4db915c0e75e7d27f1438a6ad5036", "sources": [ - "cake_dsv4_a2be51a1fd34ec0d9a3b_kernel.cu", - "cake_dsv4_a2be51a1fd34ec0d9a3b_binding.cu", + "cake_dsv4_868bb2291a1464f14e05_kernel.cu", + "cake_dsv4_868bb2291a1464f14e05_binding.cu", ], "tma_workspace_bytes": 384, }, @@ -2424,7 +2423,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "2a8255d11ab33428fbaa167fb8b2a10bf0267969da6ddfb81d9974cd7b62296f", + "identity": "a1b1f15c277a8778c33e5288abfbbb1dbacf1e2c6625e3e1fe39277c65023443", "sources": [ "cake_dsv4_230143c89d42c6c89387_kernel.cu", "cake_dsv4_230143c89d42c6c89387_binding.cu", @@ -2446,7 +2445,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "3f909a9786eea6a97a37b095198212ee64c0fddc5e3d9e362fa48443f946c528", + "identity": "ebb9afd30b8660b2b019a33b9d32b1468d7dd2f86458522b4b720506188ed321", "sources": [ "cake_dsv4_397f817e2931a6cc7c92_kernel.cu", "cake_dsv4_397f817e2931a6cc7c92_binding.cu", From 1042f9c1bbcb4064fd7c4a564e03f595e75288cf Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Wed, 30 Sep 2026 00:31:20 -0700 Subject: [PATCH 2/3] perf(cake_sparse_mla): regenerate sm_100a programs from Cake 56f72c78456 (round-7 export) Co-Authored-By: Claude Fable 5.1 --- .../cake_dsv4_09e14f1d9cbd6739db61_binding.cu | 454 +++ .../cake_dsv4_09e14f1d9cbd6739db61_kernel.cu | 1678 ++++++++++ .../cake_dsv4_0c8ab955ac7b62e7bfee_binding.cu | 389 +++ .../cake_dsv4_0c8ab955ac7b62e7bfee_kernel.cu | 1743 ++++++++++ .../cake_dsv4_12ac507a3c73f2784621_binding.cu | 385 +++ .../cake_dsv4_12ac507a3c73f2784621_kernel.cu | 1735 ++++++++++ .../cake_dsv4_15eabf5d2a5f50db201b_binding.cu | 436 +++ .../cake_dsv4_15eabf5d2a5f50db201b_kernel.cu | 2557 ++++++++++++++ .../cake_dsv4_21a93d40caefcf5b252d_binding.cu | 446 +++ .../cake_dsv4_21a93d40caefcf5b252d_kernel.cu | 2646 +++++++++++++++ .../cake_dsv4_42fa8e98da5fa46d1a94_binding.cu | 606 ++++ .../cake_dsv4_42fa8e98da5fa46d1a94_kernel.cu | 2364 +++++++++++++ .../cake_dsv4_530fe7e6cb8c7d09744d_binding.cu | 445 +++ .../cake_dsv4_530fe7e6cb8c7d09744d_kernel.cu | 2912 ++++++++++++++++ .../cake_dsv4_5703afd5b398f2b6c980_binding.cu | 461 +++ .../cake_dsv4_5703afd5b398f2b6c980_kernel.cu | 2803 ++++++++++++++++ .../cake_dsv4_714df747e199d0adbe08_binding.cu | 606 ++++ .../cake_dsv4_714df747e199d0adbe08_kernel.cu | 2378 +++++++++++++ .../cake_dsv4_9c3b9b4fa7cba6e44ffc_binding.cu | 461 +++ .../cake_dsv4_9c3b9b4fa7cba6e44ffc_kernel.cu | 2826 ++++++++++++++++ .../cake_dsv4_bab2852cc0ea4a6d4860_binding.cu | 441 +++ .../cake_dsv4_bab2852cc0ea4a6d4860_kernel.cu | 2968 +++++++++++++++++ ...sv4_family_5383a615a7e502f6fdac_binding.cu | 703 ++++ ...sv4_family_9c96528827a52c42bbd7_binding.cu | 703 ++++ flashinfer/jit/cake_dsv4.py | 110 +- 25 files changed, 33201 insertions(+), 55 deletions(-) create mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_09e14f1d9cbd6739db61_binding.cu create mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_09e14f1d9cbd6739db61_kernel.cu create mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_0c8ab955ac7b62e7bfee_binding.cu create mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_0c8ab955ac7b62e7bfee_kernel.cu create mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_12ac507a3c73f2784621_binding.cu create mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_12ac507a3c73f2784621_kernel.cu create mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_15eabf5d2a5f50db201b_binding.cu create mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_15eabf5d2a5f50db201b_kernel.cu create mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_21a93d40caefcf5b252d_binding.cu create mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_21a93d40caefcf5b252d_kernel.cu create mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_42fa8e98da5fa46d1a94_binding.cu create mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_42fa8e98da5fa46d1a94_kernel.cu create mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_530fe7e6cb8c7d09744d_binding.cu create mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_530fe7e6cb8c7d09744d_kernel.cu create mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_5703afd5b398f2b6c980_binding.cu create mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_5703afd5b398f2b6c980_kernel.cu create mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_714df747e199d0adbe08_binding.cu create mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_714df747e199d0adbe08_kernel.cu create mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_9c3b9b4fa7cba6e44ffc_binding.cu create mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_9c3b9b4fa7cba6e44ffc_kernel.cu create mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_bab2852cc0ea4a6d4860_binding.cu create mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_bab2852cc0ea4a6d4860_kernel.cu create mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_family_5383a615a7e502f6fdac_binding.cu create mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_family_9c96528827a52c42bbd7_binding.cu diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_09e14f1d9cbd6739db61_binding.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_09e14f1d9cbd6739db61_binding.cu new file mode 100644 index 00000000000..20cc47f3686 --- /dev/null +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_09e14f1d9cbd6739db61_binding.cu @@ -0,0 +1,454 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +// Generated by the Cake source exporter — do not edit. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_09e14f1d9cbd6739db61'. +#include +#include +#include +#include + +#include "tvm_ffi_utils.h" + +#include +#include +#include +#include +#include +#include +#include + +extern "C" __global__ void kernel_cake_dsv4_09e14f1d9cbd6739db61(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int batch_size, int max_q_len, int ragged_query, int has_sinks); + + +namespace cake_host_shim_db4b534bf8fd17b7 { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t CakeDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +// 4D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + int64_t d3 = t.size(t.ndim() - 3); + TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(d3)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 8u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 64, 8, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 64u, 8u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_swa_kv' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_swa_kv' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_kv'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_compressed_kv' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_compressed_kv(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_compressed_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_compressed_kv' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_compressed_kv' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_compressed_kv"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_compressed_kv' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_kv'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_kv') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_compressed_kv, TensorView arg_O, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_seq_lens, TensorView arg_cum_seq_lens_q, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_batch_size, int64_t arg_max_q_len, int64_t arg_ragged_query, int64_t arg_has_sinks, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckContiguous(arg_tmap_q, "tmap_q"); + CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); + CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 4, 16, 1); + CheckCudaTensor(arg_tmap_compressed_kv, "tmap_compressed_kv"); + CheckDtype(arg_tmap_compressed_kv, "tmap_compressed_kv", 4, 16, 1); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_swa_indices, "swa_indices"); + CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); + CheckContiguous(arg_swa_indices, "swa_indices"); + CheckCudaTensor(arg_compressed_indices, "compressed_indices"); + CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); + CheckContiguous(arg_compressed_indices, "compressed_indices"); + CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); + CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckCudaTensor(arg_seq_lens, "seq_lens"); + CheckDtype(arg_seq_lens, "seq_lens", 0, 32, 1); + CheckContiguous(arg_seq_lens, "seq_lens"); + CheckCudaTensor(arg_cum_seq_lens_q, "cum_seq_lens_q"); + CheckDtype(arg_cum_seq_lens_q, "cum_seq_lens_q", 0, 32, 1); + CheckContiguous(arg_cum_seq_lens_q, "cum_seq_lens_q"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); + CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); + CheckContiguous(arg_bmm1_scale, "bmm1_scale"); + CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); + CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); + CheckContiguous(arg_bmm2_scale, "bmm2_scale"); + TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) + << "scalar 'num_heads' value " << arg_num_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) + << "scalar 'swa_index_stride' value " << arg_swa_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) + << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) + << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) + << "scalar 'sparse_topk' value " << arg_sparse_topk + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_batch_size >= -2147483648LL && arg_batch_size <= 2147483647LL, ValueError) + << "scalar 'batch_size' value " << arg_batch_size + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_q_len >= -2147483648LL && arg_max_q_len <= 2147483647LL, ValueError) + << "scalar 'max_q_len' value " << arg_max_q_len + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_ragged_query >= -2147483648LL && arg_ragged_query <= 2147483647LL, ValueError) + << "scalar 'ragged_query' value " << arg_ragged_query + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) + << "scalar 'has_sinks' value " << arg_has_sinks + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); + CheckSameCudaDevice(arg_tmap_compressed_kv, arg_tmap_q, "tmap_compressed_kv", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); + CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); + CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); + CheckSameCudaDevice(arg_seq_lens, arg_tmap_q, "seq_lens", "tmap_q"); + CheckSameCudaDevice(arg_cum_seq_lens_q, arg_tmap_q, "cum_seq_lens_q", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); + CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + + cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); + CUtensorMap p_tmap_compressed_kv = EncodeTma_tmap_compressed_kv(arg_tmap_compressed_kv); + __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); + int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); + int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); + int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); + int* p_seq_lens = static_cast(arg_seq_lens.data_ptr()); + int* p_cum_seq_lens_q = static_cast(arg_cum_seq_lens_q.data_ptr()); + float* p_sinks = static_cast(arg_sinks.data_ptr()); + float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); + float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); + int32_t v_num_heads = (int32_t)arg_num_heads; + int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; + int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; + int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; + int32_t v_sparse_topk = (int32_t)arg_sparse_topk; + int32_t v_batch_size = (int32_t)arg_batch_size; + int32_t v_max_q_len = (int32_t)arg_max_q_len; + int32_t v_ragged_query = (int32_t)arg_ragged_query; + int32_t v_has_sinks = (int32_t)arg_has_sinks; + void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_tmap_compressed_kv, &p_O, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_seq_lens, &p_cum_seq_lens_q, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_batch_size, &v_max_q_len, &v_ragged_query, &v_has_sinks}; + + static std::mutex smem_mu; + static auto* smem_status_by_device = + new std::unordered_map(); + cudaError_t smem_status = cudaSuccess; + { + std::lock_guard lock(smem_mu); + auto it = smem_status_by_device->find(dev.device_id); + if (it == smem_status_by_device->end()) { + int current_device = -1; + smem_status = cudaGetDevice(¤t_device); + if (smem_status == cudaSuccess && current_device != dev.device_id) { + smem_status = cudaErrorInvalidDevice; + } + if (smem_status == cudaSuccess) { + smem_status = cudaFuncSetAttribute( + reinterpret_cast(kernel_cake_dsv4_09e14f1d9cbd6739db61), + cudaFuncAttributeMaxDynamicSharedMemorySize, + 198144); + } + smem_status_by_device->emplace(dev.device_id, smem_status); + } else { + smem_status = it->second; + } + } + TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) + << "cudaFuncSetAttribute for kernel_cake_dsv4_09e14f1d9cbd6739db61 failed: " + << cudaGetErrorString(smem_status); + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(256u, 1u, 1u); + + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes = 198144u; + config.stream = stream; + cudaError_t launch_status = cudaLaunchKernelExC( + &config, reinterpret_cast(kernel_cake_dsv4_09e14f1d9cbd6739db61), kargs); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernelExC for kernel_cake_dsv4_09e14f1d9cbd6739db61 failed: " + << cudaGetErrorString(launch_status); +} + +} // namespace cake_host_shim_db4b534bf8fd17b7 + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_db4b534bf8fd17b7::Run); diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_09e14f1d9cbd6739db61_kernel.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_09e14f1d9cbd6739db61_kernel.cu new file mode 100644 index 00000000000..bdf35286adc --- /dev/null +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_09e14f1d9cbd6739db61_kernel.cu @@ -0,0 +1,1678 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +#define CAKE_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_OFFSET 0 +#define NUM_KV_PIPE_STAGES 4 +#define SMEM_SMEM_Q_OFF 1024 +#define SMEM_SMEM_Q_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_STRIDE 8192 +#define SMEM_SMEM_KV_OFF 66560 +#define SMEM_SMEM_KV_STAGE_BYTES 32768 +#define SMEM_SMEM_KV_STRIDE 32768 +#define SMEM_SMEM_V_OFF 66560 +#define SMEM_SMEM_V_STAGE_BYTES 32768 +#define SMEM_SMEM_V_STRIDE 32768 +#define SMEM_TOTAL 198144 +#define THREADS 256 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count) : "memory"); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ float2 ex2_emulation_f32x2_value(float2 value) { + const float c0 = 1.0f, c1 = 0.695146143436431884765625f; + const float c2 = 0.227564394474029541015625f, c3 = 0.077119089663028717041015625f; + const float magic = 12582912.0f; + float x0 = max_noftz(value.x, -127.0f), x1 = max_noftz(value.y, -127.0f); + float2 xc2 = make_float2(x0, x1), magic2 = make_float2(magic, magic); + float2 xr2; + asm("add.rm.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xr2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&magic2)); + float2 c3_2 = make_float2(c3, c3), c2_2 = make_float2(c2, c2); + float2 c1_2 = make_float2(c1, c1), c0_2 = make_float2(c0, c0); + float2 xrb2, xfrac2; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xrb2) + : "l"(*(unsigned long long*)&xr2), "l"(*(unsigned long long*)&magic2)); + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xfrac2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&xrb2)); + float2 poly2; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&c3_2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c2_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c1_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c0_2)); + int x0r_i, x1r_i, p0_i, p1_i; + asm("mov.b64 {%0, %1}, %2;" : "=r"(x0r_i), "=r"(x1r_i) : "l"(*(unsigned long long*)&xr2)); + asm("mov.b64 {%0, %1}, %2;" : "=r"(p0_i), "=r"(p1_i) : "l"(*(unsigned long long*)&poly2)); + float r0, r1; + asm("mov.b32 %0, %1;" : "=f"(r0) : "r"((x0r_i << 23) + p0_i)); + asm("mov.b32 %0, %1;" : "=f"(r1) : "r"((x1r_i << 23) + p1_i)); + return make_float2(r0, r1); +} + +__device__ __forceinline__ void ex2_emulation_f32x2(float* x0_ptr, float* x1_ptr) { + float2 result = ex2_emulation_f32x2_value(make_float2(*x0_ptr, *x1_ptr)); + *x0_ptr = result.x; *x1_ptr = result.y; +} + +__device__ __forceinline__ void softmax_frag_exp2_cast( + float* sv, uint32_t* pv, int use_emu) +{ + #pragma unroll + for (int j = 0; j < 16; j++) { + if (use_emu && j >= 12) + ex2_emulation_f32x2(&sv[j*2], &sv[j*2+1]); + else { + sv[j*2] = approx_exp2(sv[j*2]); + sv[j*2+1] = approx_exp2(sv[j*2+1]); + } + } + #pragma unroll + for (int j = 0; j < 16; j++) { + __nv_bfloat162 bf = __float22bfloat162_rn({sv[j*2], sv[j*2+1]}); + pv[j] = reinterpret_cast(bf); + } +} + + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_2d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3}], [%4];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr) { + // Canonical .shared::cta form for non-multicast gather4, matching + // trtllm-gen / cuda_ptx and the PTX ISA qualifier order + // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, + // .shared::cluster is reserved for the multicast variant (ctaMask). + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(256, 1) void +kernel_cake_dsv4_09e14f1d9cbd6739db61(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int batch_size, int max_q_len, int ragged_query, int has_sinks) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); + + const int mbar_base = smem; + #define q_full_addr (mbar_base + 0) + #define kv_full_addr (mbar_base + 8) + #define s_full_addr (mbar_base + 40) + #define p_full_addr (mbar_base + 56) + #define pv_done_addr (mbar_base + 72) + #define tmem_dealloc_addr (mbar_base + 80) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + + // Kernel setup ops + __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_q_addr = smem + 1024; + __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_kv_addr = smem + 66560; + __nv_bfloat16* smem_v = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_v_addr = smem + 66560; + + // Mbarrier init (6 pipeline groups, 0 ordered-sequence groups, 11 barriers) + // Mbarriers at smem_raw[0..88) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // q_full: 1 barriers, init_count=1 + mbarrier_init(smem + 0, 1); + // --- pipeline 'kv_pipe' --- + // kv_full: 4 barriers, init_count=3 + mbarrier_init(smem + 8, 3); + mbarrier_init(smem + 16, 3); + mbarrier_init(smem + 24, 3); + mbarrier_init(smem + 32, 3); + // s_full: 2 barriers, init_count=1 + mbarrier_init(smem + 40, 1); + mbarrier_init(smem + 48, 1); + // p_full: 2 barriers, init_count=4 + mbarrier_init(smem + 56, 4); + mbarrier_init(smem + 64, 4); + // pv_done: 1 barriers, init_count=1 + mbarrier_init(smem + 72, 1); + // tmem_dealloc: 1 barriers, init_count=7 + mbarrier_init(smem + 80, 7); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncwarp(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 88); + if (warp == 0) { + int _tmem_hold = smem + 88; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_tmem = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 64;"); + } + + // ---- Role: softmax ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 176;"); + { // softmax_main + float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; + int query_idx = blockIdx.x; + int _max_2 = ((sparse_topk_lens[query_idx] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx] + sparse_topk_lens_offset) : (0)); + int _min_2 = ((_max_2) < (sparse_topk) ? (_max_2) : (sparse_topk)); + int active_topk = _min_2; + int tile_count = (active_topk + 128 - 1) / 128; + int query_batch = query_idx / max_q_len; + int query_offset = query_idx - query_batch * max_q_len; + int query_length = max_q_len; + if (ragged_query != 0) { + query_batch = 0; + #pragma unroll 1 + for (int batch = 0; batch < batch_size; batch++) { + if (query_idx >= cum_seq_lens_q[batch + 1]) { + query_batch = batch + 1; + } + } + int query_begin = cum_seq_lens_q[query_batch]; + query_length = cum_seq_lens_q[query_batch + 1] - query_begin; + query_offset = query_idx - query_begin; + } + int swa_valid_columns = seq_lens[query_batch] - query_length + query_offset + 1; + if (swa_valid_columns < 0) { + swa_valid_columns = 0; + } + if (swa_valid_columns > 128) { + swa_valid_columns = 128; + } + const int warp_in_compute = warp; + const int tmem_row_origin = warp_in_compute * 32; + const int logical_row_origin = warp_in_compute * 16; + const int my_row = logical_row_origin + lane % 16; + const int col_half = lane / 16; + float row_max = -CAKE_INF; + float row_sum = 0.0f; + if (has_sinks != 0 && my_row < num_heads) { + row_max = sinks[my_row] * 1.4426950408889634f / softmax_scale_log2; + row_sum = 1.0f; + } + #pragma unroll + for (int tile = 0; tile < 1; tile++) { + if (tile_count > tile) { + int phase = tile & 1; + int wait_phase = tile >> 1 & 1; + mbarrier_wait(s_full_addr + (phase) * 8, wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int score_col = ((phase != 0) ? 128 : 0); + int score_addr = taddr + (unsigned int)score_col + (unsigned int)(tmem_row_origin << 16); + float _tmem_load_0[64]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_0[0])), "=r"(*reinterpret_cast(&_tmem_load_0[1])), "=r"(*reinterpret_cast(&_tmem_load_0[2])), "=r"(*reinterpret_cast(&_tmem_load_0[3])), "=r"(*reinterpret_cast(&_tmem_load_0[4])), "=r"(*reinterpret_cast(&_tmem_load_0[5])), "=r"(*reinterpret_cast(&_tmem_load_0[6])), "=r"(*reinterpret_cast(&_tmem_load_0[7])), "=r"(*reinterpret_cast(&_tmem_load_0[8])), "=r"(*reinterpret_cast(&_tmem_load_0[9])), "=r"(*reinterpret_cast(&_tmem_load_0[10])), "=r"(*reinterpret_cast(&_tmem_load_0[11])), "=r"(*reinterpret_cast(&_tmem_load_0[12])), "=r"(*reinterpret_cast(&_tmem_load_0[13])), "=r"(*reinterpret_cast(&_tmem_load_0[14])), "=r"(*reinterpret_cast(&_tmem_load_0[15])), "=r"(*reinterpret_cast(&_tmem_load_0[16])), "=r"(*reinterpret_cast(&_tmem_load_0[17])), "=r"(*reinterpret_cast(&_tmem_load_0[18])), "=r"(*reinterpret_cast(&_tmem_load_0[19])), "=r"(*reinterpret_cast(&_tmem_load_0[20])), "=r"(*reinterpret_cast(&_tmem_load_0[21])), "=r"(*reinterpret_cast(&_tmem_load_0[22])), "=r"(*reinterpret_cast(&_tmem_load_0[23])), "=r"(*reinterpret_cast(&_tmem_load_0[24])), "=r"(*reinterpret_cast(&_tmem_load_0[25])), "=r"(*reinterpret_cast(&_tmem_load_0[26])), "=r"(*reinterpret_cast(&_tmem_load_0[27])), "=r"(*reinterpret_cast(&_tmem_load_0[28])), "=r"(*reinterpret_cast(&_tmem_load_0[29])), "=r"(*reinterpret_cast(&_tmem_load_0[30])), "=r"(*reinterpret_cast(&_tmem_load_0[31])) + : "r"(score_addr)); + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_0[32])), "=r"(*reinterpret_cast(&_tmem_load_0[33])), "=r"(*reinterpret_cast(&_tmem_load_0[34])), "=r"(*reinterpret_cast(&_tmem_load_0[35])), "=r"(*reinterpret_cast(&_tmem_load_0[36])), "=r"(*reinterpret_cast(&_tmem_load_0[37])), "=r"(*reinterpret_cast(&_tmem_load_0[38])), "=r"(*reinterpret_cast(&_tmem_load_0[39])), "=r"(*reinterpret_cast(&_tmem_load_0[40])), "=r"(*reinterpret_cast(&_tmem_load_0[41])), "=r"(*reinterpret_cast(&_tmem_load_0[42])), "=r"(*reinterpret_cast(&_tmem_load_0[43])), "=r"(*reinterpret_cast(&_tmem_load_0[44])), "=r"(*reinterpret_cast(&_tmem_load_0[45])), "=r"(*reinterpret_cast(&_tmem_load_0[46])), "=r"(*reinterpret_cast(&_tmem_load_0[47])), "=r"(*reinterpret_cast(&_tmem_load_0[48])), "=r"(*reinterpret_cast(&_tmem_load_0[49])), "=r"(*reinterpret_cast(&_tmem_load_0[50])), "=r"(*reinterpret_cast(&_tmem_load_0[51])), "=r"(*reinterpret_cast(&_tmem_load_0[52])), "=r"(*reinterpret_cast(&_tmem_load_0[53])), "=r"(*reinterpret_cast(&_tmem_load_0[54])), "=r"(*reinterpret_cast(&_tmem_load_0[55])), "=r"(*reinterpret_cast(&_tmem_load_0[56])), "=r"(*reinterpret_cast(&_tmem_load_0[57])), "=r"(*reinterpret_cast(&_tmem_load_0[58])), "=r"(*reinterpret_cast(&_tmem_load_0[59])), "=r"(*reinterpret_cast(&_tmem_load_0[60])), "=r"(*reinterpret_cast(&_tmem_load_0[61])), "=r"(*reinterpret_cast(&_tmem_load_0[62])), "=r"(*reinterpret_cast(&_tmem_load_0[63])) + : "r"(score_addr + 32)); + int valid_cols = active_topk - tile * 128 - col_half * 64; + if (valid_cols < 0) { + valid_cols = 0; + } + if (valid_cols > 64) { + valid_cols = 64; + } + if (tile == 0) { + int valid_swa_cols = swa_valid_columns - col_half * 64; + if (valid_swa_cols < 0) { + valid_swa_cols = 0; + } + if (valid_swa_cols < valid_cols) { + valid_cols = valid_swa_cols; + } + } + uint32_t _slice_lo_mask_0; + { + int _lim_0 = valid_cols; + if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_0 & (1u << 0))) _tmem_load_0[0] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 1))) _tmem_load_0[1] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 2))) _tmem_load_0[2] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 3))) _tmem_load_0[3] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 4))) _tmem_load_0[4] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 5))) _tmem_load_0[5] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 6))) _tmem_load_0[6] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 7))) _tmem_load_0[7] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 8))) _tmem_load_0[8] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 9))) _tmem_load_0[9] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 10))) _tmem_load_0[10] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 11))) _tmem_load_0[11] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 12))) _tmem_load_0[12] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 13))) _tmem_load_0[13] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 14))) _tmem_load_0[14] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 15))) _tmem_load_0[15] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 16))) _tmem_load_0[16] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 17))) _tmem_load_0[17] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 18))) _tmem_load_0[18] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 19))) _tmem_load_0[19] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 20))) _tmem_load_0[20] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 21))) _tmem_load_0[21] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 22))) _tmem_load_0[22] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 23))) _tmem_load_0[23] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 24))) _tmem_load_0[24] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 25))) _tmem_load_0[25] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 26))) _tmem_load_0[26] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 27))) _tmem_load_0[27] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 28))) _tmem_load_0[28] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 29))) _tmem_load_0[29] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 30))) _tmem_load_0[30] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 31))) _tmem_load_0[31] = -CAKE_INF; + uint32_t _slice_lo_mask_1; + { + int _lim_1 = valid_cols - 32; + if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) _tmem_load_0[32] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) _tmem_load_0[33] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) _tmem_load_0[34] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) _tmem_load_0[35] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) _tmem_load_0[36] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) _tmem_load_0[37] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) _tmem_load_0[38] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) _tmem_load_0[39] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) _tmem_load_0[40] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) _tmem_load_0[41] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) _tmem_load_0[42] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) _tmem_load_0[43] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) _tmem_load_0[44] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) _tmem_load_0[45] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) _tmem_load_0[46] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) _tmem_load_0[47] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) _tmem_load_0[48] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) _tmem_load_0[49] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) _tmem_load_0[50] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) _tmem_load_0[51] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) _tmem_load_0[52] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) _tmem_load_0[53] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) _tmem_load_0[54] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) _tmem_load_0[55] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) _tmem_load_0[56] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) _tmem_load_0[57] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) _tmem_load_0[58] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) _tmem_load_0[59] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) _tmem_load_0[60] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) _tmem_load_0[61] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) _tmem_load_0[62] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) _tmem_load_0[63] = -CAKE_INF; + int* index_row = swa_indices + (query_idx * swa_index_stride + tile * 128 + col_half * 64); + #pragma unroll + for (int chunk = 0; chunk < 64; chunk += 4) { + if (valid_cols > chunk) { + int _vec_load_1[4]; + { + const int4* _ivptr_2 = reinterpret_cast(index_row + chunk + 0); + int4 _ivld_2; + _ivld_2 = *_ivptr_2; + _vec_load_1[0 + 0] = _ivld_2.x; + _vec_load_1[0 + 1] = _ivld_2.y; + _vec_load_1[0 + 2] = _ivld_2.z; + _vec_load_1[0 + 3] = _ivld_2.w; + } + #pragma unroll + for (int i = 0; i < 4; i++) { + if (_vec_load_1[i] < 0) { + _tmem_load_0[chunk + i] = -CAKE_INF; + } + } + } + } + float2 _reg_reduce_max2_3 = {-CAKE_INF, -CAKE_INF}; + row_max_x32_accum(&_tmem_load_0[0], _reg_reduce_max2_3); + row_max_x32_accum(&_tmem_load_0[32], _reg_reduce_max2_3); + float _tmem_load_0_max = row_max_reduce(_reg_reduce_max2_3); + float tile_max = _tmem_load_0_max; + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, tile_max, 16); + float _max_3 = max_noftz(tile_max, _shfl_xor_0); + tile_max = _max_3; + float _max_4 = max_noftz(row_max, tile_max); + float new_max = _max_4; + float _fma_0 = __fmaf_rn(row_max, softmax_scale_log2, (-new_max) * softmax_scale_log2); + float delta = _fma_0; + float _exp2_0 = approx_exp2(delta); + float acc_scale = ((row_max > -CAKE_INF) ? _exp2_0 : 1.0f); + float safe_max = ((new_max == -CAKE_INF) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + const float2 _fma_b2_4 = {softmax_scale_log2, softmax_scale_log2}; + const float2 _fma_c2_5 = {-max_scaled, -max_scaled}; + #pragma unroll + for (int _lf = 0; _lf < 32; _lf++) + fma_f32x2_inplace(&reinterpret_cast(_tmem_load_0)[_lf], _fma_b2_4, _fma_c2_5); + #pragma unroll + for (int _le = 0; _le < 64; _le++) { + _tmem_load_0[_le] = approx_exp2(_tmem_load_0[_le]); + } + unsigned int packed_p[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_0[_lp*2 + 0], _tmem_load_0[_lp*2+1 + 0])); + packed_p[_lp] = *(uint32_t*)&_bf2; + } + int p_addr = taddr + (unsigned int)score_col + 64 + (unsigned int)(tmem_row_origin << 16); + asm volatile( + "tcgen05.st.sync.aligned.16x32bx2.x16.b32" + " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(p_addr), "r"(*reinterpret_cast(&packed_p[0])), "r"(*reinterpret_cast(&packed_p[1])), "r"(*reinterpret_cast(&packed_p[2])), "r"(*reinterpret_cast(&packed_p[3])), "r"(*reinterpret_cast(&packed_p[4])), "r"(*reinterpret_cast(&packed_p[5])), "r"(*reinterpret_cast(&packed_p[6])), "r"(*reinterpret_cast(&packed_p[7])), "r"(*reinterpret_cast(&packed_p[8])), "r"(*reinterpret_cast(&packed_p[9])), "r"(*reinterpret_cast(&packed_p[10])), "r"(*reinterpret_cast(&packed_p[11])), "r"(*reinterpret_cast(&packed_p[12])), "r"(*reinterpret_cast(&packed_p[13])), "r"(*reinterpret_cast(&packed_p[14])), "r"(*reinterpret_cast(&packed_p[15]))); + asm volatile( + "tcgen05.st.sync.aligned.16x32bx2.x16.b32" + " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(p_addr + 16), "r"(*reinterpret_cast(&(packed_p + 16)[0])), "r"(*reinterpret_cast(&(packed_p + 16)[1])), "r"(*reinterpret_cast(&(packed_p + 16)[2])), "r"(*reinterpret_cast(&(packed_p + 16)[3])), "r"(*reinterpret_cast(&(packed_p + 16)[4])), "r"(*reinterpret_cast(&(packed_p + 16)[5])), "r"(*reinterpret_cast(&(packed_p + 16)[6])), "r"(*reinterpret_cast(&(packed_p + 16)[7])), "r"(*reinterpret_cast(&(packed_p + 16)[8])), "r"(*reinterpret_cast(&(packed_p + 16)[9])), "r"(*reinterpret_cast(&(packed_p + 16)[10])), "r"(*reinterpret_cast(&(packed_p + 16)[11])), "r"(*reinterpret_cast(&(packed_p + 16)[12])), "r"(*reinterpret_cast(&(packed_p + 16)[13])), "r"(*reinterpret_cast(&(packed_p + 16)[14])), "r"(*reinterpret_cast(&(packed_p + 16)[15]))); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + if (elect_sync()) { + mbarrier_arrive(p_full_addr + (phase) * 8); + } + float2 _reg_reduce_sum2_6 = make_float2(0.0f, 0.0f); + softmax_block_sum(&_tmem_load_0[0], &_reg_reduce_sum2_6); + softmax_block_sum(&_tmem_load_0[32], &_reg_reduce_sum2_6); + float _tmem_load_0_sum = _reg_reduce_sum2_6.x + _reg_reduce_sum2_6.y; + float block_sum = _tmem_load_0_sum; + float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, block_sum, 16); + block_sum = block_sum + _shfl_xor_1; + float _fma_1 = __fmaf_rn(row_sum, acc_scale, block_sum); + row_sum = _fma_1; + row_max = new_max; + } + } + if (sparse_topk == 128) { + mbarrier_wait(pv_done_addr, 0); + asm volatile("tcgen05.fence::after_thread_sync;"); + float output_scale = bmm2_scale[0]; + float _rcp_0 = approx_rcp(row_sum); + float inv_sum = ((row_sum > 0.0f) ? _rcp_0 : 0.0f); + int output_base = (query_idx * num_heads + my_row) * 512 + blockIdx.y * 256; + #pragma unroll + for (int v_stage = 0; v_stage < 2; v_stage++) { + int output_col = 256 + v_stage * 128; + int output_addr = taddr + (unsigned int)output_col + (unsigned int)(tmem_row_origin << 16); + float _tmem_load_1[64]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_1[0])), "=r"(*reinterpret_cast(&_tmem_load_1[1])), "=r"(*reinterpret_cast(&_tmem_load_1[2])), "=r"(*reinterpret_cast(&_tmem_load_1[3])), "=r"(*reinterpret_cast(&_tmem_load_1[4])), "=r"(*reinterpret_cast(&_tmem_load_1[5])), "=r"(*reinterpret_cast(&_tmem_load_1[6])), "=r"(*reinterpret_cast(&_tmem_load_1[7])), "=r"(*reinterpret_cast(&_tmem_load_1[8])), "=r"(*reinterpret_cast(&_tmem_load_1[9])), "=r"(*reinterpret_cast(&_tmem_load_1[10])), "=r"(*reinterpret_cast(&_tmem_load_1[11])), "=r"(*reinterpret_cast(&_tmem_load_1[12])), "=r"(*reinterpret_cast(&_tmem_load_1[13])), "=r"(*reinterpret_cast(&_tmem_load_1[14])), "=r"(*reinterpret_cast(&_tmem_load_1[15])), "=r"(*reinterpret_cast(&_tmem_load_1[16])), "=r"(*reinterpret_cast(&_tmem_load_1[17])), "=r"(*reinterpret_cast(&_tmem_load_1[18])), "=r"(*reinterpret_cast(&_tmem_load_1[19])), "=r"(*reinterpret_cast(&_tmem_load_1[20])), "=r"(*reinterpret_cast(&_tmem_load_1[21])), "=r"(*reinterpret_cast(&_tmem_load_1[22])), "=r"(*reinterpret_cast(&_tmem_load_1[23])), "=r"(*reinterpret_cast(&_tmem_load_1[24])), "=r"(*reinterpret_cast(&_tmem_load_1[25])), "=r"(*reinterpret_cast(&_tmem_load_1[26])), "=r"(*reinterpret_cast(&_tmem_load_1[27])), "=r"(*reinterpret_cast(&_tmem_load_1[28])), "=r"(*reinterpret_cast(&_tmem_load_1[29])), "=r"(*reinterpret_cast(&_tmem_load_1[30])), "=r"(*reinterpret_cast(&_tmem_load_1[31])) + : "r"(output_addr)); + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_1[32])), "=r"(*reinterpret_cast(&_tmem_load_1[33])), "=r"(*reinterpret_cast(&_tmem_load_1[34])), "=r"(*reinterpret_cast(&_tmem_load_1[35])), "=r"(*reinterpret_cast(&_tmem_load_1[36])), "=r"(*reinterpret_cast(&_tmem_load_1[37])), "=r"(*reinterpret_cast(&_tmem_load_1[38])), "=r"(*reinterpret_cast(&_tmem_load_1[39])), "=r"(*reinterpret_cast(&_tmem_load_1[40])), "=r"(*reinterpret_cast(&_tmem_load_1[41])), "=r"(*reinterpret_cast(&_tmem_load_1[42])), "=r"(*reinterpret_cast(&_tmem_load_1[43])), "=r"(*reinterpret_cast(&_tmem_load_1[44])), "=r"(*reinterpret_cast(&_tmem_load_1[45])), "=r"(*reinterpret_cast(&_tmem_load_1[46])), "=r"(*reinterpret_cast(&_tmem_load_1[47])), "=r"(*reinterpret_cast(&_tmem_load_1[48])), "=r"(*reinterpret_cast(&_tmem_load_1[49])), "=r"(*reinterpret_cast(&_tmem_load_1[50])), "=r"(*reinterpret_cast(&_tmem_load_1[51])), "=r"(*reinterpret_cast(&_tmem_load_1[52])), "=r"(*reinterpret_cast(&_tmem_load_1[53])), "=r"(*reinterpret_cast(&_tmem_load_1[54])), "=r"(*reinterpret_cast(&_tmem_load_1[55])), "=r"(*reinterpret_cast(&_tmem_load_1[56])), "=r"(*reinterpret_cast(&_tmem_load_1[57])), "=r"(*reinterpret_cast(&_tmem_load_1[58])), "=r"(*reinterpret_cast(&_tmem_load_1[59])), "=r"(*reinterpret_cast(&_tmem_load_1[60])), "=r"(*reinterpret_cast(&_tmem_load_1[61])), "=r"(*reinterpret_cast(&_tmem_load_1[62])), "=r"(*reinterpret_cast(&_tmem_load_1[63])) + : "r"(output_addr + 32)); + int col_base = v_stage * 128 + col_half * 64; + if (my_row < num_heads) { + #pragma unroll + for (int offset = 0; offset < 64; offset += 8) { + { + const float2 _prescale2_7 = {inv_sum * output_scale, inv_sum * output_scale}; + #if __CUDA_ARCH__ >= 1000 + #pragma unroll + for (int _ps = 0; _ps < 4; _ps++) + mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_1[offset])[_ps], _prescale2_7); + #else + #pragma unroll + for (int _ps = 0; _ps < 8; _ps++) + _tmem_load_1[offset + _ps] *= inv_sum * output_scale; + #endif + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_tmem_load_1[offset + 0], _tmem_load_1[offset + 1]); + _pk[1] = __floats2bfloat162_rn(_tmem_load_1[offset + 2], _tmem_load_1[offset + 3]); + _pk[2] = __floats2bfloat162_rn(_tmem_load_1[offset + 4], _tmem_load_1[offset + 5]); + _pk[3] = __floats2bfloat162_rn(_tmem_load_1[offset + 6], _tmem_load_1[offset + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(O + (output_base + col_base + offset)))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + } + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + if (elect_sync()) { + mbarrier_arrive(tmem_dealloc_addr); + } + } + } + // ---- Role: mma_warp ---- + if (warp == 4) { + { // mma_warp_main + int query_idx_1 = blockIdx.x; + int _max_1 = ((sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) : (0)); + int _min_1 = ((_max_1) < (sparse_topk) ? (_max_1) : (sparse_topk)); + int active_topk_1 = _min_1; + int tile_count_1 = (active_topk_1 + 128 - 1) / 128; + unsigned int _phase_q_full_0 = 0; + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int tile_1 = 0; tile_1 < 1; tile_1++) { + if (tile_count_1 > tile_1) { + int phase_1 = tile_1 & 1; + int score_col_1 = ((phase_1 != 0) ? 128 : 0); + #pragma unroll + for (int k_stage = 0; k_stage < 4; k_stage++) { + mbarrier_wait(kv_full_addr + (k_stage) * 8, phase_1); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_0 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (k_stage * 2) * 512); + int _mma_b_lo_0 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (k_stage) * 2048); + { + uint64_t _mma_ss_a_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_0); + uint64_t _mma_ss_b_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_0); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, ((k_stage == 0) ? 0 : 1)); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); + } + } + int _mma_a_lo_1 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (k_stage * 2 + 1) * 512); + int _mma_b_lo_1 = make_warp_uniform((((smem_kv_addr + 16384) >> 4) & 0x3FFF) + (k_stage) * 2048); + { + uint64_t _mma_ss_a_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_1); + uint64_t _mma_ss_b_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_1); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); + } + } + if (k_stage == 3) { + elect_commit(s_full_addr + (phase_1) * 8); + } + } + int wait_phase_1 = tile_1 >> 1 & 1; + mbarrier_wait(p_full_addr + (phase_1) * 8, wait_phase_1); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int v_stage_1 = 0; v_stage_1 < 2; v_stage_1++) { + int _mma_b_lo_2 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (blockIdx.y * 2 + v_stage_1) * 2048); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 dhi, blo, ta, id;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 69272720;\n\t" + "mov.b32 ta, %2;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p0;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "}\n" + :: "r"((tmem_tmem + (256 + v_stage_1 * 128))), "r"(_mma_b_lo_2), "r"(tmem_tmem + (score_col_1 + 64)), "r"(0)); + } + elect_commit(pv_done_addr); + } + } + unsigned int _phase_tmem_dealloc_0 = 0; + mbarrier_wait(tmem_dealloc_addr, _phase_tmem_dealloc_0); + _phase_tmem_dealloc_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + } + // ---- Role: load_warp ---- + if (warp >= 5 && warp <= 7) { + { // load_warp_main + const int load_warp_rank = warp - 5; + int query_idx_2 = blockIdx.x; + int _max_0 = ((sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) : (0)); + int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); + int active_topk_2 = _min_0; + int tile_count_2 = (active_topk_2 + 128 - 1) / 128; + unsigned int load_kv_stage = 0; + unsigned int load_kv_phase = 1; + int group = load_warp_rank + lane * 3; + int row0 = 0; + int row1 = 0; + int row2 = 0; + int row3 = 0; + if (group < 32) { + int index_base = group * 4; + int* index_row_1 = swa_indices + (query_idx_2 * swa_index_stride); + int raw_rows[4]; + if (index_base + 3 < sparse_topk) { + int _vec_load_0[4]; + { + const int4* _ivptr_0 = reinterpret_cast(index_row_1 + index_base + 0); + int4 _ivld_0; + _ivld_0 = *_ivptr_0; + _vec_load_0[0 + 0] = _ivld_0.x; + _vec_load_0[0 + 1] = _ivld_0.y; + _vec_load_0[0 + 2] = _ivld_0.z; + _vec_load_0[0 + 3] = _ivld_0.w; + } + #pragma unroll + for (int i_1 = 0; i_1 < 4; i_1++) { + raw_rows[i_1] = _vec_load_0[i_1]; + } + } else { + #pragma unroll + for (int i_2 = 0; i_2 < 4; i_2++) { + raw_rows[i_2] = -1; + if (index_base + i_2 < sparse_topk) { + raw_rows[i_2] = index_row_1[index_base + i_2]; + } + } + } + row0 = ((raw_rows[0] >= 0) ? raw_rows[0] : 0); + row1 = ((raw_rows[1] >= 0) ? raw_rows[1] : 0); + row2 = ((raw_rows[2] >= 0) ? raw_rows[2] : 0); + row3 = ((raw_rows[3] >= 0) ? raw_rows[3] : 0); + } + if (load_warp_rank == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(q_full_addr, 65536); + #pragma unroll + for (int q_stage = 0; q_stage < 8; q_stage += 8) { + tma_4d_gmem2smem(smem_q_addr + (unsigned int)(q_stage * 8192), (&tmap_q), 0, 0, q_stage, query_idx_2, q_full_addr); + } + } + } + #pragma unroll + for (int tile_2 = 0; tile_2 < 1; tile_2++) { + if (tile_count_2 > tile_2) { + #pragma unroll + for (int k_stage_1 = 0; k_stage_1 < 4; k_stage_1++) { + if (elect_sync()) { + int producer_groups = ((load_warp_rank < 2) ? 11 : 10); + mbarrier_arrive_expect_tx(kv_full_addr + (load_kv_stage) * 8, producer_groups * 512 * 2); + } + if (group < 32) { + int dst_k = smem_kv_addr + load_kv_stage * 32768; + if (tile_2 == 0) { + tma_gather4_gmem2smem(dst_k + group * 512, (&tmap_swa_kv), k_stage_1 * 128, row0, row1, row2, row3, kv_full_addr + (load_kv_stage) * 8); + tma_gather4_gmem2smem(dst_k + 16384 + group * 512, (&tmap_swa_kv), k_stage_1 * 128 + 64, row0, row1, row2, row3, kv_full_addr + (load_kv_stage) * 8); + } else { + tma_gather4_gmem2smem(dst_k + group * 512, (&tmap_compressed_kv), k_stage_1 * 128, row0, row1, row2, row3, kv_full_addr + (load_kv_stage) * 8); + tma_gather4_gmem2smem(dst_k + 16384 + group * 512, (&tmap_compressed_kv), k_stage_1 * 128 + 64, row0, row1, row2, row3, kv_full_addr + (load_kv_stage) * 8); + } + } + load_kv_stage += 1; + if (load_kv_stage == 4) { load_kv_stage = 0; load_kv_phase ^= 1; } + } + } + } + if (elect_sync()) { + mbarrier_arrive(tmem_dealloc_addr); + } + } + } + + // Cleanup +} + +} // extern "C" diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_0c8ab955ac7b62e7bfee_binding.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_0c8ab955ac7b62e7bfee_binding.cu new file mode 100644 index 00000000000..2f4858f771a --- /dev/null +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_0c8ab955ac7b62e7bfee_binding.cu @@ -0,0 +1,389 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +// Generated by the Cake source exporter — do not edit. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_0c8ab955ac7b62e7bfee'. +#include +#include +#include +#include + +#include "tvm_ffi_utils.h" + +#include +#include +#include +#include +#include +#include +#include + +extern "C" __global__ void kernel_cake_dsv4_0c8ab955ac7b62e7bfee(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int num_head_tiles, int has_sinks); + + +namespace cake_host_shim_900ed6f94bca9ff0 { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t CakeDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +// 4D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + int64_t d3 = t.size(t.ndim() - 3); + TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(d3)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 64u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_swa_kv' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_swa_kv' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_kv'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_O, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_num_head_tiles, int64_t arg_has_sinks, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckContiguous(arg_tmap_q, "tmap_q"); + CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); + CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 4, 16, 1); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_swa_indices, "swa_indices"); + CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); + CheckContiguous(arg_swa_indices, "swa_indices"); + CheckCudaTensor(arg_compressed_indices, "compressed_indices"); + CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); + CheckContiguous(arg_compressed_indices, "compressed_indices"); + CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); + CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); + CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); + CheckContiguous(arg_bmm1_scale, "bmm1_scale"); + CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); + CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); + CheckContiguous(arg_bmm2_scale, "bmm2_scale"); + TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) + << "scalar 'num_heads' value " << arg_num_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) + << "scalar 'swa_index_stride' value " << arg_swa_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) + << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) + << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) + << "scalar 'sparse_topk' value " << arg_sparse_topk + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_head_tiles >= -2147483648LL && arg_num_head_tiles <= 2147483647LL, ValueError) + << "scalar 'num_head_tiles' value " << arg_num_head_tiles + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) + << "scalar 'has_sinks' value " << arg_has_sinks + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); + CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); + CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); + CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + + cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); + __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); + int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); + int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); + int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); + float* p_sinks = static_cast(arg_sinks.data_ptr()); + float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); + float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); + int32_t v_num_heads = (int32_t)arg_num_heads; + int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; + int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; + int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; + int32_t v_sparse_topk = (int32_t)arg_sparse_topk; + int32_t v_num_head_tiles = (int32_t)arg_num_head_tiles; + int32_t v_has_sinks = (int32_t)arg_has_sinks; + void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_O, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_num_head_tiles, &v_has_sinks}; + + static std::mutex smem_mu; + static auto* smem_status_by_device = + new std::unordered_map(); + cudaError_t smem_status = cudaSuccess; + { + std::lock_guard lock(smem_mu); + auto it = smem_status_by_device->find(dev.device_id); + if (it == smem_status_by_device->end()) { + int current_device = -1; + smem_status = cudaGetDevice(¤t_device); + if (smem_status == cudaSuccess && current_device != dev.device_id) { + smem_status = cudaErrorInvalidDevice; + } + if (smem_status == cudaSuccess) { + smem_status = cudaFuncSetAttribute( + reinterpret_cast(kernel_cake_dsv4_0c8ab955ac7b62e7bfee), + cudaFuncAttributeMaxDynamicSharedMemorySize, + 198528); + } + smem_status_by_device->emplace(dev.device_id, smem_status); + } else { + smem_status = it->second; + } + } + TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) + << "cudaFuncSetAttribute for kernel_cake_dsv4_0c8ab955ac7b62e7bfee failed: " + << cudaGetErrorString(smem_status); + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(512u, 1u, 1u); + + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes = 198528u; + config.stream = stream; + cudaError_t launch_status = cudaLaunchKernelExC( + &config, reinterpret_cast(kernel_cake_dsv4_0c8ab955ac7b62e7bfee), kargs); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernelExC for kernel_cake_dsv4_0c8ab955ac7b62e7bfee failed: " + << cudaGetErrorString(launch_status); +} + +} // namespace cake_host_shim_900ed6f94bca9ff0 + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_900ed6f94bca9ff0::Run); diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_0c8ab955ac7b62e7bfee_kernel.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_0c8ab955ac7b62e7bfee_kernel.cu new file mode 100644 index 00000000000..c36df67ae57 --- /dev/null +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_0c8ab955ac7b62e7bfee_kernel.cu @@ -0,0 +1,1743 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +#define CAKE_INF CUDART_INF_F +#define TMEM_NCOLS 384 +#define TMEM_TMEM_OFFSET 0 +#define NUM_KV_PIPE_STAGES 4 +#define SMEM_SMEM_Q_OFF 1024 +#define SMEM_SMEM_Q_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_STRIDE 8192 +#define SMEM_SMEM_KV_OFF 66560 +#define SMEM_SMEM_KV_STAGE_BYTES 32768 +#define SMEM_SMEM_KV_STRIDE 32768 +#define SMEM_SMEM_V_OFF 66560 +#define SMEM_SMEM_V_STAGE_BYTES 32768 +#define SMEM_SMEM_V_STRIDE 32768 +#define SMEM_SMEM_SUM_OFF 197632 +#define SMEM_SMEM_SUM_STAGE_BYTES 256 +#define SMEM_SMEM_SUM_STRIDE 256 +#define SMEM_SMEM_INDEX_FLAGS_OFF 198400 +#define SMEM_SMEM_INDEX_FLAGS_STAGE_BYTES 16 +#define SMEM_SMEM_INDEX_FLAGS_STRIDE 16 +#define SMEM_SMEM_INDICES_OFF 197888 +#define SMEM_SMEM_INDICES_STAGE_BYTES 512 +#define SMEM_SMEM_INDICES_STRIDE 512 +#define SMEM_TOTAL 198528 +#define THREADS 512 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count) : "memory"); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_2d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3}], [%4];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr) { + // Canonical .shared::cta form for non-multicast gather4, matching + // trtllm-gen / cuda_ptx and the PTX ISA qualifier order + // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, + // .shared::cluster is reserved for the multicast variant (ctaMask). + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(512, 1) void +kernel_cake_dsv4_0c8ab955ac7b62e7bfee(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int num_head_tiles, int has_sinks) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); + + const int mbar_base = smem; + #define q_full_addr (mbar_base + 0) + #define kv_full_addr (mbar_base + 8) + #define kv_empty_addr (mbar_base + 40) + #define s_full_addr (mbar_base + 72) + #define p_full_addr (mbar_base + 80) + #define sum_ready_addr (mbar_base + 88) + #define o_done_addr (mbar_base + 96) + #define tmem_dealloc_addr (mbar_base + 104) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + + // Kernel setup ops + __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_q_addr = smem + 1024; + __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_kv_addr = smem + 66560; + __nv_bfloat16* smem_v = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_v_addr = smem + 66560; + float* smem_sum = reinterpret_cast(smem_raw + 197632); + const int smem_sum_addr = smem + 197632; + int* smem_index_flags = reinterpret_cast(smem_raw + 198400); + const int smem_index_flags_addr = smem + 198400; + int* smem_indices = reinterpret_cast(smem_raw + 197888); + const int smem_indices_addr = smem + 197888; + + // Mbarrier init (8 pipeline groups, 0 ordered-sequence groups, 14 barriers) + // Mbarriers at smem_raw[0..112) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // q_full: 1 barriers, init_count=1 + mbarrier_init(smem + 0, 1); + // --- pipeline 'kv_pipe' --- + // kv_full: 4 barriers, init_count=4 + mbarrier_init(smem + 8, 4); + mbarrier_init(smem + 16, 4); + mbarrier_init(smem + 24, 4); + mbarrier_init(smem + 32, 4); + // kv_empty: 4 barriers, init_count=1 + mbarrier_init(smem + 40, 1); + mbarrier_init(smem + 48, 1); + mbarrier_init(smem + 56, 1); + mbarrier_init(smem + 64, 1); + // s_full: 1 barriers, init_count=1 + mbarrier_init(smem + 72, 1); + // p_full: 1 barriers, init_count=128 + mbarrier_init(smem + 80, 128); + // sum_ready: 1 barriers, init_count=128 + mbarrier_init(smem + 88, 128); + // o_done: 1 barriers, init_count=1 + mbarrier_init(smem + 96, 1); + // tmem_dealloc: 1 barriers, init_count=256 + mbarrier_init(smem + 104, 256); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncwarp(); + + // TMEM alloc (512 columns, 384 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 112); + if (warp == 0) { + int _tmem_hold = smem + 112; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_tmem = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 64;"); + } + + // ---- Role: softmax ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); + { // softmax_main + float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; + int query_head_work = blockIdx.x >> 2; + int query_idx = query_head_work / num_head_tiles; + int head_tile = query_head_work % num_head_tiles; + int _max_0 = ((sparse_topk_lens[query_idx] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx] + sparse_topk_lens_offset) : (0)); + int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); + int active_topk = _min_0; + const int warp_in_compute = warp; + const int tmem_row_origin = warp_in_compute * 32; + const int logical_row_origin = warp_in_compute * 16; + const int my_row = logical_row_origin + lane % 16; + int head_idx = head_tile * 64 + my_row; + const int col_half = lane / 16; + unsigned int _phase_s_full_0 = 0; + mbarrier_wait(s_full_addr, _phase_s_full_0); + _phase_s_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int score_addr = taddr + (unsigned int)(tmem_row_origin << 16); + float _tmem_load_0[64]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_0[0])), "=r"(*reinterpret_cast(&_tmem_load_0[1])), "=r"(*reinterpret_cast(&_tmem_load_0[2])), "=r"(*reinterpret_cast(&_tmem_load_0[3])), "=r"(*reinterpret_cast(&_tmem_load_0[4])), "=r"(*reinterpret_cast(&_tmem_load_0[5])), "=r"(*reinterpret_cast(&_tmem_load_0[6])), "=r"(*reinterpret_cast(&_tmem_load_0[7])), "=r"(*reinterpret_cast(&_tmem_load_0[8])), "=r"(*reinterpret_cast(&_tmem_load_0[9])), "=r"(*reinterpret_cast(&_tmem_load_0[10])), "=r"(*reinterpret_cast(&_tmem_load_0[11])), "=r"(*reinterpret_cast(&_tmem_load_0[12])), "=r"(*reinterpret_cast(&_tmem_load_0[13])), "=r"(*reinterpret_cast(&_tmem_load_0[14])), "=r"(*reinterpret_cast(&_tmem_load_0[15])), "=r"(*reinterpret_cast(&_tmem_load_0[16])), "=r"(*reinterpret_cast(&_tmem_load_0[17])), "=r"(*reinterpret_cast(&_tmem_load_0[18])), "=r"(*reinterpret_cast(&_tmem_load_0[19])), "=r"(*reinterpret_cast(&_tmem_load_0[20])), "=r"(*reinterpret_cast(&_tmem_load_0[21])), "=r"(*reinterpret_cast(&_tmem_load_0[22])), "=r"(*reinterpret_cast(&_tmem_load_0[23])), "=r"(*reinterpret_cast(&_tmem_load_0[24])), "=r"(*reinterpret_cast(&_tmem_load_0[25])), "=r"(*reinterpret_cast(&_tmem_load_0[26])), "=r"(*reinterpret_cast(&_tmem_load_0[27])), "=r"(*reinterpret_cast(&_tmem_load_0[28])), "=r"(*reinterpret_cast(&_tmem_load_0[29])), "=r"(*reinterpret_cast(&_tmem_load_0[30])), "=r"(*reinterpret_cast(&_tmem_load_0[31])) + : "r"(score_addr)); + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_0[32])), "=r"(*reinterpret_cast(&_tmem_load_0[33])), "=r"(*reinterpret_cast(&_tmem_load_0[34])), "=r"(*reinterpret_cast(&_tmem_load_0[35])), "=r"(*reinterpret_cast(&_tmem_load_0[36])), "=r"(*reinterpret_cast(&_tmem_load_0[37])), "=r"(*reinterpret_cast(&_tmem_load_0[38])), "=r"(*reinterpret_cast(&_tmem_load_0[39])), "=r"(*reinterpret_cast(&_tmem_load_0[40])), "=r"(*reinterpret_cast(&_tmem_load_0[41])), "=r"(*reinterpret_cast(&_tmem_load_0[42])), "=r"(*reinterpret_cast(&_tmem_load_0[43])), "=r"(*reinterpret_cast(&_tmem_load_0[44])), "=r"(*reinterpret_cast(&_tmem_load_0[45])), "=r"(*reinterpret_cast(&_tmem_load_0[46])), "=r"(*reinterpret_cast(&_tmem_load_0[47])), "=r"(*reinterpret_cast(&_tmem_load_0[48])), "=r"(*reinterpret_cast(&_tmem_load_0[49])), "=r"(*reinterpret_cast(&_tmem_load_0[50])), "=r"(*reinterpret_cast(&_tmem_load_0[51])), "=r"(*reinterpret_cast(&_tmem_load_0[52])), "=r"(*reinterpret_cast(&_tmem_load_0[53])), "=r"(*reinterpret_cast(&_tmem_load_0[54])), "=r"(*reinterpret_cast(&_tmem_load_0[55])), "=r"(*reinterpret_cast(&_tmem_load_0[56])), "=r"(*reinterpret_cast(&_tmem_load_0[57])), "=r"(*reinterpret_cast(&_tmem_load_0[58])), "=r"(*reinterpret_cast(&_tmem_load_0[59])), "=r"(*reinterpret_cast(&_tmem_load_0[60])), "=r"(*reinterpret_cast(&_tmem_load_0[61])), "=r"(*reinterpret_cast(&_tmem_load_0[62])), "=r"(*reinterpret_cast(&_tmem_load_0[63])) + : "r"(score_addr + 32)); + int valid_cols = active_topk - col_half * 64; + int causal_cols = smem_index_flags[col_half * 2 + 1]; + valid_cols = ((valid_cols < causal_cols) ? valid_cols : causal_cols); + uint32_t _slice_lo_mask_0; + { + int _lim_0 = valid_cols; + if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_0 & (1u << 0))) _tmem_load_0[0] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 1))) _tmem_load_0[1] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 2))) _tmem_load_0[2] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 3))) _tmem_load_0[3] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 4))) _tmem_load_0[4] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 5))) _tmem_load_0[5] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 6))) _tmem_load_0[6] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 7))) _tmem_load_0[7] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 8))) _tmem_load_0[8] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 9))) _tmem_load_0[9] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 10))) _tmem_load_0[10] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 11))) _tmem_load_0[11] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 12))) _tmem_load_0[12] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 13))) _tmem_load_0[13] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 14))) _tmem_load_0[14] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 15))) _tmem_load_0[15] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 16))) _tmem_load_0[16] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 17))) _tmem_load_0[17] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 18))) _tmem_load_0[18] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 19))) _tmem_load_0[19] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 20))) _tmem_load_0[20] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 21))) _tmem_load_0[21] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 22))) _tmem_load_0[22] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 23))) _tmem_load_0[23] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 24))) _tmem_load_0[24] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 25))) _tmem_load_0[25] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 26))) _tmem_load_0[26] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 27))) _tmem_load_0[27] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 28))) _tmem_load_0[28] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 29))) _tmem_load_0[29] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 30))) _tmem_load_0[30] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 31))) _tmem_load_0[31] = -CAKE_INF; + uint32_t _slice_lo_mask_1; + { + int _lim_1 = valid_cols - 32; + if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) _tmem_load_0[32] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) _tmem_load_0[33] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) _tmem_load_0[34] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) _tmem_load_0[35] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) _tmem_load_0[36] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) _tmem_load_0[37] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) _tmem_load_0[38] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) _tmem_load_0[39] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) _tmem_load_0[40] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) _tmem_load_0[41] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) _tmem_load_0[42] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) _tmem_load_0[43] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) _tmem_load_0[44] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) _tmem_load_0[45] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) _tmem_load_0[46] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) _tmem_load_0[47] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) _tmem_load_0[48] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) _tmem_load_0[49] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) _tmem_load_0[50] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) _tmem_load_0[51] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) _tmem_load_0[52] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) _tmem_load_0[53] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) _tmem_load_0[54] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) _tmem_load_0[55] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) _tmem_load_0[56] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) _tmem_load_0[57] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) _tmem_load_0[58] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) _tmem_load_0[59] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) _tmem_load_0[60] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) _tmem_load_0[61] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) _tmem_load_0[62] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) _tmem_load_0[63] = -CAKE_INF; + float2 _reg_reduce_max2_2 = {-CAKE_INF, -CAKE_INF}; + row_max_x32_accum(&_tmem_load_0[0], _reg_reduce_max2_2); + row_max_x32_accum(&_tmem_load_0[32], _reg_reduce_max2_2); + float _tmem_load_0_max = row_max_reduce(_reg_reduce_max2_2); + float row_max = _tmem_load_0_max; + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, row_max, 16); + float _max_1 = max_noftz(row_max, _shfl_xor_0); + row_max = _max_1; + if (has_sinks != 0 && head_idx < num_heads) { + float sink_unscaled = sinks[head_idx] * 1.4426950408889634f / softmax_scale_log2; + float _max_2 = max_noftz(row_max, sink_unscaled); + row_max = _max_2; + } + bool empty_row = row_max == -CAKE_INF; + float safe_max = ((empty_row) ? 0.0f : row_max); + float max_scaled = safe_max * softmax_scale_log2; + float2 _f2_0 = make_float2(softmax_scale_log2, softmax_scale_log2); + float2 _f2_1 = make_float2(-max_scaled, -max_scaled); + float2 _f2_2 = make_float2(0.0f, 0.0f); + float2 sum0 = _f2_2; + float2 _f2_3 = make_float2(0.0f, 0.0f); + float2 sum1 = _f2_3; + float2 _f2_4 = make_float2(0.0f, 0.0f); + float2 sum2 = _f2_4; + float2 _f2_5 = make_float2(0.0f, 0.0f); + float2 sum3 = _f2_5; + #pragma unroll + for (int i = 0; i < 64; i += 8) { + float2 _f2_6 = make_float2(_tmem_load_0[i], _tmem_load_0[i + 1]); + float2 _f2_7 = make_float2(_tmem_load_0[i + 2], _tmem_load_0[i + 3]); + float2 _f2_8 = make_float2(_tmem_load_0[i + 4], _tmem_load_0[i + 5]); + float2 _f2_9 = make_float2(_tmem_load_0[i + 6], _tmem_load_0[i + 7]); + float2 affine0 = fma_f32x2_rn_ftz(_f2_6, _f2_0, _f2_1); + float2 affine1 = fma_f32x2_rn_ftz(_f2_7, _f2_0, _f2_1); + float2 affine2 = fma_f32x2_rn_ftz(_f2_8, _f2_0, _f2_1); + float2 affine3 = fma_f32x2_rn_ftz(_f2_9, _f2_0, _f2_1); + float _exp2_0 = approx_exp2(affine0.x); + float exp0 = _exp2_0; + float _exp2_1 = approx_exp2(affine0.y); + float exp1 = _exp2_1; + float _exp2_2 = approx_exp2(affine1.x); + float exp2 = _exp2_2; + float _exp2_3 = approx_exp2(affine1.y); + float exp3 = _exp2_3; + float _exp2_4 = approx_exp2(affine2.x); + float exp4 = _exp2_4; + float _exp2_5 = approx_exp2(affine2.y); + float exp5 = _exp2_5; + float _exp2_6 = approx_exp2(affine3.x); + float exp6 = _exp2_6; + float _exp2_7 = approx_exp2(affine3.y); + float exp7 = _exp2_7; + _tmem_load_0[i] = exp0; + _tmem_load_0[i + 1] = exp1; + _tmem_load_0[i + 2] = exp2; + _tmem_load_0[i + 3] = exp3; + _tmem_load_0[i + 4] = exp4; + _tmem_load_0[i + 5] = exp5; + _tmem_load_0[i + 6] = exp6; + _tmem_load_0[i + 7] = exp7; + float2 _f2_10 = make_float2(exp0, exp1); + sum0 = add_f32x2(sum0, _f2_10); + float2 _f2_11 = make_float2(exp2, exp3); + sum1 = add_f32x2(sum1, _f2_11); + float2 _f2_12 = make_float2(exp4, exp5); + sum2 = add_f32x2(sum2, _f2_12); + float2 _f2_13 = make_float2(exp6, exp7); + sum3 = add_f32x2(sum3, _f2_13); + } + float2 sum01 = add_f32x2(sum0, sum1); + float2 sum23 = add_f32x2(sum2, sum3); + float2 row_sum_pair = add_f32x2(sum01, sum23); + float row_sum = row_sum_pair.x + row_sum_pair.y; + float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, row_sum, 16); + row_sum = row_sum + _shfl_xor_1; + if (has_sinks != 0 && head_idx < num_heads) { + float _exp2_8 = approx_exp2(sinks[head_idx] * 1.4426950408889634f - max_scaled); + row_sum = row_sum + _exp2_8; + } + unsigned int packed_p[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_0[_lp*2 + 0], _tmem_load_0[_lp*2+1 + 0])); + packed_p[_lp] = *(uint32_t*)&_bf2; + } + int p_addr = taddr + 128 + (unsigned int)(tmem_row_origin << 16); + asm volatile( + "tcgen05.st.sync.aligned.16x32bx2.x16.b32" + " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(p_addr), "r"(*reinterpret_cast(&packed_p[0])), "r"(*reinterpret_cast(&packed_p[1])), "r"(*reinterpret_cast(&packed_p[2])), "r"(*reinterpret_cast(&packed_p[3])), "r"(*reinterpret_cast(&packed_p[4])), "r"(*reinterpret_cast(&packed_p[5])), "r"(*reinterpret_cast(&packed_p[6])), "r"(*reinterpret_cast(&packed_p[7])), "r"(*reinterpret_cast(&packed_p[8])), "r"(*reinterpret_cast(&packed_p[9])), "r"(*reinterpret_cast(&packed_p[10])), "r"(*reinterpret_cast(&packed_p[11])), "r"(*reinterpret_cast(&packed_p[12])), "r"(*reinterpret_cast(&packed_p[13])), "r"(*reinterpret_cast(&packed_p[14])), "r"(*reinterpret_cast(&packed_p[15]))); + asm volatile( + "tcgen05.st.sync.aligned.16x32bx2.x16.b32" + " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(p_addr + 16), "r"(*reinterpret_cast(&(packed_p + 16)[0])), "r"(*reinterpret_cast(&(packed_p + 16)[1])), "r"(*reinterpret_cast(&(packed_p + 16)[2])), "r"(*reinterpret_cast(&(packed_p + 16)[3])), "r"(*reinterpret_cast(&(packed_p + 16)[4])), "r"(*reinterpret_cast(&(packed_p + 16)[5])), "r"(*reinterpret_cast(&(packed_p + 16)[6])), "r"(*reinterpret_cast(&(packed_p + 16)[7])), "r"(*reinterpret_cast(&(packed_p + 16)[8])), "r"(*reinterpret_cast(&(packed_p + 16)[9])), "r"(*reinterpret_cast(&(packed_p + 16)[10])), "r"(*reinterpret_cast(&(packed_p + 16)[11])), "r"(*reinterpret_cast(&(packed_p + 16)[12])), "r"(*reinterpret_cast(&(packed_p + 16)[13])), "r"(*reinterpret_cast(&(packed_p + 16)[14])), "r"(*reinterpret_cast(&(packed_p + 16)[15]))); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + if (col_half == 0) { + float _max_3 = max_noftz(row_sum, 1.1754943508222875e-38f); + float published_sum = _max_3; + smem_sum[my_row] = published_sum; + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(p_full_addr); + mbarrier_arrive(sum_ready_addr); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_dealloc_addr); + } + } + // ---- Role: epilogue ---- + if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 128;"); + { // epilogue_main + float output_scale = bmm2_scale[0]; + int work_idx = blockIdx.x; + int query_head_work_1 = work_idx >> 2; + int query_idx_1 = query_head_work_1 / num_head_tiles; + int head_tile_1 = query_head_work_1 % num_head_tiles; + int v_chunk = work_idx & 3; + const int warp_in_role = warp - 4; + const int tmem_row_origin_1 = warp_in_role * 32; + const int logical_row_origin_1 = warp_in_role * 16; + const int my_row_1 = logical_row_origin_1 + lane % 16; + int head_idx_1 = head_tile_1 * 64 + my_row_1; + const int col_half_1 = lane / 16; + const int row_addr = tmem_row_origin_1 << 16; + unsigned int _phase_o_done_0 = 0; + mbarrier_wait(o_done_addr, _phase_o_done_0); + _phase_o_done_0 ^= 1; + unsigned int _phase_sum_ready_0 = 0; + mbarrier_wait(sum_ready_addr, _phase_sum_ready_0); + _phase_sum_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + float row_sum_1 = smem_sum[my_row_1]; + float _rcp_0 = approx_rcp(row_sum_1); + float inv_sum = _rcp_0; + int output_base = (query_idx_1 * num_heads + head_idx_1) * 512 + v_chunk * 128; + float _tmem_load_1[64]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_1[0])), "=r"(*reinterpret_cast(&_tmem_load_1[1])), "=r"(*reinterpret_cast(&_tmem_load_1[2])), "=r"(*reinterpret_cast(&_tmem_load_1[3])), "=r"(*reinterpret_cast(&_tmem_load_1[4])), "=r"(*reinterpret_cast(&_tmem_load_1[5])), "=r"(*reinterpret_cast(&_tmem_load_1[6])), "=r"(*reinterpret_cast(&_tmem_load_1[7])), "=r"(*reinterpret_cast(&_tmem_load_1[8])), "=r"(*reinterpret_cast(&_tmem_load_1[9])), "=r"(*reinterpret_cast(&_tmem_load_1[10])), "=r"(*reinterpret_cast(&_tmem_load_1[11])), "=r"(*reinterpret_cast(&_tmem_load_1[12])), "=r"(*reinterpret_cast(&_tmem_load_1[13])), "=r"(*reinterpret_cast(&_tmem_load_1[14])), "=r"(*reinterpret_cast(&_tmem_load_1[15])), "=r"(*reinterpret_cast(&_tmem_load_1[16])), "=r"(*reinterpret_cast(&_tmem_load_1[17])), "=r"(*reinterpret_cast(&_tmem_load_1[18])), "=r"(*reinterpret_cast(&_tmem_load_1[19])), "=r"(*reinterpret_cast(&_tmem_load_1[20])), "=r"(*reinterpret_cast(&_tmem_load_1[21])), "=r"(*reinterpret_cast(&_tmem_load_1[22])), "=r"(*reinterpret_cast(&_tmem_load_1[23])), "=r"(*reinterpret_cast(&_tmem_load_1[24])), "=r"(*reinterpret_cast(&_tmem_load_1[25])), "=r"(*reinterpret_cast(&_tmem_load_1[26])), "=r"(*reinterpret_cast(&_tmem_load_1[27])), "=r"(*reinterpret_cast(&_tmem_load_1[28])), "=r"(*reinterpret_cast(&_tmem_load_1[29])), "=r"(*reinterpret_cast(&_tmem_load_1[30])), "=r"(*reinterpret_cast(&_tmem_load_1[31])) + : "r"(taddr + 256 + (unsigned int)row_addr)); + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_1[32])), "=r"(*reinterpret_cast(&_tmem_load_1[33])), "=r"(*reinterpret_cast(&_tmem_load_1[34])), "=r"(*reinterpret_cast(&_tmem_load_1[35])), "=r"(*reinterpret_cast(&_tmem_load_1[36])), "=r"(*reinterpret_cast(&_tmem_load_1[37])), "=r"(*reinterpret_cast(&_tmem_load_1[38])), "=r"(*reinterpret_cast(&_tmem_load_1[39])), "=r"(*reinterpret_cast(&_tmem_load_1[40])), "=r"(*reinterpret_cast(&_tmem_load_1[41])), "=r"(*reinterpret_cast(&_tmem_load_1[42])), "=r"(*reinterpret_cast(&_tmem_load_1[43])), "=r"(*reinterpret_cast(&_tmem_load_1[44])), "=r"(*reinterpret_cast(&_tmem_load_1[45])), "=r"(*reinterpret_cast(&_tmem_load_1[46])), "=r"(*reinterpret_cast(&_tmem_load_1[47])), "=r"(*reinterpret_cast(&_tmem_load_1[48])), "=r"(*reinterpret_cast(&_tmem_load_1[49])), "=r"(*reinterpret_cast(&_tmem_load_1[50])), "=r"(*reinterpret_cast(&_tmem_load_1[51])), "=r"(*reinterpret_cast(&_tmem_load_1[52])), "=r"(*reinterpret_cast(&_tmem_load_1[53])), "=r"(*reinterpret_cast(&_tmem_load_1[54])), "=r"(*reinterpret_cast(&_tmem_load_1[55])), "=r"(*reinterpret_cast(&_tmem_load_1[56])), "=r"(*reinterpret_cast(&_tmem_load_1[57])), "=r"(*reinterpret_cast(&_tmem_load_1[58])), "=r"(*reinterpret_cast(&_tmem_load_1[59])), "=r"(*reinterpret_cast(&_tmem_load_1[60])), "=r"(*reinterpret_cast(&_tmem_load_1[61])), "=r"(*reinterpret_cast(&_tmem_load_1[62])), "=r"(*reinterpret_cast(&_tmem_load_1[63])) + : "r"(taddr + 256 + (unsigned int)row_addr + 32)); + float2 _f2_14 = make_float2(inv_sum * output_scale, inv_sum * output_scale); + #pragma unroll + for (int i_1 = 0; i_1 < 64; i_1 += 2) { + float2 _f2_15 = make_float2(_tmem_load_1[i_1], _tmem_load_1[i_1 + 1]); + float2 _mul_f32x2_0; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_0) : "l"(*(const unsigned long long*)&_f2_15), "l"(*(const unsigned long long*)&_f2_14)); + _tmem_load_1[i_1] = _mul_f32x2_0.x; + _tmem_load_1[i_1 + 1] = _mul_f32x2_0.y; + } + unsigned int packed_o[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_1[_lp*2 + 0], _tmem_load_1[_lp*2+1 + 0])); + packed_o[_lp] = *(uint32_t*)&_bf2; + } + bool stage_hi_1 = (lane & 1) != 0; + bool stage_hi_2 = (lane & 2) != 0; + bool stage_hi_4 = (lane & 4) != 0; + unsigned int _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[0] : packed_o[4]), 1); + packed_o[0] = ((stage_hi_1) ? _shfl_xor_2 : packed_o[0]); + packed_o[4] = ((stage_hi_1) ? packed_o[4] : _shfl_xor_2); + unsigned int _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[1] : packed_o[5]), 1); + packed_o[1] = ((stage_hi_1) ? _shfl_xor_3 : packed_o[1]); + packed_o[5] = ((stage_hi_1) ? packed_o[5] : _shfl_xor_3); + unsigned int _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[2] : packed_o[6]), 1); + packed_o[2] = ((stage_hi_1) ? _shfl_xor_4 : packed_o[2]); + packed_o[6] = ((stage_hi_1) ? packed_o[6] : _shfl_xor_4); + unsigned int _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[3] : packed_o[7]), 1); + packed_o[3] = ((stage_hi_1) ? _shfl_xor_5 : packed_o[3]); + packed_o[7] = ((stage_hi_1) ? packed_o[7] : _shfl_xor_5); + unsigned int _shfl_xor_6 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[8] : packed_o[12]), 1); + packed_o[8] = ((stage_hi_1) ? _shfl_xor_6 : packed_o[8]); + packed_o[12] = ((stage_hi_1) ? packed_o[12] : _shfl_xor_6); + unsigned int _shfl_xor_7 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[9] : packed_o[13]), 1); + packed_o[9] = ((stage_hi_1) ? _shfl_xor_7 : packed_o[9]); + packed_o[13] = ((stage_hi_1) ? packed_o[13] : _shfl_xor_7); + unsigned int _shfl_xor_8 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[10] : packed_o[14]), 1); + packed_o[10] = ((stage_hi_1) ? _shfl_xor_8 : packed_o[10]); + packed_o[14] = ((stage_hi_1) ? packed_o[14] : _shfl_xor_8); + unsigned int _shfl_xor_9 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[11] : packed_o[15]), 1); + packed_o[11] = ((stage_hi_1) ? _shfl_xor_9 : packed_o[11]); + packed_o[15] = ((stage_hi_1) ? packed_o[15] : _shfl_xor_9); + unsigned int _shfl_xor_10 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[16] : packed_o[20]), 1); + packed_o[16] = ((stage_hi_1) ? _shfl_xor_10 : packed_o[16]); + packed_o[20] = ((stage_hi_1) ? packed_o[20] : _shfl_xor_10); + unsigned int _shfl_xor_11 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[17] : packed_o[21]), 1); + packed_o[17] = ((stage_hi_1) ? _shfl_xor_11 : packed_o[17]); + packed_o[21] = ((stage_hi_1) ? packed_o[21] : _shfl_xor_11); + unsigned int _shfl_xor_12 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[18] : packed_o[22]), 1); + packed_o[18] = ((stage_hi_1) ? _shfl_xor_12 : packed_o[18]); + packed_o[22] = ((stage_hi_1) ? packed_o[22] : _shfl_xor_12); + unsigned int _shfl_xor_13 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[19] : packed_o[23]), 1); + packed_o[19] = ((stage_hi_1) ? _shfl_xor_13 : packed_o[19]); + packed_o[23] = ((stage_hi_1) ? packed_o[23] : _shfl_xor_13); + unsigned int _shfl_xor_14 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[24] : packed_o[28]), 1); + packed_o[24] = ((stage_hi_1) ? _shfl_xor_14 : packed_o[24]); + packed_o[28] = ((stage_hi_1) ? packed_o[28] : _shfl_xor_14); + unsigned int _shfl_xor_15 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[25] : packed_o[29]), 1); + packed_o[25] = ((stage_hi_1) ? _shfl_xor_15 : packed_o[25]); + packed_o[29] = ((stage_hi_1) ? packed_o[29] : _shfl_xor_15); + unsigned int _shfl_xor_16 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[26] : packed_o[30]), 1); + packed_o[26] = ((stage_hi_1) ? _shfl_xor_16 : packed_o[26]); + packed_o[30] = ((stage_hi_1) ? packed_o[30] : _shfl_xor_16); + unsigned int _shfl_xor_17 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[27] : packed_o[31]), 1); + packed_o[27] = ((stage_hi_1) ? _shfl_xor_17 : packed_o[27]); + packed_o[31] = ((stage_hi_1) ? packed_o[31] : _shfl_xor_17); + int lane_t = lane & 1; + unsigned int chunk_words[4]; + chunk_words[0] = packed_o[0]; + chunk_words[1] = packed_o[1]; + chunk_words[2] = packed_o[2]; + chunk_words[3] = packed_o[3]; + if (head_tile_1 * 64 + (my_row_1 & -2 | 0) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + head_tile_1 * 64 + (my_row_1 & -2 | 0)) * 512 + v_chunk * 128 + col_half_1 * 64 + (0 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[4]; + chunk_words[1] = packed_o[5]; + chunk_words[2] = packed_o[6]; + chunk_words[3] = packed_o[7]; + if (head_tile_1 * 64 + (my_row_1 & -2 | 1) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + head_tile_1 * 64 + (my_row_1 & -2 | 1)) * 512 + v_chunk * 128 + col_half_1 * 64 + (0 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[8]; + chunk_words[1] = packed_o[9]; + chunk_words[2] = packed_o[10]; + chunk_words[3] = packed_o[11]; + if (head_tile_1 * 64 + (my_row_1 & -2 | 0) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + head_tile_1 * 64 + (my_row_1 & -2 | 0)) * 512 + v_chunk * 128 + col_half_1 * 64 + (2 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[12]; + chunk_words[1] = packed_o[13]; + chunk_words[2] = packed_o[14]; + chunk_words[3] = packed_o[15]; + if (head_tile_1 * 64 + (my_row_1 & -2 | 1) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + head_tile_1 * 64 + (my_row_1 & -2 | 1)) * 512 + v_chunk * 128 + col_half_1 * 64 + (2 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[16]; + chunk_words[1] = packed_o[17]; + chunk_words[2] = packed_o[18]; + chunk_words[3] = packed_o[19]; + if (head_tile_1 * 64 + (my_row_1 & -2 | 0) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + head_tile_1 * 64 + (my_row_1 & -2 | 0)) * 512 + v_chunk * 128 + col_half_1 * 64 + (4 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[20]; + chunk_words[1] = packed_o[21]; + chunk_words[2] = packed_o[22]; + chunk_words[3] = packed_o[23]; + if (head_tile_1 * 64 + (my_row_1 & -2 | 1) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + head_tile_1 * 64 + (my_row_1 & -2 | 1)) * 512 + v_chunk * 128 + col_half_1 * 64 + (4 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[24]; + chunk_words[1] = packed_o[25]; + chunk_words[2] = packed_o[26]; + chunk_words[3] = packed_o[27]; + if (head_tile_1 * 64 + (my_row_1 & -2 | 0) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + head_tile_1 * 64 + (my_row_1 & -2 | 0)) * 512 + v_chunk * 128 + col_half_1 * 64 + (6 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[28]; + chunk_words[1] = packed_o[29]; + chunk_words[2] = packed_o[30]; + chunk_words[3] = packed_o[31]; + if (head_tile_1 * 64 + (my_row_1 & -2 | 1) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + head_tile_1 * 64 + (my_row_1 & -2 | 1)) * 512 + v_chunk * 128 + col_half_1 * 64 + (6 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_dealloc_addr); + } + } + // ---- Role: mma_warp ---- + if (warp == 8) { + { // mma_warp_main + unsigned int _phase_q_full_0 = 0; + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + unsigned int _phase_kv_full = 0; + #pragma unroll + for (int k_stage = 0; k_stage < 4; k_stage++) { + mbarrier_wait(kv_full_addr + (k_stage) * 8, _phase_kv_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_0 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (k_stage * 2) * 512); + int _mma_b_lo_0 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (k_stage) * 2048); + { + uint64_t _mma_ss_a_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_0); + uint64_t _mma_ss_b_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_0); + if (elect_sync()) { + tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, ((k_stage == 0) ? 0 : 1)); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); + } + } + int _mma_a_lo_1 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (k_stage * 2 + 1) * 512); + int _mma_b_lo_1 = make_warp_uniform((((smem_kv_addr + 16384) >> 4) & 0x3FFF) + (k_stage) * 2048); + { + uint64_t _mma_ss_a_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_1); + uint64_t _mma_ss_b_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_1); + if (elect_sync()) { + tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); + } + } + if (k_stage == 3) { + elect_commit(s_full_addr); + } + elect_commit(kv_empty_addr + (k_stage) * 8); + } + int _mma_prepared_b_lo_2 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (0) * 2048); + unsigned int _phase_p_full_0 = 0; + mbarrier_wait(p_full_addr, _phase_p_full_0); + _phase_p_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + mbarrier_wait(kv_full_addr, 1); + asm volatile("tcgen05.fence::after_thread_sync;"); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 dhi, blo, ta, id;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 69272720;\n\t" + "mov.b32 ta, %2;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p0;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "}\n" + :: "r"((tmem_tmem + (256))), "r"(_mma_prepared_b_lo_2), "r"(tmem_tmem + 128), "r"(0)); + elect_commit(o_done_addr); + elect_commit(kv_empty_addr); + unsigned int _phase_tmem_dealloc_0 = 0; + mbarrier_wait(tmem_dealloc_addr, _phase_tmem_dealloc_0); + _phase_tmem_dealloc_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + } + // ---- Role: empty ---- + if (warp >= 9 && warp <= 11) { + // idle — no tasks assigned + } + // ---- Role: load_warp ---- + if (warp >= 12 && warp <= 15) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 64;"); + { // load_warp_main + const int load_warp_rank = warp - 12; + int work_idx_1 = blockIdx.x; + int query_head_work_2 = work_idx_1 >> 2; + int query_idx_2 = query_head_work_2 / num_head_tiles; + int head_tile_2 = query_head_work_2 % num_head_tiles; + int v_chunk_1 = work_idx_1 & 3; + int* row_ptr = swa_indices + (query_idx_2 * swa_index_stride); + int index_offset = load_warp_rank * 32 + lane; + int sparse_row = -1; + if (index_offset < sparse_topk) { + sparse_row = row_ptr[index_offset]; + } + if (load_warp_rank == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(q_full_addr, 65536); + #pragma unroll + for (int q_stage = 0; q_stage < 8; q_stage++) { + tma_4d_gmem2smem(smem_q_addr + (unsigned int)(q_stage * 8192), (&tmap_q), 0, head_tile_2 * 64, q_stage, query_idx_2, q_full_addr); + } + } + } + smem_indices[index_offset] = sparse_row; + unsigned int _vote_0 = __ballot_sync(0xFFFFFFFF, sparse_row < 0); + unsigned int invalid_indices = _vote_0; + int valid_prefix = 32; + if (invalid_indices != 0) { + int _ffs_0 = __ffs(invalid_indices); + valid_prefix = _ffs_0 - 1; + } + if (elect_sync()) { + smem_index_flags[load_warp_rank] = valid_prefix; + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("barrier.sync 8, 128;" ::: "memory"); + if ((load_warp_rank & 1) != 0) { + if (elect_sync()) { + int lower_prefix = smem_index_flags[load_warp_rank - 1]; + int upper_prefix = smem_index_flags[load_warp_rank]; + smem_index_flags[load_warp_rank] = ((lower_prefix < 32) ? lower_prefix : 32 + upper_prefix); + } + } + unsigned int _phase_kv_empty = 1; + #pragma unroll + for (int k_stage_1 = 0; k_stage_1 < 4; k_stage_1++) { + mbarrier_wait(kv_empty_addr + (k_stage_1) * 8, _phase_kv_empty); + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (k_stage_1) * 8, 8192); + } + #pragma unroll + for (int local_group = 0; local_group < 8; local_group++) { + int dst_k = smem_kv_addr + (unsigned int)(k_stage_1 * 32768); + int group = load_warp_rank * 8 + local_group; + int group_offset = group * 4; + int raw_rows[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows[0])), "=r"(*reinterpret_cast(&raw_rows[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows[(0) + 3])) + : "r"(smem_indices_addr + (unsigned int)(group_offset * 4))); + int raw0 = raw_rows[0]; + int raw1 = raw_rows[1]; + int raw2 = raw_rows[2]; + int raw3 = raw_rows[3]; + int row0 = ((raw0 >= 0) ? raw0 : 0); + int row1 = ((raw1 >= 0) ? raw1 : 0); + int row2 = ((raw2 >= 0) ? raw2 : 0); + int row3 = ((raw3 >= 0) ? raw3 : 0); + if (elect_sync()) { + tma_gather4_gmem2smem(dst_k + group * 512, (&tmap_swa_kv), k_stage_1 * 128, row0, row1, row2, row3, kv_full_addr + (k_stage_1) * 8); + tma_gather4_gmem2smem(dst_k + 16384 + group * 512, (&tmap_swa_kv), k_stage_1 * 128 + 64, row0, row1, row2, row3, kv_full_addr + (k_stage_1) * 8); + } + } + } + #pragma unroll + for (int v_stage = 0; v_stage < 1; v_stage++) { + mbarrier_wait(kv_empty_addr + (v_stage) * 8, 0); + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (v_stage) * 8, 8192); + } + #pragma unroll + for (int local_group_1 = 0; local_group_1 < 8; local_group_1++) { + int dst_v = smem_v_addr + (unsigned int)(v_stage * 32768); + int v_col = v_chunk_1 * 128 + v_stage * 64; + int group_1 = load_warp_rank * 8 + local_group_1; + int group_offset_1 = group_1 * 4; + int raw_rows_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_1[0])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 3])) + : "r"(smem_indices_addr + (unsigned int)(group_offset_1 * 4))); + int raw0_1 = raw_rows_1[0]; + int raw1_1 = raw_rows_1[1]; + int raw2_1 = raw_rows_1[2]; + int raw3_1 = raw_rows_1[3]; + int row0_1 = ((raw0_1 >= 0) ? raw0_1 : 0); + int row1_1 = ((raw1_1 >= 0) ? raw1_1 : 0); + int row2_1 = ((raw2_1 >= 0) ? raw2_1 : 0); + int row3_1 = ((raw3_1 >= 0) ? raw3_1 : 0); + if (elect_sync()) { + tma_gather4_gmem2smem(dst_v + group_1 * 512, (&tmap_swa_kv), v_col, row0_1, row1_1, row2_1, row3_1, kv_full_addr + (v_stage) * 8); + tma_gather4_gmem2smem(dst_v + 16384 + group_1 * 512, (&tmap_swa_kv), v_col + 64, row0_1, row1_1, row2_1, row3_1, kv_full_addr + (v_stage) * 8); + } + } + } + } + } + + // Cleanup +} + +} // extern "C" diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_12ac507a3c73f2784621_binding.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_12ac507a3c73f2784621_binding.cu new file mode 100644 index 00000000000..579d030d38f --- /dev/null +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_12ac507a3c73f2784621_binding.cu @@ -0,0 +1,385 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +// Generated by the Cake source exporter — do not edit. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_12ac507a3c73f2784621'. +#include +#include +#include +#include + +#include "tvm_ffi_utils.h" + +#include +#include +#include +#include +#include +#include +#include + +extern "C" __global__ void kernel_cake_dsv4_12ac507a3c73f2784621(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks); + + +namespace cake_host_shim_96d072de9e3f7f61 { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t CakeDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +// 4D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + int64_t d3 = t.size(t.ndim() - 3); + TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(d3)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 64u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_swa_kv' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_swa_kv' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_kv'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_O, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckContiguous(arg_tmap_q, "tmap_q"); + CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); + CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 4, 16, 1); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_swa_indices, "swa_indices"); + CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); + CheckContiguous(arg_swa_indices, "swa_indices"); + CheckCudaTensor(arg_compressed_indices, "compressed_indices"); + CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); + CheckContiguous(arg_compressed_indices, "compressed_indices"); + CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); + CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); + CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); + CheckContiguous(arg_bmm1_scale, "bmm1_scale"); + CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); + CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); + CheckContiguous(arg_bmm2_scale, "bmm2_scale"); + TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) + << "scalar 'num_heads' value " << arg_num_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) + << "scalar 'swa_index_stride' value " << arg_swa_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) + << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) + << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) + << "scalar 'sparse_topk' value " << arg_sparse_topk + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) + << "scalar 'has_sinks' value " << arg_has_sinks + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); + CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); + CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); + CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + + cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); + __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); + int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); + int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); + int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); + float* p_sinks = static_cast(arg_sinks.data_ptr()); + float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); + float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); + int32_t v_num_heads = (int32_t)arg_num_heads; + int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; + int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; + int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; + int32_t v_sparse_topk = (int32_t)arg_sparse_topk; + int32_t v_has_sinks = (int32_t)arg_has_sinks; + void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_O, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_has_sinks}; + + static std::mutex smem_mu; + static auto* smem_status_by_device = + new std::unordered_map(); + cudaError_t smem_status = cudaSuccess; + { + std::lock_guard lock(smem_mu); + auto it = smem_status_by_device->find(dev.device_id); + if (it == smem_status_by_device->end()) { + int current_device = -1; + smem_status = cudaGetDevice(¤t_device); + if (smem_status == cudaSuccess && current_device != dev.device_id) { + smem_status = cudaErrorInvalidDevice; + } + if (smem_status == cudaSuccess) { + smem_status = cudaFuncSetAttribute( + reinterpret_cast(kernel_cake_dsv4_12ac507a3c73f2784621), + cudaFuncAttributeMaxDynamicSharedMemorySize, + 198528); + } + smem_status_by_device->emplace(dev.device_id, smem_status); + } else { + smem_status = it->second; + } + } + TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) + << "cudaFuncSetAttribute for kernel_cake_dsv4_12ac507a3c73f2784621 failed: " + << cudaGetErrorString(smem_status); + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(512u, 1u, 1u); + + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes = 198528u; + config.stream = stream; + cudaError_t launch_status = cudaLaunchKernelExC( + &config, reinterpret_cast(kernel_cake_dsv4_12ac507a3c73f2784621), kargs); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernelExC for kernel_cake_dsv4_12ac507a3c73f2784621 failed: " + << cudaGetErrorString(launch_status); +} + +} // namespace cake_host_shim_96d072de9e3f7f61 + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_96d072de9e3f7f61::Run); diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_12ac507a3c73f2784621_kernel.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_12ac507a3c73f2784621_kernel.cu new file mode 100644 index 00000000000..0fa750be2db --- /dev/null +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_12ac507a3c73f2784621_kernel.cu @@ -0,0 +1,1735 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +#define CAKE_INF CUDART_INF_F +#define TMEM_NCOLS 384 +#define TMEM_TMEM_OFFSET 0 +#define NUM_KV_PIPE_STAGES 4 +#define SMEM_SMEM_Q_OFF 1024 +#define SMEM_SMEM_Q_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_STRIDE 8192 +#define SMEM_SMEM_KV_OFF 66560 +#define SMEM_SMEM_KV_STAGE_BYTES 32768 +#define SMEM_SMEM_KV_STRIDE 32768 +#define SMEM_SMEM_V_OFF 66560 +#define SMEM_SMEM_V_STAGE_BYTES 32768 +#define SMEM_SMEM_V_STRIDE 32768 +#define SMEM_SMEM_SUM_OFF 197632 +#define SMEM_SMEM_SUM_STAGE_BYTES 256 +#define SMEM_SMEM_SUM_STRIDE 256 +#define SMEM_SMEM_INDEX_FLAGS_OFF 198400 +#define SMEM_SMEM_INDEX_FLAGS_STAGE_BYTES 16 +#define SMEM_SMEM_INDEX_FLAGS_STRIDE 16 +#define SMEM_SMEM_INDICES_OFF 197888 +#define SMEM_SMEM_INDICES_STAGE_BYTES 512 +#define SMEM_SMEM_INDICES_STRIDE 512 +#define SMEM_TOTAL 198528 +#define THREADS 512 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count) : "memory"); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_2d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3}], [%4];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr) { + // Canonical .shared::cta form for non-multicast gather4, matching + // trtllm-gen / cuda_ptx and the PTX ISA qualifier order + // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, + // .shared::cluster is reserved for the multicast variant (ctaMask). + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(512, 1) void +kernel_cake_dsv4_12ac507a3c73f2784621(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); + + const int mbar_base = smem; + #define q_full_addr (mbar_base + 0) + #define kv_full_addr (mbar_base + 8) + #define kv_empty_addr (mbar_base + 40) + #define s_full_addr (mbar_base + 72) + #define p_full_addr (mbar_base + 80) + #define sum_ready_addr (mbar_base + 88) + #define o_done_addr (mbar_base + 96) + #define tmem_dealloc_addr (mbar_base + 104) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + + // Kernel setup ops + __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_q_addr = smem + 1024; + __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_kv_addr = smem + 66560; + __nv_bfloat16* smem_v = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_v_addr = smem + 66560; + float* smem_sum = reinterpret_cast(smem_raw + 197632); + const int smem_sum_addr = smem + 197632; + int* smem_index_flags = reinterpret_cast(smem_raw + 198400); + const int smem_index_flags_addr = smem + 198400; + int* smem_indices = reinterpret_cast(smem_raw + 197888); + const int smem_indices_addr = smem + 197888; + + // Mbarrier init (8 pipeline groups, 0 ordered-sequence groups, 14 barriers) + // Mbarriers at smem_raw[0..112) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // q_full: 1 barriers, init_count=1 + mbarrier_init(smem + 0, 1); + // --- pipeline 'kv_pipe' --- + // kv_full: 4 barriers, init_count=4 + mbarrier_init(smem + 8, 4); + mbarrier_init(smem + 16, 4); + mbarrier_init(smem + 24, 4); + mbarrier_init(smem + 32, 4); + // kv_empty: 4 barriers, init_count=1 + mbarrier_init(smem + 40, 1); + mbarrier_init(smem + 48, 1); + mbarrier_init(smem + 56, 1); + mbarrier_init(smem + 64, 1); + // s_full: 1 barriers, init_count=1 + mbarrier_init(smem + 72, 1); + // p_full: 1 barriers, init_count=128 + mbarrier_init(smem + 80, 128); + // sum_ready: 1 barriers, init_count=128 + mbarrier_init(smem + 88, 128); + // o_done: 1 barriers, init_count=1 + mbarrier_init(smem + 96, 1); + // tmem_dealloc: 1 barriers, init_count=256 + mbarrier_init(smem + 104, 256); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncwarp(); + + // TMEM alloc (512 columns, 384 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 112); + if (warp == 0) { + int _tmem_hold = smem + 112; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_tmem = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 64;"); + } + + // ---- Role: softmax ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); + { // softmax_main + float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; + int query_idx = blockIdx.x >> 2; + int _max_0 = ((sparse_topk_lens[query_idx] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx] + sparse_topk_lens_offset) : (0)); + int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); + int active_topk = _min_0; + const int warp_in_compute = warp; + const int tmem_row_origin = warp_in_compute * 32; + const int logical_row_origin = warp_in_compute * 16; + const int my_row = logical_row_origin + lane % 16; + const int col_half = lane / 16; + unsigned int _phase_s_full_0 = 0; + mbarrier_wait(s_full_addr, _phase_s_full_0); + _phase_s_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int score_addr = taddr + (unsigned int)(tmem_row_origin << 16); + float _tmem_load_0[64]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_0[0])), "=r"(*reinterpret_cast(&_tmem_load_0[1])), "=r"(*reinterpret_cast(&_tmem_load_0[2])), "=r"(*reinterpret_cast(&_tmem_load_0[3])), "=r"(*reinterpret_cast(&_tmem_load_0[4])), "=r"(*reinterpret_cast(&_tmem_load_0[5])), "=r"(*reinterpret_cast(&_tmem_load_0[6])), "=r"(*reinterpret_cast(&_tmem_load_0[7])), "=r"(*reinterpret_cast(&_tmem_load_0[8])), "=r"(*reinterpret_cast(&_tmem_load_0[9])), "=r"(*reinterpret_cast(&_tmem_load_0[10])), "=r"(*reinterpret_cast(&_tmem_load_0[11])), "=r"(*reinterpret_cast(&_tmem_load_0[12])), "=r"(*reinterpret_cast(&_tmem_load_0[13])), "=r"(*reinterpret_cast(&_tmem_load_0[14])), "=r"(*reinterpret_cast(&_tmem_load_0[15])), "=r"(*reinterpret_cast(&_tmem_load_0[16])), "=r"(*reinterpret_cast(&_tmem_load_0[17])), "=r"(*reinterpret_cast(&_tmem_load_0[18])), "=r"(*reinterpret_cast(&_tmem_load_0[19])), "=r"(*reinterpret_cast(&_tmem_load_0[20])), "=r"(*reinterpret_cast(&_tmem_load_0[21])), "=r"(*reinterpret_cast(&_tmem_load_0[22])), "=r"(*reinterpret_cast(&_tmem_load_0[23])), "=r"(*reinterpret_cast(&_tmem_load_0[24])), "=r"(*reinterpret_cast(&_tmem_load_0[25])), "=r"(*reinterpret_cast(&_tmem_load_0[26])), "=r"(*reinterpret_cast(&_tmem_load_0[27])), "=r"(*reinterpret_cast(&_tmem_load_0[28])), "=r"(*reinterpret_cast(&_tmem_load_0[29])), "=r"(*reinterpret_cast(&_tmem_load_0[30])), "=r"(*reinterpret_cast(&_tmem_load_0[31])) + : "r"(score_addr)); + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_0[32])), "=r"(*reinterpret_cast(&_tmem_load_0[33])), "=r"(*reinterpret_cast(&_tmem_load_0[34])), "=r"(*reinterpret_cast(&_tmem_load_0[35])), "=r"(*reinterpret_cast(&_tmem_load_0[36])), "=r"(*reinterpret_cast(&_tmem_load_0[37])), "=r"(*reinterpret_cast(&_tmem_load_0[38])), "=r"(*reinterpret_cast(&_tmem_load_0[39])), "=r"(*reinterpret_cast(&_tmem_load_0[40])), "=r"(*reinterpret_cast(&_tmem_load_0[41])), "=r"(*reinterpret_cast(&_tmem_load_0[42])), "=r"(*reinterpret_cast(&_tmem_load_0[43])), "=r"(*reinterpret_cast(&_tmem_load_0[44])), "=r"(*reinterpret_cast(&_tmem_load_0[45])), "=r"(*reinterpret_cast(&_tmem_load_0[46])), "=r"(*reinterpret_cast(&_tmem_load_0[47])), "=r"(*reinterpret_cast(&_tmem_load_0[48])), "=r"(*reinterpret_cast(&_tmem_load_0[49])), "=r"(*reinterpret_cast(&_tmem_load_0[50])), "=r"(*reinterpret_cast(&_tmem_load_0[51])), "=r"(*reinterpret_cast(&_tmem_load_0[52])), "=r"(*reinterpret_cast(&_tmem_load_0[53])), "=r"(*reinterpret_cast(&_tmem_load_0[54])), "=r"(*reinterpret_cast(&_tmem_load_0[55])), "=r"(*reinterpret_cast(&_tmem_load_0[56])), "=r"(*reinterpret_cast(&_tmem_load_0[57])), "=r"(*reinterpret_cast(&_tmem_load_0[58])), "=r"(*reinterpret_cast(&_tmem_load_0[59])), "=r"(*reinterpret_cast(&_tmem_load_0[60])), "=r"(*reinterpret_cast(&_tmem_load_0[61])), "=r"(*reinterpret_cast(&_tmem_load_0[62])), "=r"(*reinterpret_cast(&_tmem_load_0[63])) + : "r"(score_addr + 32)); + int valid_cols = active_topk - col_half * 64; + int causal_cols = smem_index_flags[col_half * 2 + 1]; + valid_cols = ((valid_cols < causal_cols) ? valid_cols : causal_cols); + uint32_t _slice_lo_mask_0; + { + int _lim_0 = valid_cols; + if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_0 & (1u << 0))) _tmem_load_0[0] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 1))) _tmem_load_0[1] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 2))) _tmem_load_0[2] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 3))) _tmem_load_0[3] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 4))) _tmem_load_0[4] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 5))) _tmem_load_0[5] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 6))) _tmem_load_0[6] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 7))) _tmem_load_0[7] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 8))) _tmem_load_0[8] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 9))) _tmem_load_0[9] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 10))) _tmem_load_0[10] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 11))) _tmem_load_0[11] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 12))) _tmem_load_0[12] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 13))) _tmem_load_0[13] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 14))) _tmem_load_0[14] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 15))) _tmem_load_0[15] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 16))) _tmem_load_0[16] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 17))) _tmem_load_0[17] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 18))) _tmem_load_0[18] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 19))) _tmem_load_0[19] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 20))) _tmem_load_0[20] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 21))) _tmem_load_0[21] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 22))) _tmem_load_0[22] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 23))) _tmem_load_0[23] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 24))) _tmem_load_0[24] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 25))) _tmem_load_0[25] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 26))) _tmem_load_0[26] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 27))) _tmem_load_0[27] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 28))) _tmem_load_0[28] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 29))) _tmem_load_0[29] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 30))) _tmem_load_0[30] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 31))) _tmem_load_0[31] = -CAKE_INF; + uint32_t _slice_lo_mask_1; + { + int _lim_1 = valid_cols - 32; + if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) _tmem_load_0[32] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) _tmem_load_0[33] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) _tmem_load_0[34] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) _tmem_load_0[35] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) _tmem_load_0[36] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) _tmem_load_0[37] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) _tmem_load_0[38] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) _tmem_load_0[39] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) _tmem_load_0[40] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) _tmem_load_0[41] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) _tmem_load_0[42] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) _tmem_load_0[43] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) _tmem_load_0[44] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) _tmem_load_0[45] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) _tmem_load_0[46] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) _tmem_load_0[47] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) _tmem_load_0[48] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) _tmem_load_0[49] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) _tmem_load_0[50] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) _tmem_load_0[51] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) _tmem_load_0[52] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) _tmem_load_0[53] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) _tmem_load_0[54] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) _tmem_load_0[55] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) _tmem_load_0[56] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) _tmem_load_0[57] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) _tmem_load_0[58] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) _tmem_load_0[59] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) _tmem_load_0[60] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) _tmem_load_0[61] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) _tmem_load_0[62] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) _tmem_load_0[63] = -CAKE_INF; + float2 _reg_reduce_max2_2 = {-CAKE_INF, -CAKE_INF}; + row_max_x32_accum(&_tmem_load_0[0], _reg_reduce_max2_2); + row_max_x32_accum(&_tmem_load_0[32], _reg_reduce_max2_2); + float _tmem_load_0_max = row_max_reduce(_reg_reduce_max2_2); + float row_max = _tmem_load_0_max; + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, row_max, 16); + float _max_1 = max_noftz(row_max, _shfl_xor_0); + row_max = _max_1; + if (has_sinks != 0 && my_row < num_heads) { + float sink_unscaled = sinks[my_row] * 1.4426950408889634f / softmax_scale_log2; + float _max_2 = max_noftz(row_max, sink_unscaled); + row_max = _max_2; + } + bool empty_row = row_max == -CAKE_INF; + float safe_max = ((empty_row) ? 0.0f : row_max); + float max_scaled = safe_max * softmax_scale_log2; + float2 _f2_0 = make_float2(softmax_scale_log2, softmax_scale_log2); + float2 _f2_1 = make_float2(-max_scaled, -max_scaled); + float2 _f2_2 = make_float2(0.0f, 0.0f); + float2 sum0 = _f2_2; + float2 _f2_3 = make_float2(0.0f, 0.0f); + float2 sum1 = _f2_3; + float2 _f2_4 = make_float2(0.0f, 0.0f); + float2 sum2 = _f2_4; + float2 _f2_5 = make_float2(0.0f, 0.0f); + float2 sum3 = _f2_5; + #pragma unroll + for (int i = 0; i < 64; i += 8) { + float2 _f2_6 = make_float2(_tmem_load_0[i], _tmem_load_0[i + 1]); + float2 _f2_7 = make_float2(_tmem_load_0[i + 2], _tmem_load_0[i + 3]); + float2 _f2_8 = make_float2(_tmem_load_0[i + 4], _tmem_load_0[i + 5]); + float2 _f2_9 = make_float2(_tmem_load_0[i + 6], _tmem_load_0[i + 7]); + float2 affine0 = fma_f32x2_rn_ftz(_f2_6, _f2_0, _f2_1); + float2 affine1 = fma_f32x2_rn_ftz(_f2_7, _f2_0, _f2_1); + float2 affine2 = fma_f32x2_rn_ftz(_f2_8, _f2_0, _f2_1); + float2 affine3 = fma_f32x2_rn_ftz(_f2_9, _f2_0, _f2_1); + float _exp2_0 = approx_exp2(affine0.x); + float exp0 = _exp2_0; + float _exp2_1 = approx_exp2(affine0.y); + float exp1 = _exp2_1; + float _exp2_2 = approx_exp2(affine1.x); + float exp2 = _exp2_2; + float _exp2_3 = approx_exp2(affine1.y); + float exp3 = _exp2_3; + float _exp2_4 = approx_exp2(affine2.x); + float exp4 = _exp2_4; + float _exp2_5 = approx_exp2(affine2.y); + float exp5 = _exp2_5; + float _exp2_6 = approx_exp2(affine3.x); + float exp6 = _exp2_6; + float _exp2_7 = approx_exp2(affine3.y); + float exp7 = _exp2_7; + _tmem_load_0[i] = exp0; + _tmem_load_0[i + 1] = exp1; + _tmem_load_0[i + 2] = exp2; + _tmem_load_0[i + 3] = exp3; + _tmem_load_0[i + 4] = exp4; + _tmem_load_0[i + 5] = exp5; + _tmem_load_0[i + 6] = exp6; + _tmem_load_0[i + 7] = exp7; + float2 _f2_10 = make_float2(exp0, exp1); + sum0 = add_f32x2(sum0, _f2_10); + float2 _f2_11 = make_float2(exp2, exp3); + sum1 = add_f32x2(sum1, _f2_11); + float2 _f2_12 = make_float2(exp4, exp5); + sum2 = add_f32x2(sum2, _f2_12); + float2 _f2_13 = make_float2(exp6, exp7); + sum3 = add_f32x2(sum3, _f2_13); + } + float2 sum01 = add_f32x2(sum0, sum1); + float2 sum23 = add_f32x2(sum2, sum3); + float2 row_sum_pair = add_f32x2(sum01, sum23); + float row_sum = row_sum_pair.x + row_sum_pair.y; + float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, row_sum, 16); + row_sum = row_sum + _shfl_xor_1; + if (has_sinks != 0 && my_row < num_heads) { + float _exp2_8 = approx_exp2(sinks[my_row] * 1.4426950408889634f - max_scaled); + row_sum = row_sum + _exp2_8; + } + unsigned int packed_p[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_0[_lp*2 + 0], _tmem_load_0[_lp*2+1 + 0])); + packed_p[_lp] = *(uint32_t*)&_bf2; + } + int p_addr = taddr + 128 + (unsigned int)(tmem_row_origin << 16); + asm volatile( + "tcgen05.st.sync.aligned.16x32bx2.x16.b32" + " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(p_addr), "r"(*reinterpret_cast(&packed_p[0])), "r"(*reinterpret_cast(&packed_p[1])), "r"(*reinterpret_cast(&packed_p[2])), "r"(*reinterpret_cast(&packed_p[3])), "r"(*reinterpret_cast(&packed_p[4])), "r"(*reinterpret_cast(&packed_p[5])), "r"(*reinterpret_cast(&packed_p[6])), "r"(*reinterpret_cast(&packed_p[7])), "r"(*reinterpret_cast(&packed_p[8])), "r"(*reinterpret_cast(&packed_p[9])), "r"(*reinterpret_cast(&packed_p[10])), "r"(*reinterpret_cast(&packed_p[11])), "r"(*reinterpret_cast(&packed_p[12])), "r"(*reinterpret_cast(&packed_p[13])), "r"(*reinterpret_cast(&packed_p[14])), "r"(*reinterpret_cast(&packed_p[15]))); + asm volatile( + "tcgen05.st.sync.aligned.16x32bx2.x16.b32" + " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(p_addr + 16), "r"(*reinterpret_cast(&(packed_p + 16)[0])), "r"(*reinterpret_cast(&(packed_p + 16)[1])), "r"(*reinterpret_cast(&(packed_p + 16)[2])), "r"(*reinterpret_cast(&(packed_p + 16)[3])), "r"(*reinterpret_cast(&(packed_p + 16)[4])), "r"(*reinterpret_cast(&(packed_p + 16)[5])), "r"(*reinterpret_cast(&(packed_p + 16)[6])), "r"(*reinterpret_cast(&(packed_p + 16)[7])), "r"(*reinterpret_cast(&(packed_p + 16)[8])), "r"(*reinterpret_cast(&(packed_p + 16)[9])), "r"(*reinterpret_cast(&(packed_p + 16)[10])), "r"(*reinterpret_cast(&(packed_p + 16)[11])), "r"(*reinterpret_cast(&(packed_p + 16)[12])), "r"(*reinterpret_cast(&(packed_p + 16)[13])), "r"(*reinterpret_cast(&(packed_p + 16)[14])), "r"(*reinterpret_cast(&(packed_p + 16)[15]))); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(p_full_addr); + if (col_half == 0) { + float _max_3 = max_noftz(row_sum, 1.1754943508222875e-38f); + float published_sum = _max_3; + smem_sum[my_row] = published_sum; + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(sum_ready_addr); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_dealloc_addr); + } + } + // ---- Role: epilogue ---- + if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 128;"); + { // epilogue_main + float output_scale = bmm2_scale[0]; + int work_idx = blockIdx.x; + int query_idx_1 = work_idx >> 2; + int v_chunk = work_idx & 3; + const int warp_in_role = warp - 4; + const int tmem_row_origin_1 = warp_in_role * 32; + const int logical_row_origin_1 = warp_in_role * 16; + const int my_row_1 = logical_row_origin_1 + lane % 16; + const int col_half_1 = lane / 16; + const int row_addr = tmem_row_origin_1 << 16; + unsigned int _phase_o_done_0 = 0; + mbarrier_wait(o_done_addr, _phase_o_done_0); + _phase_o_done_0 ^= 1; + unsigned int _phase_sum_ready_0 = 0; + mbarrier_wait(sum_ready_addr, _phase_sum_ready_0); + _phase_sum_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + float row_sum_1 = smem_sum[my_row_1]; + float _rcp_0 = approx_rcp(row_sum_1); + float inv_sum = _rcp_0; + int output_base = (query_idx_1 * num_heads + my_row_1) * 512 + v_chunk * 128; + float _tmem_load_1[64]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_1[0])), "=r"(*reinterpret_cast(&_tmem_load_1[1])), "=r"(*reinterpret_cast(&_tmem_load_1[2])), "=r"(*reinterpret_cast(&_tmem_load_1[3])), "=r"(*reinterpret_cast(&_tmem_load_1[4])), "=r"(*reinterpret_cast(&_tmem_load_1[5])), "=r"(*reinterpret_cast(&_tmem_load_1[6])), "=r"(*reinterpret_cast(&_tmem_load_1[7])), "=r"(*reinterpret_cast(&_tmem_load_1[8])), "=r"(*reinterpret_cast(&_tmem_load_1[9])), "=r"(*reinterpret_cast(&_tmem_load_1[10])), "=r"(*reinterpret_cast(&_tmem_load_1[11])), "=r"(*reinterpret_cast(&_tmem_load_1[12])), "=r"(*reinterpret_cast(&_tmem_load_1[13])), "=r"(*reinterpret_cast(&_tmem_load_1[14])), "=r"(*reinterpret_cast(&_tmem_load_1[15])), "=r"(*reinterpret_cast(&_tmem_load_1[16])), "=r"(*reinterpret_cast(&_tmem_load_1[17])), "=r"(*reinterpret_cast(&_tmem_load_1[18])), "=r"(*reinterpret_cast(&_tmem_load_1[19])), "=r"(*reinterpret_cast(&_tmem_load_1[20])), "=r"(*reinterpret_cast(&_tmem_load_1[21])), "=r"(*reinterpret_cast(&_tmem_load_1[22])), "=r"(*reinterpret_cast(&_tmem_load_1[23])), "=r"(*reinterpret_cast(&_tmem_load_1[24])), "=r"(*reinterpret_cast(&_tmem_load_1[25])), "=r"(*reinterpret_cast(&_tmem_load_1[26])), "=r"(*reinterpret_cast(&_tmem_load_1[27])), "=r"(*reinterpret_cast(&_tmem_load_1[28])), "=r"(*reinterpret_cast(&_tmem_load_1[29])), "=r"(*reinterpret_cast(&_tmem_load_1[30])), "=r"(*reinterpret_cast(&_tmem_load_1[31])) + : "r"(taddr + 256 + (unsigned int)row_addr)); + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_1[32])), "=r"(*reinterpret_cast(&_tmem_load_1[33])), "=r"(*reinterpret_cast(&_tmem_load_1[34])), "=r"(*reinterpret_cast(&_tmem_load_1[35])), "=r"(*reinterpret_cast(&_tmem_load_1[36])), "=r"(*reinterpret_cast(&_tmem_load_1[37])), "=r"(*reinterpret_cast(&_tmem_load_1[38])), "=r"(*reinterpret_cast(&_tmem_load_1[39])), "=r"(*reinterpret_cast(&_tmem_load_1[40])), "=r"(*reinterpret_cast(&_tmem_load_1[41])), "=r"(*reinterpret_cast(&_tmem_load_1[42])), "=r"(*reinterpret_cast(&_tmem_load_1[43])), "=r"(*reinterpret_cast(&_tmem_load_1[44])), "=r"(*reinterpret_cast(&_tmem_load_1[45])), "=r"(*reinterpret_cast(&_tmem_load_1[46])), "=r"(*reinterpret_cast(&_tmem_load_1[47])), "=r"(*reinterpret_cast(&_tmem_load_1[48])), "=r"(*reinterpret_cast(&_tmem_load_1[49])), "=r"(*reinterpret_cast(&_tmem_load_1[50])), "=r"(*reinterpret_cast(&_tmem_load_1[51])), "=r"(*reinterpret_cast(&_tmem_load_1[52])), "=r"(*reinterpret_cast(&_tmem_load_1[53])), "=r"(*reinterpret_cast(&_tmem_load_1[54])), "=r"(*reinterpret_cast(&_tmem_load_1[55])), "=r"(*reinterpret_cast(&_tmem_load_1[56])), "=r"(*reinterpret_cast(&_tmem_load_1[57])), "=r"(*reinterpret_cast(&_tmem_load_1[58])), "=r"(*reinterpret_cast(&_tmem_load_1[59])), "=r"(*reinterpret_cast(&_tmem_load_1[60])), "=r"(*reinterpret_cast(&_tmem_load_1[61])), "=r"(*reinterpret_cast(&_tmem_load_1[62])), "=r"(*reinterpret_cast(&_tmem_load_1[63])) + : "r"(taddr + 256 + (unsigned int)row_addr + 32)); + float2 _f2_14 = make_float2(inv_sum * output_scale, inv_sum * output_scale); + #pragma unroll + for (int i_1 = 0; i_1 < 64; i_1 += 2) { + float2 _f2_15 = make_float2(_tmem_load_1[i_1], _tmem_load_1[i_1 + 1]); + float2 _mul_f32x2_0; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&_mul_f32x2_0) : "l"(*(const unsigned long long*)&_f2_15), "l"(*(const unsigned long long*)&_f2_14)); + _tmem_load_1[i_1] = _mul_f32x2_0.x; + _tmem_load_1[i_1 + 1] = _mul_f32x2_0.y; + } + unsigned int packed_o[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_1[_lp*2 + 0], _tmem_load_1[_lp*2+1 + 0])); + packed_o[_lp] = *(uint32_t*)&_bf2; + } + bool stage_hi_1 = (lane & 1) != 0; + bool stage_hi_2 = (lane & 2) != 0; + bool stage_hi_4 = (lane & 4) != 0; + unsigned int _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[0] : packed_o[4]), 1); + packed_o[0] = ((stage_hi_1) ? _shfl_xor_2 : packed_o[0]); + packed_o[4] = ((stage_hi_1) ? packed_o[4] : _shfl_xor_2); + unsigned int _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[1] : packed_o[5]), 1); + packed_o[1] = ((stage_hi_1) ? _shfl_xor_3 : packed_o[1]); + packed_o[5] = ((stage_hi_1) ? packed_o[5] : _shfl_xor_3); + unsigned int _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[2] : packed_o[6]), 1); + packed_o[2] = ((stage_hi_1) ? _shfl_xor_4 : packed_o[2]); + packed_o[6] = ((stage_hi_1) ? packed_o[6] : _shfl_xor_4); + unsigned int _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[3] : packed_o[7]), 1); + packed_o[3] = ((stage_hi_1) ? _shfl_xor_5 : packed_o[3]); + packed_o[7] = ((stage_hi_1) ? packed_o[7] : _shfl_xor_5); + unsigned int _shfl_xor_6 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[8] : packed_o[12]), 1); + packed_o[8] = ((stage_hi_1) ? _shfl_xor_6 : packed_o[8]); + packed_o[12] = ((stage_hi_1) ? packed_o[12] : _shfl_xor_6); + unsigned int _shfl_xor_7 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[9] : packed_o[13]), 1); + packed_o[9] = ((stage_hi_1) ? _shfl_xor_7 : packed_o[9]); + packed_o[13] = ((stage_hi_1) ? packed_o[13] : _shfl_xor_7); + unsigned int _shfl_xor_8 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[10] : packed_o[14]), 1); + packed_o[10] = ((stage_hi_1) ? _shfl_xor_8 : packed_o[10]); + packed_o[14] = ((stage_hi_1) ? packed_o[14] : _shfl_xor_8); + unsigned int _shfl_xor_9 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[11] : packed_o[15]), 1); + packed_o[11] = ((stage_hi_1) ? _shfl_xor_9 : packed_o[11]); + packed_o[15] = ((stage_hi_1) ? packed_o[15] : _shfl_xor_9); + unsigned int _shfl_xor_10 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[16] : packed_o[20]), 1); + packed_o[16] = ((stage_hi_1) ? _shfl_xor_10 : packed_o[16]); + packed_o[20] = ((stage_hi_1) ? packed_o[20] : _shfl_xor_10); + unsigned int _shfl_xor_11 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[17] : packed_o[21]), 1); + packed_o[17] = ((stage_hi_1) ? _shfl_xor_11 : packed_o[17]); + packed_o[21] = ((stage_hi_1) ? packed_o[21] : _shfl_xor_11); + unsigned int _shfl_xor_12 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[18] : packed_o[22]), 1); + packed_o[18] = ((stage_hi_1) ? _shfl_xor_12 : packed_o[18]); + packed_o[22] = ((stage_hi_1) ? packed_o[22] : _shfl_xor_12); + unsigned int _shfl_xor_13 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[19] : packed_o[23]), 1); + packed_o[19] = ((stage_hi_1) ? _shfl_xor_13 : packed_o[19]); + packed_o[23] = ((stage_hi_1) ? packed_o[23] : _shfl_xor_13); + unsigned int _shfl_xor_14 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[24] : packed_o[28]), 1); + packed_o[24] = ((stage_hi_1) ? _shfl_xor_14 : packed_o[24]); + packed_o[28] = ((stage_hi_1) ? packed_o[28] : _shfl_xor_14); + unsigned int _shfl_xor_15 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[25] : packed_o[29]), 1); + packed_o[25] = ((stage_hi_1) ? _shfl_xor_15 : packed_o[25]); + packed_o[29] = ((stage_hi_1) ? packed_o[29] : _shfl_xor_15); + unsigned int _shfl_xor_16 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[26] : packed_o[30]), 1); + packed_o[26] = ((stage_hi_1) ? _shfl_xor_16 : packed_o[26]); + packed_o[30] = ((stage_hi_1) ? packed_o[30] : _shfl_xor_16); + unsigned int _shfl_xor_17 = __shfl_xor_sync(0xFFFFFFFF, ((stage_hi_1) ? packed_o[27] : packed_o[31]), 1); + packed_o[27] = ((stage_hi_1) ? _shfl_xor_17 : packed_o[27]); + packed_o[31] = ((stage_hi_1) ? packed_o[31] : _shfl_xor_17); + int lane_t = lane & 1; + unsigned int chunk_words[4]; + chunk_words[0] = packed_o[0]; + chunk_words[1] = packed_o[1]; + chunk_words[2] = packed_o[2]; + chunk_words[3] = packed_o[3]; + if ((my_row_1 & -2 | 0) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + (my_row_1 & -2 | 0)) * 512 + v_chunk * 128 + col_half_1 * 64 + (0 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[4]; + chunk_words[1] = packed_o[5]; + chunk_words[2] = packed_o[6]; + chunk_words[3] = packed_o[7]; + if ((my_row_1 & -2 | 1) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + (my_row_1 & -2 | 1)) * 512 + v_chunk * 128 + col_half_1 * 64 + (0 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[8]; + chunk_words[1] = packed_o[9]; + chunk_words[2] = packed_o[10]; + chunk_words[3] = packed_o[11]; + if ((my_row_1 & -2 | 0) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + (my_row_1 & -2 | 0)) * 512 + v_chunk * 128 + col_half_1 * 64 + (2 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[12]; + chunk_words[1] = packed_o[13]; + chunk_words[2] = packed_o[14]; + chunk_words[3] = packed_o[15]; + if ((my_row_1 & -2 | 1) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + (my_row_1 & -2 | 1)) * 512 + v_chunk * 128 + col_half_1 * 64 + (2 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[16]; + chunk_words[1] = packed_o[17]; + chunk_words[2] = packed_o[18]; + chunk_words[3] = packed_o[19]; + if ((my_row_1 & -2 | 0) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + (my_row_1 & -2 | 0)) * 512 + v_chunk * 128 + col_half_1 * 64 + (4 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[20]; + chunk_words[1] = packed_o[21]; + chunk_words[2] = packed_o[22]; + chunk_words[3] = packed_o[23]; + if ((my_row_1 & -2 | 1) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + (my_row_1 & -2 | 1)) * 512 + v_chunk * 128 + col_half_1 * 64 + (4 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[24]; + chunk_words[1] = packed_o[25]; + chunk_words[2] = packed_o[26]; + chunk_words[3] = packed_o[27]; + if ((my_row_1 & -2 | 0) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + (my_row_1 & -2 | 0)) * 512 + v_chunk * 128 + col_half_1 * 64 + (6 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + chunk_words[0] = packed_o[28]; + chunk_words[1] = packed_o[29]; + chunk_words[2] = packed_o[30]; + chunk_words[3] = packed_o[31]; + if ((my_row_1 & -2 | 1) < num_heads) { + reinterpret_cast(O + ((query_idx_1 * num_heads + (my_row_1 & -2 | 1)) * 512 + v_chunk * 128 + col_half_1 * 64 + (6 | lane_t) * 8))[0] = reinterpret_cast(chunk_words)[0]; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_dealloc_addr); + } + } + // ---- Role: mma_warp ---- + if (warp == 8) { + { // mma_warp_main + unsigned int _phase_q_full_0 = 0; + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + unsigned int _phase_kv_full = 0; + #pragma unroll + for (int k_stage = 0; k_stage < 4; k_stage++) { + mbarrier_wait(kv_full_addr + (k_stage) * 8, _phase_kv_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_0 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (k_stage * 2) * 512); + int _mma_b_lo_0 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (k_stage) * 2048); + { + uint64_t _mma_ss_a_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_0); + uint64_t _mma_ss_b_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_0); + if (elect_sync()) { + tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, ((k_stage == 0) ? 0 : 1)); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); + } + } + int _mma_a_lo_1 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (k_stage * 2 + 1) * 512); + int _mma_b_lo_1 = make_warp_uniform((((smem_kv_addr + 16384) >> 4) & 0x3FFF) + (k_stage) * 2048); + { + uint64_t _mma_ss_a_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_1); + uint64_t _mma_ss_b_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_1); + if (elect_sync()) { + tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); + } + } + if (k_stage == 3) { + elect_commit(s_full_addr); + } + elect_commit(kv_empty_addr + (k_stage) * 8); + } + int _mma_prepared_b_lo_2 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (0) * 2048); + unsigned int _phase_p_full_0 = 0; + mbarrier_wait(p_full_addr, _phase_p_full_0); + _phase_p_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + mbarrier_wait(kv_full_addr, 1); + asm volatile("tcgen05.fence::after_thread_sync;"); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 dhi, blo, ta, id;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 69272720;\n\t" + "mov.b32 ta, %2;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p0;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "}\n" + :: "r"((tmem_tmem + (256))), "r"(_mma_prepared_b_lo_2), "r"(tmem_tmem + 128), "r"(0)); + elect_commit(o_done_addr); + elect_commit(kv_empty_addr); + unsigned int _phase_tmem_dealloc_0 = 0; + mbarrier_wait(tmem_dealloc_addr, _phase_tmem_dealloc_0); + _phase_tmem_dealloc_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + } + // ---- Role: empty ---- + if (warp >= 9 && warp <= 11) { + // idle — no tasks assigned + } + // ---- Role: load_warp ---- + if (warp >= 12 && warp <= 15) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 64;"); + { // load_warp_main + const int load_warp_rank = warp - 12; + int work_idx_1 = blockIdx.x; + int query_idx_2 = work_idx_1 >> 2; + int v_chunk_1 = work_idx_1 & 3; + int* row_ptr = swa_indices + (query_idx_2 * swa_index_stride); + int index_offset = load_warp_rank * 32 + lane; + int sparse_row = -1; + if (index_offset < sparse_topk) { + sparse_row = row_ptr[index_offset]; + } + if (load_warp_rank == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(q_full_addr, 65536); + #pragma unroll + for (int q_stage = 0; q_stage < 8; q_stage++) { + tma_4d_gmem2smem(smem_q_addr + (unsigned int)(q_stage * 8192), (&tmap_q), 0, 0, q_stage, query_idx_2, q_full_addr); + } + } + } + smem_indices[index_offset] = sparse_row; + unsigned int _vote_0 = __ballot_sync(0xFFFFFFFF, sparse_row < 0); + unsigned int invalid_indices = _vote_0; + int valid_prefix = 32; + if (invalid_indices != 0) { + int _ffs_0 = __ffs(invalid_indices); + valid_prefix = _ffs_0 - 1; + } + if (elect_sync()) { + smem_index_flags[load_warp_rank] = valid_prefix; + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("barrier.sync 8, 128;" ::: "memory"); + if ((load_warp_rank & 1) != 0) { + if (elect_sync()) { + int lower_prefix = smem_index_flags[load_warp_rank - 1]; + int upper_prefix = smem_index_flags[load_warp_rank]; + smem_index_flags[load_warp_rank] = ((lower_prefix < 32) ? lower_prefix : 32 + upper_prefix); + } + } + unsigned int _phase_kv_empty = 1; + #pragma unroll + for (int k_stage_1 = 0; k_stage_1 < 4; k_stage_1++) { + mbarrier_wait(kv_empty_addr + (k_stage_1) * 8, _phase_kv_empty); + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (k_stage_1) * 8, 8192); + } + #pragma unroll + for (int local_group = 0; local_group < 8; local_group++) { + int dst_k = smem_kv_addr + (unsigned int)(k_stage_1 * 32768); + int group = load_warp_rank * 8 + local_group; + int group_offset = group * 4; + int raw_rows[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows[0])), "=r"(*reinterpret_cast(&raw_rows[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows[(0) + 3])) + : "r"(smem_indices_addr + (unsigned int)(group_offset * 4))); + int raw0 = raw_rows[0]; + int raw1 = raw_rows[1]; + int raw2 = raw_rows[2]; + int raw3 = raw_rows[3]; + int row0 = ((raw0 >= 0) ? raw0 : 0); + int row1 = ((raw1 >= 0) ? raw1 : 0); + int row2 = ((raw2 >= 0) ? raw2 : 0); + int row3 = ((raw3 >= 0) ? raw3 : 0); + if (elect_sync()) { + tma_gather4_gmem2smem(dst_k + group * 512, (&tmap_swa_kv), k_stage_1 * 128, row0, row1, row2, row3, kv_full_addr + (k_stage_1) * 8); + tma_gather4_gmem2smem(dst_k + 16384 + group * 512, (&tmap_swa_kv), k_stage_1 * 128 + 64, row0, row1, row2, row3, kv_full_addr + (k_stage_1) * 8); + } + } + } + #pragma unroll + for (int v_stage = 0; v_stage < 1; v_stage++) { + mbarrier_wait(kv_empty_addr + (v_stage) * 8, 0); + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (v_stage) * 8, 8192); + } + #pragma unroll + for (int local_group_1 = 0; local_group_1 < 8; local_group_1++) { + int dst_v = smem_v_addr + (unsigned int)(v_stage * 32768); + int v_col = v_chunk_1 * 128 + v_stage * 64; + int group_1 = load_warp_rank * 8 + local_group_1; + int group_offset_1 = group_1 * 4; + int raw_rows_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_1[0])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 3])) + : "r"(smem_indices_addr + (unsigned int)(group_offset_1 * 4))); + int raw0_1 = raw_rows_1[0]; + int raw1_1 = raw_rows_1[1]; + int raw2_1 = raw_rows_1[2]; + int raw3_1 = raw_rows_1[3]; + int row0_1 = ((raw0_1 >= 0) ? raw0_1 : 0); + int row1_1 = ((raw1_1 >= 0) ? raw1_1 : 0); + int row2_1 = ((raw2_1 >= 0) ? raw2_1 : 0); + int row3_1 = ((raw3_1 >= 0) ? raw3_1 : 0); + if (elect_sync()) { + tma_gather4_gmem2smem(dst_v + group_1 * 512, (&tmap_swa_kv), v_col, row0_1, row1_1, row2_1, row3_1, kv_full_addr + (v_stage) * 8); + tma_gather4_gmem2smem(dst_v + 16384 + group_1 * 512, (&tmap_swa_kv), v_col + 64, row0_1, row1_1, row2_1, row3_1, kv_full_addr + (v_stage) * 8); + } + } + } + } + } + + // Cleanup +} + +} // extern "C" diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_15eabf5d2a5f50db201b_binding.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_15eabf5d2a5f50db201b_binding.cu new file mode 100644 index 00000000000..d58a59b6a60 --- /dev/null +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_15eabf5d2a5f50db201b_binding.cu @@ -0,0 +1,436 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +// Generated by the Cake source exporter — do not edit. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_15eabf5d2a5f50db201b'. +#include +#include +#include +#include + +#include "tvm_ffi_utils.h" + +#include +#include +#include +#include +#include +#include +#include + +extern "C" __global__ void kernel_cake_dsv4_15eabf5d2a5f50db201b(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int num_query_tokens, int has_sinks); + + +namespace cake_host_shim_6c292f2d365fae68 { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t CakeDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +// 4D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + int64_t d3 = t.size(t.ndim() - 3); + TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(d3)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 64u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_swa_kv' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_swa_kv' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_kv'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_compressed_kv' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_compressed_kv(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_compressed_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_compressed_kv' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_compressed_kv' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_compressed_kv"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_compressed_kv' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_kv'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_kv') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_compressed_kv, TensorView arg_O, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_num_query_tokens, int64_t arg_has_sinks, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckContiguous(arg_tmap_q, "tmap_q"); + CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); + CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 4, 16, 1); + CheckCudaTensor(arg_tmap_compressed_kv, "tmap_compressed_kv"); + CheckDtype(arg_tmap_compressed_kv, "tmap_compressed_kv", 4, 16, 1); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_swa_indices, "swa_indices"); + CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); + CheckContiguous(arg_swa_indices, "swa_indices"); + CheckCudaTensor(arg_compressed_indices, "compressed_indices"); + CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); + CheckContiguous(arg_compressed_indices, "compressed_indices"); + CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); + CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); + CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); + CheckContiguous(arg_bmm1_scale, "bmm1_scale"); + CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); + CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); + CheckContiguous(arg_bmm2_scale, "bmm2_scale"); + TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) + << "scalar 'num_heads' value " << arg_num_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) + << "scalar 'swa_index_stride' value " << arg_swa_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) + << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) + << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) + << "scalar 'sparse_topk' value " << arg_sparse_topk + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_query_tokens >= -2147483648LL && arg_num_query_tokens <= 2147483647LL, ValueError) + << "scalar 'num_query_tokens' value " << arg_num_query_tokens + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) + << "scalar 'has_sinks' value " << arg_has_sinks + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); + CheckSameCudaDevice(arg_tmap_compressed_kv, arg_tmap_q, "tmap_compressed_kv", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); + CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); + CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); + CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + + cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); + CUtensorMap p_tmap_compressed_kv = EncodeTma_tmap_compressed_kv(arg_tmap_compressed_kv); + __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); + int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); + int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); + int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); + float* p_sinks = static_cast(arg_sinks.data_ptr()); + float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); + float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); + int32_t v_num_heads = (int32_t)arg_num_heads; + int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; + int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; + int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; + int32_t v_sparse_topk = (int32_t)arg_sparse_topk; + int32_t v_num_query_tokens = (int32_t)arg_num_query_tokens; + int32_t v_has_sinks = (int32_t)arg_has_sinks; + void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_tmap_compressed_kv, &p_O, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_num_query_tokens, &v_has_sinks}; + + static std::mutex smem_mu; + static auto* smem_status_by_device = + new std::unordered_map(); + cudaError_t smem_status = cudaSuccess; + { + std::lock_guard lock(smem_mu); + auto it = smem_status_by_device->find(dev.device_id); + if (it == smem_status_by_device->end()) { + int current_device = -1; + smem_status = cudaGetDevice(¤t_device); + if (smem_status == cudaSuccess && current_device != dev.device_id) { + smem_status = cudaErrorInvalidDevice; + } + if (smem_status == cudaSuccess) { + smem_status = cudaFuncSetAttribute( + reinterpret_cast(kernel_cake_dsv4_15eabf5d2a5f50db201b), + cudaFuncAttributeMaxDynamicSharedMemorySize, + 203904); + } + smem_status_by_device->emplace(dev.device_id, smem_status); + } else { + smem_status = it->second; + } + } + TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) + << "cudaFuncSetAttribute for kernel_cake_dsv4_15eabf5d2a5f50db201b failed: " + << cudaGetErrorString(smem_status); + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(512u, 1u, 1u); + + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes = 203904u; + config.stream = stream; + cudaError_t launch_status = cudaLaunchKernelExC( + &config, reinterpret_cast(kernel_cake_dsv4_15eabf5d2a5f50db201b), kargs); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernelExC for kernel_cake_dsv4_15eabf5d2a5f50db201b failed: " + << cudaGetErrorString(launch_status); +} + +} // namespace cake_host_shim_6c292f2d365fae68 + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_6c292f2d365fae68::Run); diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_15eabf5d2a5f50db201b_kernel.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_15eabf5d2a5f50db201b_kernel.cu new file mode 100644 index 00000000000..265088b7d34 --- /dev/null +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_15eabf5d2a5f50db201b_kernel.cu @@ -0,0 +1,2557 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +#define CAKE_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_OFFSET 0 +#define NUM_KV_PIPE_STAGES 4 +#define NUM_PAGE_PIPE_STAGES 6 +#define NUM_WORK_PIPE_STAGES 2 +#define NUM_THROTTLE_PIPE_STAGES 2 +#define SMEM_SMEM_Q_OFF 1024 +#define SMEM_SMEM_Q_STAGE_BYTES 16384 +#define SMEM_SMEM_Q_STRIDE 16384 +#define SMEM_SMEM_KV_OFF 66560 +#define SMEM_SMEM_KV_STAGE_BYTES 32768 +#define SMEM_SMEM_KV_STRIDE 32768 +#define SMEM_SMEM_V_OFF 66560 +#define SMEM_SMEM_V_STAGE_BYTES 32768 +#define SMEM_SMEM_V_STRIDE 32768 +#define SMEM_SMEM_PAGE_OFFSETS_OFF 197632 +#define SMEM_SMEM_PAGE_OFFSETS_STAGE_BYTES 1024 +#define SMEM_SMEM_PAGE_OFFSETS_STRIDE 1024 +#define SMEM_WORK_RESPONSE_OFF 203776 +#define SMEM_WORK_RESPONSE_STAGE_BYTES 16 +#define SMEM_WORK_RESPONSE_STRIDE 16 +#define SMEM_TOTAL 203904 +#define THREADS 512 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count) : "memory"); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ float2 ex2_emulation_f32x2_value(float2 value) { + const float c0 = 1.0f, c1 = 0.695146143436431884765625f; + const float c2 = 0.227564394474029541015625f, c3 = 0.077119089663028717041015625f; + const float magic = 12582912.0f; + float x0 = max_noftz(value.x, -127.0f), x1 = max_noftz(value.y, -127.0f); + float2 xc2 = make_float2(x0, x1), magic2 = make_float2(magic, magic); + float2 xr2; + asm("add.rm.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xr2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&magic2)); + float2 c3_2 = make_float2(c3, c3), c2_2 = make_float2(c2, c2); + float2 c1_2 = make_float2(c1, c1), c0_2 = make_float2(c0, c0); + float2 xrb2, xfrac2; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xrb2) + : "l"(*(unsigned long long*)&xr2), "l"(*(unsigned long long*)&magic2)); + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xfrac2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&xrb2)); + float2 poly2; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&c3_2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c2_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c1_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c0_2)); + int x0r_i, x1r_i, p0_i, p1_i; + asm("mov.b64 {%0, %1}, %2;" : "=r"(x0r_i), "=r"(x1r_i) : "l"(*(unsigned long long*)&xr2)); + asm("mov.b64 {%0, %1}, %2;" : "=r"(p0_i), "=r"(p1_i) : "l"(*(unsigned long long*)&poly2)); + float r0, r1; + asm("mov.b32 %0, %1;" : "=f"(r0) : "r"((x0r_i << 23) + p0_i)); + asm("mov.b32 %0, %1;" : "=f"(r1) : "r"((x1r_i << 23) + p1_i)); + return make_float2(r0, r1); +} + +__device__ __forceinline__ void ex2_emulation_f32x2(float* x0_ptr, float* x1_ptr) { + float2 result = ex2_emulation_f32x2_value(make_float2(*x0_ptr, *x1_ptr)); + *x0_ptr = result.x; *x1_ptr = result.y; +} + +__device__ __forceinline__ void softmax_frag_exp2_cast( + float* sv, uint32_t* pv, int use_emu) +{ + #pragma unroll + for (int j = 0; j < 16; j++) { + if (use_emu && j >= 12) + ex2_emulation_f32x2(&sv[j*2], &sv[j*2+1]); + else { + sv[j*2] = approx_exp2(sv[j*2]); + sv[j*2+1] = approx_exp2(sv[j*2+1]); + } + } + #pragma unroll + for (int j = 0; j < 16; j++) { + __nv_bfloat162 bf = __float22bfloat162_rn({sv[j*2], sv[j*2+1]}); + pv[j] = reinterpret_cast(bf); + } +} + + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_2d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3}], [%4];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr) { + // Canonical .shared::cta form for non-multicast gather4, matching + // trtllm-gen / cuda_ptx and the PTX ISA qualifier order + // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, + // .shared::cluster is reserved for the multicast variant (ctaMask). + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(512, 1) void +kernel_cake_dsv4_15eabf5d2a5f50db201b(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int num_query_tokens, int has_sinks) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); + + const int mbar_base = smem; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define kv_full_addr (mbar_base + 16) + #define kv_empty_addr (mbar_base + 48) + #define s_full_addr (mbar_base + 80) + #define p_full_addr (mbar_base + 96) + #define stats_addr (mbar_base + 112) + #define stats_empty_addr (mbar_base + 128) + #define o_empty_addr (mbar_base + 144) + #define o_full_addr (mbar_base + 152) + #define page_full_addr (mbar_base + 160) + #define page_empty_addr (mbar_base + 208) + #define work_full_addr (mbar_base + 256) + #define work_empty_addr (mbar_base + 272) + #define throttle_full_addr (mbar_base + 288) + #define throttle_empty_addr (mbar_base + 304) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + + // Kernel setup ops + __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_q_addr = smem + 1024; + __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_kv_addr = smem + 66560; + __nv_bfloat16* smem_v = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_v_addr = smem + 66560; + int* smem_page_offsets = reinterpret_cast(smem_raw + 197632); + const int smem_page_offsets_addr = smem + 197632; + unsigned int* work_response = reinterpret_cast(smem_raw + 203776); + const int work_response_addr = smem + 203776; + + // Mbarrier init (16 pipeline groups, 0 ordered-sequence groups, 40 barriers) + // Mbarriers at smem_raw[0..320) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // q_full: 1 barriers, init_count=1 + mbarrier_init(smem + 0, 1); + // q_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 8, 1); + // --- pipeline 'kv_pipe' --- + // kv_full: 4 barriers, init_count=1 + mbarrier_init(smem + 16, 1); + mbarrier_init(smem + 24, 1); + mbarrier_init(smem + 32, 1); + mbarrier_init(smem + 40, 1); + // kv_empty: 4 barriers, init_count=1 + mbarrier_init(smem + 48, 1); + mbarrier_init(smem + 56, 1); + mbarrier_init(smem + 64, 1); + mbarrier_init(smem + 72, 1); + // s_full: 2 barriers, init_count=1 + mbarrier_init(smem + 80, 1); + mbarrier_init(smem + 88, 1); + // p_full: 2 barriers, init_count=128 + mbarrier_init(smem + 96, 128); + mbarrier_init(smem + 104, 128); + // stats: 2 barriers, init_count=128 + mbarrier_init(smem + 112, 128); + mbarrier_init(smem + 120, 128); + // stats_empty: 2 barriers, init_count=128 + mbarrier_init(smem + 128, 128); + mbarrier_init(smem + 136, 128); + // o_empty: 1 barriers, init_count=128 + mbarrier_init(smem + 144, 128); + // o_full: 1 barriers, init_count=1 + mbarrier_init(smem + 152, 1); + // --- pipeline 'page_pipe' --- + // page_full: 6 barriers, init_count=32 + mbarrier_init(smem + 160, 32); + mbarrier_init(smem + 168, 32); + mbarrier_init(smem + 176, 32); + mbarrier_init(smem + 184, 32); + mbarrier_init(smem + 192, 32); + mbarrier_init(smem + 200, 32); + // page_empty: 6 barriers, init_count=128 + mbarrier_init(smem + 208, 128); + mbarrier_init(smem + 216, 128); + mbarrier_init(smem + 224, 128); + mbarrier_init(smem + 232, 128); + mbarrier_init(smem + 240, 128); + mbarrier_init(smem + 248, 128); + // --- pipeline 'work_pipe' --- + // work_full: 2 barriers, init_count=1 + mbarrier_init(smem + 256, 1); + mbarrier_init(smem + 264, 1); + // work_empty: 2 barriers, init_count=512 + mbarrier_init(smem + 272, 512); + mbarrier_init(smem + 280, 512); + // --- pipeline 'throttle_pipe' --- + // throttle_full: 2 barriers, init_count=128 + mbarrier_init(smem + 288, 128); + mbarrier_init(smem + 296, 128); + // throttle_empty: 2 barriers, init_count=32 + mbarrier_init(smem + 304, 32); + mbarrier_init(smem + 312, 32); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncwarp(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 320); + if (warp == 0) { + int _tmem_hold = smem + 320; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_tmem = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Inc phase consumes the registers released above. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 136;"); + } + + // ---- Role: softmax ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 152;"); + { // softmax_main + asm volatile("griddepcontrol.wait;" ::: "memory"); + float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; + unsigned int work_stage = 0; + unsigned int query_idx = blockIdx.x; + unsigned int global_tile = 0; + unsigned int stats_cursor = 0; + const int warp_in_compute = warp; + const int tmem_row_origin = warp_in_compute * 32; + const int logical_row_origin = warp_in_compute * 16; + const int my_row = logical_row_origin + lane % 16; + const int col_half = lane / 16; + unsigned int _phase_work_full = 0; + #pragma unroll 1 + for (unsigned int work_iter = 0; work_iter < num_query_tokens; work_iter++) { + if (query_idx >= (unsigned int)num_query_tokens) { + break; + } + int _max_4 = ((sparse_topk_lens[query_idx] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx] + sparse_topk_lens_offset) : (0)); + int _min_2 = ((_max_4) < (sparse_topk) ? (_max_4) : (sparse_topk)); + int active_topk = _min_2; + int _max_5 = (((active_topk + 128 - 1) / 128) > (1) ? ((active_topk + 128 - 1) / 128) : (1)); + int num_loop_steps = _max_5; + float row_max = -CAKE_INF; + float row_sum = 0.0f; + if (has_sinks != 0 && my_row < num_heads) { + row_max = sinks[my_row] * 1.4426950408889634f / softmax_scale_log2; + row_sum = 1.0f; + } + if (num_loop_steps > 0) { + unsigned int first_stats_stage = stats_cursor & 1; + unsigned int first_stats_empty_phase = stats_cursor >> 1 & 1 ^ 1; + mbarrier_wait(stats_empty_addr + (first_stats_stage) * 8, first_stats_empty_phase); + } + #pragma unroll 1 + for (int tile = 0; tile < num_loop_steps; tile++) { + unsigned int sync_stage = global_tile & 1; + unsigned int sync_phase = global_tile >> 1 & 1; + mbarrier_wait(s_full_addr + (sync_stage) * 8, sync_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int score_col = ((sync_stage != 0) ? 128 : 0); + int score_addr = taddr + (unsigned int)score_col + (unsigned int)(tmem_row_origin << 16); + float _tmem_load_0[64]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_0[0])), "=r"(*reinterpret_cast(&_tmem_load_0[1])), "=r"(*reinterpret_cast(&_tmem_load_0[2])), "=r"(*reinterpret_cast(&_tmem_load_0[3])), "=r"(*reinterpret_cast(&_tmem_load_0[4])), "=r"(*reinterpret_cast(&_tmem_load_0[5])), "=r"(*reinterpret_cast(&_tmem_load_0[6])), "=r"(*reinterpret_cast(&_tmem_load_0[7])), "=r"(*reinterpret_cast(&_tmem_load_0[8])), "=r"(*reinterpret_cast(&_tmem_load_0[9])), "=r"(*reinterpret_cast(&_tmem_load_0[10])), "=r"(*reinterpret_cast(&_tmem_load_0[11])), "=r"(*reinterpret_cast(&_tmem_load_0[12])), "=r"(*reinterpret_cast(&_tmem_load_0[13])), "=r"(*reinterpret_cast(&_tmem_load_0[14])), "=r"(*reinterpret_cast(&_tmem_load_0[15])), "=r"(*reinterpret_cast(&_tmem_load_0[16])), "=r"(*reinterpret_cast(&_tmem_load_0[17])), "=r"(*reinterpret_cast(&_tmem_load_0[18])), "=r"(*reinterpret_cast(&_tmem_load_0[19])), "=r"(*reinterpret_cast(&_tmem_load_0[20])), "=r"(*reinterpret_cast(&_tmem_load_0[21])), "=r"(*reinterpret_cast(&_tmem_load_0[22])), "=r"(*reinterpret_cast(&_tmem_load_0[23])), "=r"(*reinterpret_cast(&_tmem_load_0[24])), "=r"(*reinterpret_cast(&_tmem_load_0[25])), "=r"(*reinterpret_cast(&_tmem_load_0[26])), "=r"(*reinterpret_cast(&_tmem_load_0[27])), "=r"(*reinterpret_cast(&_tmem_load_0[28])), "=r"(*reinterpret_cast(&_tmem_load_0[29])), "=r"(*reinterpret_cast(&_tmem_load_0[30])), "=r"(*reinterpret_cast(&_tmem_load_0[31])) + : "r"(score_addr)); + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_0[32])), "=r"(*reinterpret_cast(&_tmem_load_0[33])), "=r"(*reinterpret_cast(&_tmem_load_0[34])), "=r"(*reinterpret_cast(&_tmem_load_0[35])), "=r"(*reinterpret_cast(&_tmem_load_0[36])), "=r"(*reinterpret_cast(&_tmem_load_0[37])), "=r"(*reinterpret_cast(&_tmem_load_0[38])), "=r"(*reinterpret_cast(&_tmem_load_0[39])), "=r"(*reinterpret_cast(&_tmem_load_0[40])), "=r"(*reinterpret_cast(&_tmem_load_0[41])), "=r"(*reinterpret_cast(&_tmem_load_0[42])), "=r"(*reinterpret_cast(&_tmem_load_0[43])), "=r"(*reinterpret_cast(&_tmem_load_0[44])), "=r"(*reinterpret_cast(&_tmem_load_0[45])), "=r"(*reinterpret_cast(&_tmem_load_0[46])), "=r"(*reinterpret_cast(&_tmem_load_0[47])), "=r"(*reinterpret_cast(&_tmem_load_0[48])), "=r"(*reinterpret_cast(&_tmem_load_0[49])), "=r"(*reinterpret_cast(&_tmem_load_0[50])), "=r"(*reinterpret_cast(&_tmem_load_0[51])), "=r"(*reinterpret_cast(&_tmem_load_0[52])), "=r"(*reinterpret_cast(&_tmem_load_0[53])), "=r"(*reinterpret_cast(&_tmem_load_0[54])), "=r"(*reinterpret_cast(&_tmem_load_0[55])), "=r"(*reinterpret_cast(&_tmem_load_0[56])), "=r"(*reinterpret_cast(&_tmem_load_0[57])), "=r"(*reinterpret_cast(&_tmem_load_0[58])), "=r"(*reinterpret_cast(&_tmem_load_0[59])), "=r"(*reinterpret_cast(&_tmem_load_0[60])), "=r"(*reinterpret_cast(&_tmem_load_0[61])), "=r"(*reinterpret_cast(&_tmem_load_0[62])), "=r"(*reinterpret_cast(&_tmem_load_0[63])) + : "r"(score_addr + 32)); + int valid_cols = active_topk - tile * 128 - col_half * 64; + if (valid_cols < 0) { + valid_cols = 0; + } + if (valid_cols > 64) { + valid_cols = 64; + } + uint32_t _slice_lo_mask_0; + { + int _lim_0 = valid_cols; + if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_0 & (1u << 0))) _tmem_load_0[0] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 1))) _tmem_load_0[1] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 2))) _tmem_load_0[2] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 3))) _tmem_load_0[3] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 4))) _tmem_load_0[4] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 5))) _tmem_load_0[5] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 6))) _tmem_load_0[6] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 7))) _tmem_load_0[7] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 8))) _tmem_load_0[8] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 9))) _tmem_load_0[9] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 10))) _tmem_load_0[10] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 11))) _tmem_load_0[11] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 12))) _tmem_load_0[12] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 13))) _tmem_load_0[13] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 14))) _tmem_load_0[14] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 15))) _tmem_load_0[15] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 16))) _tmem_load_0[16] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 17))) _tmem_load_0[17] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 18))) _tmem_load_0[18] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 19))) _tmem_load_0[19] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 20))) _tmem_load_0[20] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 21))) _tmem_load_0[21] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 22))) _tmem_load_0[22] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 23))) _tmem_load_0[23] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 24))) _tmem_load_0[24] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 25))) _tmem_load_0[25] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 26))) _tmem_load_0[26] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 27))) _tmem_load_0[27] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 28))) _tmem_load_0[28] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 29))) _tmem_load_0[29] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 30))) _tmem_load_0[30] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 31))) _tmem_load_0[31] = -CAKE_INF; + uint32_t _slice_lo_mask_1; + { + int _lim_1 = valid_cols - 32; + if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) _tmem_load_0[32] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) _tmem_load_0[33] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) _tmem_load_0[34] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) _tmem_load_0[35] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) _tmem_load_0[36] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) _tmem_load_0[37] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) _tmem_load_0[38] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) _tmem_load_0[39] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) _tmem_load_0[40] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) _tmem_load_0[41] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) _tmem_load_0[42] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) _tmem_load_0[43] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) _tmem_load_0[44] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) _tmem_load_0[45] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) _tmem_load_0[46] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) _tmem_load_0[47] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) _tmem_load_0[48] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) _tmem_load_0[49] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) _tmem_load_0[50] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) _tmem_load_0[51] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) _tmem_load_0[52] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) _tmem_load_0[53] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) _tmem_load_0[54] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) _tmem_load_0[55] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) _tmem_load_0[56] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) _tmem_load_0[57] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) _tmem_load_0[58] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) _tmem_load_0[59] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) _tmem_load_0[60] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) _tmem_load_0[61] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) _tmem_load_0[62] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) _tmem_load_0[63] = -CAKE_INF; + float2 _reg_reduce_max2_2 = {-CAKE_INF, -CAKE_INF}; + row_max_x32_accum(&_tmem_load_0[0], _reg_reduce_max2_2); + row_max_x32_accum(&_tmem_load_0[32], _reg_reduce_max2_2); + float _tmem_load_0_max = row_max_reduce(_reg_reduce_max2_2); + float tile_max = _tmem_load_0_max; + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, tile_max, 16); + float _max_6 = max_noftz(tile_max, _shfl_xor_0); + tile_max = _max_6; + float old_max = row_max; + float _max_7 = max_noftz(row_max, tile_max); + float new_max = _max_7; + float _fma_0 = __fmaf_rn(old_max, softmax_scale_log2, (-new_max) * softmax_scale_log2); + float delta = _fma_0; + float _exp2_0 = approx_exp2(delta); + float acc_scale = ((old_max > -CAKE_INF) ? _exp2_0 : 1.0f); + unsigned int stats_stage = stats_cursor & 1; + int stats_col = ((stats_stage != 0) ? 128 : 0); + float max_stats[2]; + max_stats[0] = old_max; + max_stats[1] = new_max; + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(taddr + (unsigned int)stats_col), "f"(max_stats[0]), "f"(max_stats[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (stats_stage) * 8); + stats_cursor = stats_cursor + 1; + unsigned int next_stats_stage = stats_cursor & 1; + unsigned int next_stats_empty_phase = stats_cursor >> 1 & 1 ^ 1; + mbarrier_wait(stats_empty_addr + (next_stats_stage) * 8, next_stats_empty_phase); + bool empty_row = new_max == -CAKE_INF; + float safe_max = ((empty_row) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + const float2 _fma_b2_3 = {softmax_scale_log2, softmax_scale_log2}; + const float2 _fma_c2_4 = {-max_scaled, -max_scaled}; + #pragma unroll + for (int _lf = 0; _lf < 32; _lf++) + fma_f32x2_inplace(&reinterpret_cast(_tmem_load_0)[_lf], _fma_b2_3, _fma_c2_4); + #pragma unroll + for (int _le = 0; _le < 64; _le++) { + _tmem_load_0[_le] = approx_exp2(_tmem_load_0[_le]); + } + float2 _reg_reduce_sum2_5 = make_float2(0.0f, 0.0f); + softmax_block_sum(&_tmem_load_0[0], &_reg_reduce_sum2_5); + softmax_block_sum(&_tmem_load_0[32], &_reg_reduce_sum2_5); + float _tmem_load_0_sum = _reg_reduce_sum2_5.x + _reg_reduce_sum2_5.y; + float block_sum = _tmem_load_0_sum; + float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, block_sum, 16); + block_sum = block_sum + _shfl_xor_1; + float _fma_1 = __fmaf_rn(row_sum, acc_scale, block_sum); + row_sum = _fma_1; + row_max = new_max; + if (tile == num_loop_steps - 1) { + float final_stats[2]; + float _max_8 = max_noftz(row_sum, 1.1754943508222875e-38f); + final_stats[0] = _max_8; + final_stats[1] = row_max; + int final_stats_col = ((next_stats_stage != 0) ? 128 : 0); + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(taddr + (unsigned int)final_stats_col), "f"(final_stats[0]), "f"(final_stats[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (next_stats_stage) * 8); + stats_cursor = stats_cursor + 1; + unsigned int tail_stats_stage = stats_cursor & 1; + unsigned int tail_stats_empty_phase = stats_cursor >> 1 & 1 ^ 1; + mbarrier_wait(stats_empty_addr + (tail_stats_stage) * 8, tail_stats_empty_phase); + mbarrier_arrive(stats_addr + (tail_stats_stage) * 8); + stats_cursor = stats_cursor + 1; + } + unsigned int packed_p[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_0[_lp*2 + 0], _tmem_load_0[_lp*2+1 + 0])); + packed_p[_lp] = *(uint32_t*)&_bf2; + } + int p_addr = taddr + (unsigned int)score_col + 32 + (unsigned int)(tmem_row_origin << 16); + asm volatile( + "tcgen05.st.sync.aligned.16x32bx2.x16.b32" + " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(p_addr), "r"(*reinterpret_cast(&packed_p[0])), "r"(*reinterpret_cast(&packed_p[1])), "r"(*reinterpret_cast(&packed_p[2])), "r"(*reinterpret_cast(&packed_p[3])), "r"(*reinterpret_cast(&packed_p[4])), "r"(*reinterpret_cast(&packed_p[5])), "r"(*reinterpret_cast(&packed_p[6])), "r"(*reinterpret_cast(&packed_p[7])), "r"(*reinterpret_cast(&packed_p[8])), "r"(*reinterpret_cast(&packed_p[9])), "r"(*reinterpret_cast(&packed_p[10])), "r"(*reinterpret_cast(&packed_p[11])), "r"(*reinterpret_cast(&packed_p[12])), "r"(*reinterpret_cast(&packed_p[13])), "r"(*reinterpret_cast(&packed_p[14])), "r"(*reinterpret_cast(&packed_p[15]))); + asm volatile( + "tcgen05.st.sync.aligned.16x32bx2.x16.b32" + " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(p_addr + 16), "r"(*reinterpret_cast(&(packed_p + 16)[0])), "r"(*reinterpret_cast(&(packed_p + 16)[1])), "r"(*reinterpret_cast(&(packed_p + 16)[2])), "r"(*reinterpret_cast(&(packed_p + 16)[3])), "r"(*reinterpret_cast(&(packed_p + 16)[4])), "r"(*reinterpret_cast(&(packed_p + 16)[5])), "r"(*reinterpret_cast(&(packed_p + 16)[6])), "r"(*reinterpret_cast(&(packed_p + 16)[7])), "r"(*reinterpret_cast(&(packed_p + 16)[8])), "r"(*reinterpret_cast(&(packed_p + 16)[9])), "r"(*reinterpret_cast(&(packed_p + 16)[10])), "r"(*reinterpret_cast(&(packed_p + 16)[11])), "r"(*reinterpret_cast(&(packed_p + 16)[12])), "r"(*reinterpret_cast(&(packed_p + 16)[13])), "r"(*reinterpret_cast(&(packed_p + 16)[14])), "r"(*reinterpret_cast(&(packed_p + 16)[15]))); + int p_addr_high = p_addr + 1048576; + asm volatile( + "tcgen05.st.sync.aligned.16x32bx2.x16.b32" + " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(p_addr_high), "r"(*reinterpret_cast(&packed_p[0])), "r"(*reinterpret_cast(&packed_p[1])), "r"(*reinterpret_cast(&packed_p[2])), "r"(*reinterpret_cast(&packed_p[3])), "r"(*reinterpret_cast(&packed_p[4])), "r"(*reinterpret_cast(&packed_p[5])), "r"(*reinterpret_cast(&packed_p[6])), "r"(*reinterpret_cast(&packed_p[7])), "r"(*reinterpret_cast(&packed_p[8])), "r"(*reinterpret_cast(&packed_p[9])), "r"(*reinterpret_cast(&packed_p[10])), "r"(*reinterpret_cast(&packed_p[11])), "r"(*reinterpret_cast(&packed_p[12])), "r"(*reinterpret_cast(&packed_p[13])), "r"(*reinterpret_cast(&packed_p[14])), "r"(*reinterpret_cast(&packed_p[15]))); + asm volatile( + "tcgen05.st.sync.aligned.16x32bx2.x16.b32" + " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(p_addr_high + 16), "r"(*reinterpret_cast(&(packed_p + 16)[0])), "r"(*reinterpret_cast(&(packed_p + 16)[1])), "r"(*reinterpret_cast(&(packed_p + 16)[2])), "r"(*reinterpret_cast(&(packed_p + 16)[3])), "r"(*reinterpret_cast(&(packed_p + 16)[4])), "r"(*reinterpret_cast(&(packed_p + 16)[5])), "r"(*reinterpret_cast(&(packed_p + 16)[6])), "r"(*reinterpret_cast(&(packed_p + 16)[7])), "r"(*reinterpret_cast(&(packed_p + 16)[8])), "r"(*reinterpret_cast(&(packed_p + 16)[9])), "r"(*reinterpret_cast(&(packed_p + 16)[10])), "r"(*reinterpret_cast(&(packed_p + 16)[11])), "r"(*reinterpret_cast(&(packed_p + 16)[12])), "r"(*reinterpret_cast(&(packed_p + 16)[13])), "r"(*reinterpret_cast(&(packed_p + 16)[14])), "r"(*reinterpret_cast(&(packed_p + 16)[15]))); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(p_full_addr + (sync_stage) * 8); + global_tile = global_tile + 1; + } + mbarrier_wait(work_full_addr + (work_stage) * 8, _phase_work_full); + uint32_t _clc_valid_2 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_2) + : "r"(work_response_addr + work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_2 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_2) + : "r"(work_response_addr + work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(work_empty_addr + (work_stage) * 8); + work_stage += 1; + if (work_stage == 2) { work_stage = 0; _phase_work_full ^= 1; } + unsigned int valid = _clc_valid_2; + query_idx = _clc_ctaid_2; + if (valid == 0) { + break; + } + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + } + } + // ---- Role: correction ---- + if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 144;"); + { // correction_main + asm volatile("griddepcontrol.wait;" ::: "memory"); + float softmax_scale_log2_1 = bmm1_scale[0] * 1.4426950408889634f; + float output_scale = bmm2_scale[0]; + unsigned int work_stage_1 = 0; + unsigned int query_idx_1 = blockIdx.x; + unsigned int global_tile_1 = 0; + unsigned int stats_cursor_1 = 0; + const int warp_in_role = warp - 4; + const int tmem_row_origin_1 = warp_in_role * 32; + const int logical_row_origin_1 = warp_in_role * 16; + const int my_row_1 = logical_row_origin_1 + lane % 16; + const int col_half_1 = lane / 16; + unsigned int _phase_work_full_1 = 0; + #pragma unroll 1 + for (unsigned int work_iter_1 = 0; work_iter_1 < num_query_tokens; work_iter_1++) { + if (query_idx_1 >= (unsigned int)num_query_tokens) { + break; + } + int _max_9 = ((sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) : (0)); + int _min_3 = ((_max_9) < (sparse_topk) ? (_max_9) : (sparse_topk)); + int active_topk_1 = _min_3; + int _max_10 = (((active_topk_1 + 128 - 1) / 128) > (1) ? ((active_topk_1 + 128 - 1) / 128) : (1)); + int num_loop_steps_1 = _max_10; + unsigned int first_stats_stage_1 = stats_cursor_1 & 1; + unsigned int first_stats_phase = stats_cursor_1 >> 1 & 1; + mbarrier_wait(stats_addr + (first_stats_stage_1) * 8, first_stats_phase); + mbarrier_arrive(stats_empty_addr + (first_stats_stage_1) * 8); + stats_cursor_1 = stats_cursor_1 + 1; + global_tile_1 = global_tile_1 + 1; + #pragma unroll 1 + for (int tile_1 = 0; tile_1 < num_loop_steps_1 - 1; tile_1++) { + unsigned int stats_stage_1 = stats_cursor_1 & 1; + unsigned int stats_phase = stats_cursor_1 >> 1 & 1; + mbarrier_wait(stats_addr + (stats_stage_1) * 8, stats_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int stats_col_1 = ((stats_stage_1 != 0) ? 128 : 0); + float _tmem_load_1[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]) + : "r"(taddr + (unsigned int)stats_col_1)); + mbarrier_arrive(stats_empty_addr + (stats_stage_1) * 8); + stats_cursor_1 = stats_cursor_1 + 1; + float max_diff = _tmem_load_1[0] - _tmem_load_1[1]; + float _exp2_1 = approx_exp2(softmax_scale_log2_1 * max_diff); + float acc_scale_1 = ((max_diff != 0.0f) ? _exp2_1 : 1.0f); + mbarrier_wait(o_full_addr, global_tile_1 - 1 & 1); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _vote_0 = __any_sync(0xFFFFFFFF, acc_scale_1 < 1.0f); + int any_rescale = _vote_0; + if (any_rescale != 0) { + #pragma unroll + for (int v_stage = 0; v_stage < 4; v_stage++) { + int output_addr = taddr + 256 + (unsigned int)((v_stage & 1) * 128) + (unsigned int)((v_stage >> 1) * 16 << 16) + (unsigned int)(tmem_row_origin_1 << 16); + float _tmem_load_2[64]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_2[0])), "=r"(*reinterpret_cast(&_tmem_load_2[1])), "=r"(*reinterpret_cast(&_tmem_load_2[2])), "=r"(*reinterpret_cast(&_tmem_load_2[3])), "=r"(*reinterpret_cast(&_tmem_load_2[4])), "=r"(*reinterpret_cast(&_tmem_load_2[5])), "=r"(*reinterpret_cast(&_tmem_load_2[6])), "=r"(*reinterpret_cast(&_tmem_load_2[7])), "=r"(*reinterpret_cast(&_tmem_load_2[8])), "=r"(*reinterpret_cast(&_tmem_load_2[9])), "=r"(*reinterpret_cast(&_tmem_load_2[10])), "=r"(*reinterpret_cast(&_tmem_load_2[11])), "=r"(*reinterpret_cast(&_tmem_load_2[12])), "=r"(*reinterpret_cast(&_tmem_load_2[13])), "=r"(*reinterpret_cast(&_tmem_load_2[14])), "=r"(*reinterpret_cast(&_tmem_load_2[15])), "=r"(*reinterpret_cast(&_tmem_load_2[16])), "=r"(*reinterpret_cast(&_tmem_load_2[17])), "=r"(*reinterpret_cast(&_tmem_load_2[18])), "=r"(*reinterpret_cast(&_tmem_load_2[19])), "=r"(*reinterpret_cast(&_tmem_load_2[20])), "=r"(*reinterpret_cast(&_tmem_load_2[21])), "=r"(*reinterpret_cast(&_tmem_load_2[22])), "=r"(*reinterpret_cast(&_tmem_load_2[23])), "=r"(*reinterpret_cast(&_tmem_load_2[24])), "=r"(*reinterpret_cast(&_tmem_load_2[25])), "=r"(*reinterpret_cast(&_tmem_load_2[26])), "=r"(*reinterpret_cast(&_tmem_load_2[27])), "=r"(*reinterpret_cast(&_tmem_load_2[28])), "=r"(*reinterpret_cast(&_tmem_load_2[29])), "=r"(*reinterpret_cast(&_tmem_load_2[30])), "=r"(*reinterpret_cast(&_tmem_load_2[31])) + : "r"(output_addr)); + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_2[32])), "=r"(*reinterpret_cast(&_tmem_load_2[33])), "=r"(*reinterpret_cast(&_tmem_load_2[34])), "=r"(*reinterpret_cast(&_tmem_load_2[35])), "=r"(*reinterpret_cast(&_tmem_load_2[36])), "=r"(*reinterpret_cast(&_tmem_load_2[37])), "=r"(*reinterpret_cast(&_tmem_load_2[38])), "=r"(*reinterpret_cast(&_tmem_load_2[39])), "=r"(*reinterpret_cast(&_tmem_load_2[40])), "=r"(*reinterpret_cast(&_tmem_load_2[41])), "=r"(*reinterpret_cast(&_tmem_load_2[42])), "=r"(*reinterpret_cast(&_tmem_load_2[43])), "=r"(*reinterpret_cast(&_tmem_load_2[44])), "=r"(*reinterpret_cast(&_tmem_load_2[45])), "=r"(*reinterpret_cast(&_tmem_load_2[46])), "=r"(*reinterpret_cast(&_tmem_load_2[47])), "=r"(*reinterpret_cast(&_tmem_load_2[48])), "=r"(*reinterpret_cast(&_tmem_load_2[49])), "=r"(*reinterpret_cast(&_tmem_load_2[50])), "=r"(*reinterpret_cast(&_tmem_load_2[51])), "=r"(*reinterpret_cast(&_tmem_load_2[52])), "=r"(*reinterpret_cast(&_tmem_load_2[53])), "=r"(*reinterpret_cast(&_tmem_load_2[54])), "=r"(*reinterpret_cast(&_tmem_load_2[55])), "=r"(*reinterpret_cast(&_tmem_load_2[56])), "=r"(*reinterpret_cast(&_tmem_load_2[57])), "=r"(*reinterpret_cast(&_tmem_load_2[58])), "=r"(*reinterpret_cast(&_tmem_load_2[59])), "=r"(*reinterpret_cast(&_tmem_load_2[60])), "=r"(*reinterpret_cast(&_tmem_load_2[61])), "=r"(*reinterpret_cast(&_tmem_load_2[62])), "=r"(*reinterpret_cast(&_tmem_load_2[63])) + : "r"(output_addr + 32)); + const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 32; _ls++) + mul_f32x2_inplace(&reinterpret_cast(_tmem_load_2)[_ls], _scale2_0); + asm volatile( + "tcgen05.st.sync.aligned.16x32bx2.x64.b32" + " [%0], 64, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64};" + :: "r"(output_addr), "r"(*reinterpret_cast(&_tmem_load_2[0])), "r"(*reinterpret_cast(&_tmem_load_2[1])), "r"(*reinterpret_cast(&_tmem_load_2[2])), "r"(*reinterpret_cast(&_tmem_load_2[3])), "r"(*reinterpret_cast(&_tmem_load_2[4])), "r"(*reinterpret_cast(&_tmem_load_2[5])), "r"(*reinterpret_cast(&_tmem_load_2[6])), "r"(*reinterpret_cast(&_tmem_load_2[7])), "r"(*reinterpret_cast(&_tmem_load_2[8])), "r"(*reinterpret_cast(&_tmem_load_2[9])), "r"(*reinterpret_cast(&_tmem_load_2[10])), "r"(*reinterpret_cast(&_tmem_load_2[11])), "r"(*reinterpret_cast(&_tmem_load_2[12])), "r"(*reinterpret_cast(&_tmem_load_2[13])), "r"(*reinterpret_cast(&_tmem_load_2[14])), "r"(*reinterpret_cast(&_tmem_load_2[15])), "r"(*reinterpret_cast(&_tmem_load_2[16])), "r"(*reinterpret_cast(&_tmem_load_2[17])), "r"(*reinterpret_cast(&_tmem_load_2[18])), "r"(*reinterpret_cast(&_tmem_load_2[19])), "r"(*reinterpret_cast(&_tmem_load_2[20])), "r"(*reinterpret_cast(&_tmem_load_2[21])), "r"(*reinterpret_cast(&_tmem_load_2[22])), "r"(*reinterpret_cast(&_tmem_load_2[23])), "r"(*reinterpret_cast(&_tmem_load_2[24])), "r"(*reinterpret_cast(&_tmem_load_2[25])), "r"(*reinterpret_cast(&_tmem_load_2[26])), "r"(*reinterpret_cast(&_tmem_load_2[27])), "r"(*reinterpret_cast(&_tmem_load_2[28])), "r"(*reinterpret_cast(&_tmem_load_2[29])), "r"(*reinterpret_cast(&_tmem_load_2[30])), "r"(*reinterpret_cast(&_tmem_load_2[31])), "r"(*reinterpret_cast(&_tmem_load_2[32])), "r"(*reinterpret_cast(&_tmem_load_2[33])), "r"(*reinterpret_cast(&_tmem_load_2[34])), "r"(*reinterpret_cast(&_tmem_load_2[35])), "r"(*reinterpret_cast(&_tmem_load_2[36])), "r"(*reinterpret_cast(&_tmem_load_2[37])), "r"(*reinterpret_cast(&_tmem_load_2[38])), "r"(*reinterpret_cast(&_tmem_load_2[39])), "r"(*reinterpret_cast(&_tmem_load_2[40])), "r"(*reinterpret_cast(&_tmem_load_2[41])), "r"(*reinterpret_cast(&_tmem_load_2[42])), "r"(*reinterpret_cast(&_tmem_load_2[43])), "r"(*reinterpret_cast(&_tmem_load_2[44])), "r"(*reinterpret_cast(&_tmem_load_2[45])), "r"(*reinterpret_cast(&_tmem_load_2[46])), "r"(*reinterpret_cast(&_tmem_load_2[47])), "r"(*reinterpret_cast(&_tmem_load_2[48])), "r"(*reinterpret_cast(&_tmem_load_2[49])), "r"(*reinterpret_cast(&_tmem_load_2[50])), "r"(*reinterpret_cast(&_tmem_load_2[51])), "r"(*reinterpret_cast(&_tmem_load_2[52])), "r"(*reinterpret_cast(&_tmem_load_2[53])), "r"(*reinterpret_cast(&_tmem_load_2[54])), "r"(*reinterpret_cast(&_tmem_load_2[55])), "r"(*reinterpret_cast(&_tmem_load_2[56])), "r"(*reinterpret_cast(&_tmem_load_2[57])), "r"(*reinterpret_cast(&_tmem_load_2[58])), "r"(*reinterpret_cast(&_tmem_load_2[59])), "r"(*reinterpret_cast(&_tmem_load_2[60])), "r"(*reinterpret_cast(&_tmem_load_2[61])), "r"(*reinterpret_cast(&_tmem_load_2[62])), "r"(*reinterpret_cast(&_tmem_load_2[63]))); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + } + } + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(o_empty_addr); + global_tile_1 = global_tile_1 + 1; + } + unsigned int final_stats_stage = stats_cursor_1 & 1; + unsigned int final_stats_phase = stats_cursor_1 >> 1 & 1; + mbarrier_wait(stats_addr + (final_stats_stage) * 8, final_stats_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int final_stats_col_1 = ((final_stats_stage != 0) ? 128 : 0); + float _tmem_load_3[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]) + : "r"(taddr + (unsigned int)final_stats_col_1)); + mbarrier_arrive(stats_empty_addr + (final_stats_stage) * 8); + stats_cursor_1 = stats_cursor_1 + 1; + unsigned int tail_stats_stage_1 = stats_cursor_1 & 1; + unsigned int tail_stats_phase = stats_cursor_1 >> 1 & 1; + mbarrier_wait(stats_addr + (tail_stats_stage_1) * 8, tail_stats_phase); + mbarrier_arrive(stats_empty_addr + (tail_stats_stage_1) * 8); + stats_cursor_1 = stats_cursor_1 + 1; + mbarrier_wait(o_full_addr, global_tile_1 - 1 & 1); + asm volatile("tcgen05.fence::after_thread_sync;"); + float final_sum = _tmem_load_3[0]; + float _rcp_0 = approx_rcp(final_sum); + float inv_sum = _rcp_0; + int output_base = (query_idx_1 * (unsigned int)num_heads + (unsigned int)my_row_1) * 512; + #pragma unroll + for (int v_stage_1 = 0; v_stage_1 < 4; v_stage_1++) { + int output_addr_1 = taddr + 256 + (unsigned int)((v_stage_1 & 1) * 128) + (unsigned int)((v_stage_1 >> 1) * 16 << 16) + (unsigned int)(tmem_row_origin_1 << 16); + float _tmem_load_4[64]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_4[0])), "=r"(*reinterpret_cast(&_tmem_load_4[1])), "=r"(*reinterpret_cast(&_tmem_load_4[2])), "=r"(*reinterpret_cast(&_tmem_load_4[3])), "=r"(*reinterpret_cast(&_tmem_load_4[4])), "=r"(*reinterpret_cast(&_tmem_load_4[5])), "=r"(*reinterpret_cast(&_tmem_load_4[6])), "=r"(*reinterpret_cast(&_tmem_load_4[7])), "=r"(*reinterpret_cast(&_tmem_load_4[8])), "=r"(*reinterpret_cast(&_tmem_load_4[9])), "=r"(*reinterpret_cast(&_tmem_load_4[10])), "=r"(*reinterpret_cast(&_tmem_load_4[11])), "=r"(*reinterpret_cast(&_tmem_load_4[12])), "=r"(*reinterpret_cast(&_tmem_load_4[13])), "=r"(*reinterpret_cast(&_tmem_load_4[14])), "=r"(*reinterpret_cast(&_tmem_load_4[15])), "=r"(*reinterpret_cast(&_tmem_load_4[16])), "=r"(*reinterpret_cast(&_tmem_load_4[17])), "=r"(*reinterpret_cast(&_tmem_load_4[18])), "=r"(*reinterpret_cast(&_tmem_load_4[19])), "=r"(*reinterpret_cast(&_tmem_load_4[20])), "=r"(*reinterpret_cast(&_tmem_load_4[21])), "=r"(*reinterpret_cast(&_tmem_load_4[22])), "=r"(*reinterpret_cast(&_tmem_load_4[23])), "=r"(*reinterpret_cast(&_tmem_load_4[24])), "=r"(*reinterpret_cast(&_tmem_load_4[25])), "=r"(*reinterpret_cast(&_tmem_load_4[26])), "=r"(*reinterpret_cast(&_tmem_load_4[27])), "=r"(*reinterpret_cast(&_tmem_load_4[28])), "=r"(*reinterpret_cast(&_tmem_load_4[29])), "=r"(*reinterpret_cast(&_tmem_load_4[30])), "=r"(*reinterpret_cast(&_tmem_load_4[31])) + : "r"(output_addr_1)); + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_4[32])), "=r"(*reinterpret_cast(&_tmem_load_4[33])), "=r"(*reinterpret_cast(&_tmem_load_4[34])), "=r"(*reinterpret_cast(&_tmem_load_4[35])), "=r"(*reinterpret_cast(&_tmem_load_4[36])), "=r"(*reinterpret_cast(&_tmem_load_4[37])), "=r"(*reinterpret_cast(&_tmem_load_4[38])), "=r"(*reinterpret_cast(&_tmem_load_4[39])), "=r"(*reinterpret_cast(&_tmem_load_4[40])), "=r"(*reinterpret_cast(&_tmem_load_4[41])), "=r"(*reinterpret_cast(&_tmem_load_4[42])), "=r"(*reinterpret_cast(&_tmem_load_4[43])), "=r"(*reinterpret_cast(&_tmem_load_4[44])), "=r"(*reinterpret_cast(&_tmem_load_4[45])), "=r"(*reinterpret_cast(&_tmem_load_4[46])), "=r"(*reinterpret_cast(&_tmem_load_4[47])), "=r"(*reinterpret_cast(&_tmem_load_4[48])), "=r"(*reinterpret_cast(&_tmem_load_4[49])), "=r"(*reinterpret_cast(&_tmem_load_4[50])), "=r"(*reinterpret_cast(&_tmem_load_4[51])), "=r"(*reinterpret_cast(&_tmem_load_4[52])), "=r"(*reinterpret_cast(&_tmem_load_4[53])), "=r"(*reinterpret_cast(&_tmem_load_4[54])), "=r"(*reinterpret_cast(&_tmem_load_4[55])), "=r"(*reinterpret_cast(&_tmem_load_4[56])), "=r"(*reinterpret_cast(&_tmem_load_4[57])), "=r"(*reinterpret_cast(&_tmem_load_4[58])), "=r"(*reinterpret_cast(&_tmem_load_4[59])), "=r"(*reinterpret_cast(&_tmem_load_4[60])), "=r"(*reinterpret_cast(&_tmem_load_4[61])), "=r"(*reinterpret_cast(&_tmem_load_4[62])), "=r"(*reinterpret_cast(&_tmem_load_4[63])) + : "r"(output_addr_1 + 32)); + int col_base = v_stage_1 * 128 + col_half_1 * 64; + if (my_row_1 < num_heads) { + #pragma unroll + for (int offset = 0; offset < 64; offset += 16) { + { + const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; + #if __CUDA_ARCH__ >= 1000 + #pragma unroll + for (int _ps = 0; _ps < 8; _ps++) + mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_4[offset])[_ps], _prescale2_1); + #else + #pragma unroll + for (int _ps = 0; _ps < 16; _ps++) + _tmem_load_4[offset + _ps] *= inv_sum * output_scale; + #endif + { + __nv_bfloat162 _pk0 = __floats2bfloat162_rn(_tmem_load_4[offset + 0], _tmem_load_4[offset + 1]); + unsigned _pk_u0 = *reinterpret_cast(&_pk0); + __nv_bfloat162 _pk1 = __floats2bfloat162_rn(_tmem_load_4[offset + 2], _tmem_load_4[offset + 3]); + unsigned _pk_u1 = *reinterpret_cast(&_pk1); + __nv_bfloat162 _pk2 = __floats2bfloat162_rn(_tmem_load_4[offset + 4], _tmem_load_4[offset + 5]); + unsigned _pk_u2 = *reinterpret_cast(&_pk2); + __nv_bfloat162 _pk3 = __floats2bfloat162_rn(_tmem_load_4[offset + 6], _tmem_load_4[offset + 7]); + unsigned _pk_u3 = *reinterpret_cast(&_pk3); + __nv_bfloat162 _pk4 = __floats2bfloat162_rn(_tmem_load_4[offset + 8], _tmem_load_4[offset + 9]); + unsigned _pk_u4 = *reinterpret_cast(&_pk4); + __nv_bfloat162 _pk5 = __floats2bfloat162_rn(_tmem_load_4[offset + 10], _tmem_load_4[offset + 11]); + unsigned _pk_u5 = *reinterpret_cast(&_pk5); + __nv_bfloat162 _pk6 = __floats2bfloat162_rn(_tmem_load_4[offset + 12], _tmem_load_4[offset + 13]); + unsigned _pk_u6 = *reinterpret_cast(&_pk6); + __nv_bfloat162 _pk7 = __floats2bfloat162_rn(_tmem_load_4[offset + 14], _tmem_load_4[offset + 15]); + unsigned _pk_u7 = *reinterpret_cast(&_pk7); + asm volatile( + "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"((void*)(&((__nv_bfloat16*)(O + (output_base + col_base + offset)))[0])), "r"(_pk_u0), "r"(_pk_u1), "r"(_pk_u2), "r"(_pk_u3), "r"(_pk_u4), "r"(_pk_u5), "r"(_pk_u6), "r"(_pk_u7) : "memory"); + } + } + } + } + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(o_empty_addr); + mbarrier_wait(work_full_addr + (work_stage_1) * 8, _phase_work_full_1); + uint32_t _clc_valid_3 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_3) + : "r"(work_response_addr + work_stage_1 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_3 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_3) + : "r"(work_response_addr + work_stage_1 * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(work_empty_addr + (work_stage_1) * 8); + work_stage_1 += 1; + if (work_stage_1 == 2) { work_stage_1 = 0; _phase_work_full_1 ^= 1; } + unsigned int valid_1 = _clc_valid_3; + query_idx_1 = _clc_ctaid_3; + if (valid_1 == 0) { + break; + } + } + asm volatile("barrier.sync 7, 128;" ::: "memory"); + if (warp == 4) { + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + } + } + // ---- Role: mma_warp ---- + if (warp == 8) { + { // mma_warp_main + unsigned int work_stage_2 = 0; + unsigned int query_idx_2 = blockIdx.x; + unsigned int global_tile_2 = 0; + unsigned int _phase_work_full_2 = 0; + #pragma unroll 1 + for (unsigned int work_iter_2 = 0; work_iter_2 < num_query_tokens; work_iter_2++) { + if (query_idx_2 >= (unsigned int)num_query_tokens) { + break; + } + int _max_2 = ((sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) : (0)); + int _min_1 = ((_max_2) < (sparse_topk) ? (_max_2) : (sparse_topk)); + int active_topk_2 = _min_1; + int _max_3 = (((active_topk_2 + 128 - 1) / 128) > (1) ? ((active_topk_2 + 128 - 1) / 128) : (1)); + int num_loop_steps_2 = _max_3; + unsigned int query_phase = work_iter_2 & 1; + mbarrier_wait(q_full_addr, query_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + if (num_loop_steps_2 > 0) { + unsigned int score_stage = global_tile_2 & 1; + int score_col_1 = ((score_stage != 0) ? 128 : 0); + #pragma unroll + for (int k_stage = 0; k_stage < 4; k_stage++) { + mbarrier_wait(kv_full_addr + (k_stage) * 8, 0); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_0 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (k_stage) * 1024); + int _mma_b_lo_0 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (k_stage) * 2048); + { + uint64_t _mma_ss_a_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_0); + uint64_t _mma_ss_b_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_0); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, ((k_stage == 0) ? 0 : 1)); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 506U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 1018U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); + } + } + if (k_stage == 3) { + elect_commit(s_full_addr + (score_stage) * 8); + } + elect_commit(kv_empty_addr + (k_stage) * 8); + } + #pragma unroll 1 + for (int tile_2 = 0; tile_2 < num_loop_steps_2 - 1; tile_2++) { + unsigned int next_global_tile = global_tile_2 + 1; + unsigned int next_score_stage = next_global_tile & 1; + int next_score_col = ((next_score_stage != 0) ? 128 : 0); + #pragma unroll + for (int k_stage_1 = 0; k_stage_1 < 4; k_stage_1++) { + mbarrier_wait(kv_full_addr + (k_stage_1) * 8, 1); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_1 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (k_stage_1) * 1024); + int _mma_b_lo_1 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (k_stage_1) * 2048); + { + uint64_t _mma_ss_a_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_1); + uint64_t _mma_ss_b_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_1); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (next_score_col)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, ((k_stage_1 == 0) ? 0 : 1)); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (next_score_col)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (next_score_col)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (next_score_col)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 506U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 1018U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (next_score_col)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (next_score_col)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (next_score_col)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f16((tmem_tmem + (next_score_col)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); + } + } + if (k_stage_1 == 3) { + elect_commit(s_full_addr + (next_score_stage) * 8); + } + elect_commit(kv_empty_addr + (k_stage_1) * 8); + } + unsigned int sync_stage_1 = global_tile_2 & 1; + unsigned int sync_phase_1 = global_tile_2 >> 1 & 1; + int current_score_col = ((sync_stage_1 != 0) ? 128 : 0); + mbarrier_wait(p_full_addr + (sync_stage_1) * 8, sync_phase_1); + mbarrier_wait(o_empty_addr, 1 - (global_tile_2 & 1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int v_stage_2 = 0; v_stage_2 < 4; v_stage_2++) { + mbarrier_wait(kv_full_addr + (v_stage_2) * 8, 0); + asm volatile("tcgen05.fence::after_thread_sync;"); + int output_addr_2 = 256 + (v_stage_2 & 1) * 128 + ((v_stage_2 >> 1) * 16 << 16); + int _mma_b_lo_2 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (v_stage_2) * 2048); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 dhi, blo, ta, id;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 69272720;\n\t" + "mov.b32 ta, %2;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p0;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "}\n" + :: "r"((tmem_tmem + (output_addr_2))), "r"(_mma_b_lo_2), "r"(tmem_tmem + (current_score_col + 32 + ((v_stage_2 >> 1) * 16 << 16))), "r"(((tile_2 == 0) ? 0 : 1))); + elect_commit(kv_empty_addr + (v_stage_2) * 8); + } + elect_commit(o_full_addr); + global_tile_2 = next_global_tile; + } + elect_commit(q_empty_addr); + int last_tile = num_loop_steps_2 - 1; + unsigned int last_sync_stage = global_tile_2 & 1; + unsigned int last_sync_phase = global_tile_2 >> 1 & 1; + int last_score_col = ((last_sync_stage != 0) ? 128 : 0); + mbarrier_wait(p_full_addr + (last_sync_stage) * 8, last_sync_phase); + mbarrier_wait(o_empty_addr, 1 - (global_tile_2 & 1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int v_stage_3 = 0; v_stage_3 < 4; v_stage_3++) { + mbarrier_wait(kv_full_addr + (v_stage_3) * 8, 1); + asm volatile("tcgen05.fence::after_thread_sync;"); + int output_addr_3 = 256 + (v_stage_3 & 1) * 128 + ((v_stage_3 >> 1) * 16 << 16); + int _mma_b_lo_3 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (v_stage_3) * 2048); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 dhi, blo, ta, id;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 69272720;\n\t" + "mov.b32 ta, %2;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p0;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" + "}\n" + :: "r"((tmem_tmem + (output_addr_3))), "r"(_mma_b_lo_3), "r"(tmem_tmem + (last_score_col + 32 + ((v_stage_3 >> 1) * 16 << 16))), "r"(((last_tile == 0) ? 0 : 1))); + elect_commit(kv_empty_addr + (v_stage_3) * 8); + } + elect_commit(o_full_addr); + global_tile_2 = global_tile_2 + 1; + } else { + elect_commit(q_empty_addr); + } + mbarrier_wait(work_full_addr + (work_stage_2) * 8, _phase_work_full_2); + uint32_t _clc_valid_1 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_1) + : "r"(work_response_addr + work_stage_2 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_1 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_1) + : "r"(work_response_addr + work_stage_2 * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(work_empty_addr + (work_stage_2) * 8); + work_stage_2 += 1; + if (work_stage_2 == 2) { work_stage_2 = 0; _phase_work_full_2 ^= 1; } + unsigned int valid_2 = _clc_valid_1; + query_idx_2 = _clc_ctaid_1; + if (valid_2 == 0) { + break; + } + } + } + } + // ---- Role: page_warp ---- + if (warp == 9) { + { // page_warp_main + unsigned int page_stage = 0; + unsigned int work_stage_3 = 0; + unsigned int query_idx_3 = blockIdx.x; + unsigned int _phase_page_empty = 1; + unsigned int _phase_work_full_3 = 0; + #pragma unroll 1 + for (unsigned int work_iter_3 = 0; work_iter_3 < num_query_tokens; work_iter_3++) { + if (query_idx_3 >= (unsigned int)num_query_tokens) { + break; + } + int* swa_row = swa_indices + (query_idx_3 * (unsigned int)swa_index_stride); + int* compressed_row = compressed_indices + (query_idx_3 * (unsigned int)compressed_index_stride); + int _max_11 = ((sparse_topk_lens[query_idx_3] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_3] + sparse_topk_lens_offset) : (0)); + int _min_4 = ((_max_11) < (sparse_topk) ? (_max_11) : (sparse_topk)); + int active_topk_3 = _min_4; + int _max_12 = (((active_topk_3 + 128 - 1) / 128) > (1) ? ((active_topk_3 + 128 - 1) / 128) : (1)); + int num_loop_steps_3 = _max_12; + int num_page_pairs = (num_loop_steps_3 + 1) / 2; + int _max_13 = ((active_topk_3 - 1 & -4) > (0) ? (active_topk_3 - 1 & -4) : (0)); + int last_aligned = _max_13; + #pragma unroll 1 + for (int pair = 0; pair < num_page_pairs; pair++) { + int pair_base = pair * 256; + #pragma unroll + for (int duplicate = 0; duplicate < 2; duplicate++) { + mbarrier_wait(page_empty_addr + (page_stage) * 8, _phase_page_empty); + int lane_base = lane * 4; + int page0 = pair_base + lane_base; + int page1 = page0 + 128; + int safe0 = ((page0 <= last_aligned) ? page0 : last_aligned); + int safe1 = ((page1 <= last_aligned) ? page1 : last_aligned); + int* src0 = ((safe0 < 128) ? swa_row + safe0 : compressed_row + (safe0 - 128)); + int* src1 = ((safe1 < 128) ? swa_row + safe1 : compressed_row + (safe1 - 128)); + asm volatile("cp.async.ca.shared::cta.global [%0], [%1], 16;" + :: "r"(smem_page_offsets_addr + page_stage * 1024 + (unsigned int)(lane_base * 4)), "l"(src0)); + asm volatile("cp.async.ca.shared::cta.global [%0], [%1], 16;" + :: "r"(smem_page_offsets_addr + page_stage * 1024 + (unsigned int)((128 + lane_base) * 4)), "l"(src1)); + asm volatile( + "{\n\t" + "cp.async.mbarrier.arrive.noinc.shared::cta.b64 [%0];\n\t" + "}" + :: "r"(page_full_addr + (page_stage) * 8) : "memory"); + page_stage += 1; + if (page_stage == 6) { page_stage = 0; _phase_page_empty ^= 1; } + } + } + mbarrier_wait(work_full_addr + (work_stage_3) * 8, _phase_work_full_3); + uint32_t _clc_valid_4 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_4) + : "r"(work_response_addr + work_stage_3 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_4 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_4) + : "r"(work_response_addr + work_stage_3 * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(work_empty_addr + (work_stage_3) * 8); + work_stage_3 += 1; + if (work_stage_3 == 2) { work_stage_3 = 0; _phase_work_full_3 ^= 1; } + unsigned int valid_3 = _clc_valid_4; + query_idx_3 = _clc_ctaid_4; + if (valid_3 == 0) { + break; + } + } + } + } + // ---- Role: scheduler_warp ---- + if (warp == 10) { + { // scheduler_warp_main + unsigned int work_stage_4 = 0; + unsigned int throttle_stage = 0; + unsigned int query_idx_4 = blockIdx.x; + unsigned int _phase_throttle_full = 0; + unsigned int _phase_work_empty = 1; + unsigned int _phase_work_full_4 = 0; + #pragma unroll 1 + for (unsigned int work_iter_4 = 0; work_iter_4 < num_query_tokens; work_iter_4++) { + if (query_idx_4 >= (unsigned int)num_query_tokens) { + break; + } + mbarrier_wait(throttle_full_addr + (throttle_stage) * 8, _phase_throttle_full); + mbarrier_arrive(throttle_empty_addr + (throttle_stage) * 8); + throttle_stage += 1; + if (throttle_stage == 2) { throttle_stage = 0; _phase_throttle_full ^= 1; } + if (elect_sync()) { + mbarrier_wait(work_empty_addr + (work_stage_4) * 8, _phase_work_empty); + mbarrier_arrive_expect_tx(work_full_addr + (work_stage_4) * 8, 16); + asm volatile( + "fence.proxy.async.shared::cta;\n\t" + "clusterlaunchcontrol.try_cancel.async.shared::cta" + ".mbarrier::complete_tx::bytes.b128" + " [%0], [%1];" + :: "r"(work_response_addr + work_stage_4 * 16 + 0 * 16), "r"(work_full_addr + work_stage_4 * 8) + : "memory"); + } + mbarrier_wait(work_full_addr + (work_stage_4) * 8, _phase_work_full_4); + uint32_t _clc_valid_5 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_5) + : "r"(work_response_addr + work_stage_4 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_5 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_5) + : "r"(work_response_addr + work_stage_4 * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(work_empty_addr + (work_stage_4) * 8); + work_stage_4 += 1; + if (work_stage_4 == 2) { work_stage_4 = 0; _phase_work_empty ^= 1; _phase_work_full_4 ^= 1; } + unsigned int valid_4 = _clc_valid_5; + query_idx_4 = _clc_ctaid_5; + if (valid_4 == 0) { + break; + } + } + } + } + // ---- Role: padding_warp ---- + if (warp == 11) { + { // padding_warp_main + unsigned int work_stage_5 = 0; + unsigned int query_idx_5 = blockIdx.x; + unsigned int _phase_work_full_5 = 0; + #pragma unroll 1 + for (unsigned int work_iter_5 = 0; work_iter_5 < num_query_tokens; work_iter_5++) { + if (query_idx_5 >= (unsigned int)num_query_tokens) { + break; + } + mbarrier_wait(work_full_addr + (work_stage_5) * 8, _phase_work_full_5); + uint32_t _clc_valid_6 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_6) + : "r"(work_response_addr + work_stage_5 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_6 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_6) + : "r"(work_response_addr + work_stage_5 * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(work_empty_addr + (work_stage_5) * 8); + work_stage_5 += 1; + if (work_stage_5 == 2) { work_stage_5 = 0; _phase_work_full_5 ^= 1; } + unsigned int valid_5 = _clc_valid_6; + query_idx_5 = _clc_ctaid_6; + if (valid_5 == 0) { + break; + } + } + } + } + // ---- Role: load_warp ---- + if (warp >= 12 && warp <= 15) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 72;"); + { // load_warp_main + asm volatile("griddepcontrol.wait;" ::: "memory"); + const int load_warp_rank = warp - 12; + unsigned int load_page_stage = 0; + unsigned int page_k_stage = 0; + unsigned int page_v_stage = 0; + int page_k_addr = 0; + int page_v_addr = 0; + unsigned int work_stage_6 = 0; + unsigned int throttle_stage_1 = 0; + unsigned int query_idx_6 = blockIdx.x; + unsigned int _phase_throttle_empty = 1; + unsigned int _phase_page_full = 0; + unsigned int _phase_work_full_6 = 0; + #pragma unroll 1 + for (unsigned int work_iter_6 = 0; work_iter_6 < num_query_tokens; work_iter_6++) { + if (query_idx_6 >= (unsigned int)num_query_tokens) { + break; + } + int _max_0 = ((sparse_topk_lens[query_idx_6] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_6] + sparse_topk_lens_offset) : (0)); + int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); + int active_topk_4 = _min_0; + int _max_1 = (((active_topk_4 + 128 - 1) / 128) > (1) ? ((active_topk_4 + 128 - 1) / 128) : (1)); + int num_loop_steps_4 = _max_1; + mbarrier_wait(throttle_empty_addr + (throttle_stage_1) * 8, _phase_throttle_empty); + mbarrier_arrive(throttle_full_addr + (throttle_stage_1) * 8); + throttle_stage_1 += 1; + if (throttle_stage_1 == 2) { throttle_stage_1 = 0; _phase_throttle_empty ^= 1; } + unsigned int query_phase_1 = work_iter_6 & 1; + mbarrier_wait(q_empty_addr, 1 - query_phase_1); + if (warp == 12) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(q_full_addr, 65536); + #pragma unroll + for (int q_subload = 0; q_subload < 8; q_subload++) { + tma_4d_gmem2smem(smem_q_addr + (unsigned int)(q_subload * 8192), (&tmap_q), 0, 0, q_subload, query_idx_6, q_full_addr); + } + } + } + if (num_loop_steps_4 > 0) { + mbarrier_wait(page_full_addr + (load_page_stage) * 8, _phase_page_full); + page_k_stage = load_page_stage; + page_k_addr = smem_page_offsets_addr + page_k_stage * 1024; + load_page_stage += 1; + if (load_page_stage == 6) { load_page_stage = 0; _phase_page_full ^= 1; } + #pragma unroll + for (int kv_stage = 0; kv_stage < 4; kv_stage++) { + mbarrier_wait(kv_empty_addr + (kv_stage) * 8, 1); + if (warp == 12) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage) * 8, 32768); + } + } + #pragma unroll + for (int group = 0; group < 8; group++) { + int row_offset = load_warp_rank * 4 + group * 16; + int raw_rows[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows[0])), "=r"(*reinterpret_cast(&raw_rows[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows[(0) + 3])) + : "r"(page_k_addr + row_offset * 4)); + int row0 = raw_rows[0]; + int row1 = raw_rows[1]; + int row2 = raw_rows[2]; + int row3 = raw_rows[3]; + int dst = smem_kv_addr + (unsigned int)(kv_stage * 32768) + (unsigned int)(row_offset * 64 * 2); + if (elect_sync()) { + { + tma_gather4_gmem2smem(dst, (&tmap_swa_kv), kv_stage * 128, row0, row1, row2, row3, kv_full_addr + (kv_stage) * 8); + tma_gather4_gmem2smem(dst + 16384, (&tmap_swa_kv), kv_stage * 128 + 64, row0, row1, row2, row3, kv_full_addr + (kv_stage) * 8); + } + } + } + } + if (num_loop_steps_4 > 1) { + #pragma unroll + for (int kv_stage_1 = 0; kv_stage_1 < 4; kv_stage_1++) { + mbarrier_wait(kv_empty_addr + (kv_stage_1) * 8, 0); + if (warp == 12) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_1) * 8, 32768); + } + } + #pragma unroll + for (int group_1 = 0; group_1 < 8; group_1++) { + int row_offset_1 = load_warp_rank * 4 + group_1 * 16; + int raw_rows_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_1[0])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 3])) + : "r"(page_k_addr + 512 + row_offset_1 * 4)); + int row0_1 = raw_rows_1[0]; + int row1_1 = raw_rows_1[1]; + int row2_1 = raw_rows_1[2]; + int row3_1 = raw_rows_1[3]; + int dst_1 = smem_kv_addr + (unsigned int)(kv_stage_1 * 32768) + (unsigned int)(row_offset_1 * 64 * 2); + if (elect_sync()) { + { + tma_gather4_gmem2smem(dst_1, (&tmap_compressed_kv), kv_stage_1 * 128, row0_1, row1_1, row2_1, row3_1, kv_full_addr + (kv_stage_1) * 8); + tma_gather4_gmem2smem(dst_1 + 16384, (&tmap_compressed_kv), kv_stage_1 * 128 + 64, row0_1, row1_1, row2_1, row3_1, kv_full_addr + (kv_stage_1) * 8); + } + } + } + } + if (num_loop_steps_4 > 2) { + mbarrier_arrive(page_empty_addr + (page_k_stage) * 8); + } + mbarrier_wait(page_full_addr + (load_page_stage) * 8, _phase_page_full); + page_v_stage = load_page_stage; + page_v_addr = smem_page_offsets_addr + page_v_stage * 1024; + load_page_stage += 1; + if (load_page_stage == 6) { load_page_stage = 0; _phase_page_full ^= 1; } + #pragma unroll + for (int kv_stage_2 = 0; kv_stage_2 < 4; kv_stage_2++) { + mbarrier_wait(kv_empty_addr + (kv_stage_2) * 8, 1); + if (warp == 12) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_2) * 8, 32768); + } + } + #pragma unroll + for (int group_2 = 0; group_2 < 8; group_2++) { + int row_offset_2 = load_warp_rank * 4 + group_2 * 16; + int raw_rows_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_2[0])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 3])) + : "r"(page_v_addr + row_offset_2 * 4)); + int row0_2 = raw_rows_2[0]; + int row1_2 = raw_rows_2[1]; + int row2_2 = raw_rows_2[2]; + int row3_2 = raw_rows_2[3]; + int dst_2 = smem_v_addr + (unsigned int)(kv_stage_2 * 32768) + (unsigned int)(row_offset_2 * 64 * 2); + if (elect_sync()) { + { + tma_gather4_gmem2smem(dst_2, (&tmap_swa_kv), kv_stage_2 * 128, row0_2, row1_2, row2_2, row3_2, kv_full_addr + (kv_stage_2) * 8); + tma_gather4_gmem2smem(dst_2 + 16384, (&tmap_swa_kv), kv_stage_2 * 128 + 64, row0_2, row1_2, row2_2, row3_2, kv_full_addr + (kv_stage_2) * 8); + } + } + } + } + } + #pragma unroll 1 + for (int tile_3 = 1; tile_3 < num_loop_steps_4 - 1; tile_3++) { + int next_tile = tile_3 + 1; + if ((next_tile & 1) == 0) { + mbarrier_wait(page_full_addr + (load_page_stage) * 8, _phase_page_full); + page_k_stage = load_page_stage; + page_k_addr = smem_page_offsets_addr + page_k_stage * 1024; + load_page_stage += 1; + if (load_page_stage == 6) { load_page_stage = 0; _phase_page_full ^= 1; } + } + #pragma unroll + for (int kv_stage_3 = 0; kv_stage_3 < 4; kv_stage_3++) { + mbarrier_wait(kv_empty_addr + (kv_stage_3) * 8, 0); + if (warp == 12) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_3) * 8, 32768); + } + } + #pragma unroll + for (int group_3 = 0; group_3 < 8; group_3++) { + int row_offset_3 = load_warp_rank * 4 + group_3 * 16; + int raw_rows_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_3[0])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 3])) + : "r"(page_k_addr + (next_tile & 1) * 128 * 4 + row_offset_3 * 4)); + int row0_3 = raw_rows_3[0]; + int row1_3 = raw_rows_3[1]; + int row2_3 = raw_rows_3[2]; + int row3_3 = raw_rows_3[3]; + int dst_3 = smem_kv_addr + (unsigned int)(kv_stage_3 * 32768) + (unsigned int)(row_offset_3 * 64 * 2); + if (elect_sync()) { + { + tma_gather4_gmem2smem(dst_3, (&tmap_compressed_kv), kv_stage_3 * 128, row0_3, row1_3, row2_3, row3_3, kv_full_addr + (kv_stage_3) * 8); + tma_gather4_gmem2smem(dst_3 + 16384, (&tmap_compressed_kv), kv_stage_3 * 128 + 64, row0_3, row1_3, row2_3, row3_3, kv_full_addr + (kv_stage_3) * 8); + } + } + } + } + if ((next_tile & 1) != 0 && tile_3 < num_loop_steps_4 - 2) { + mbarrier_arrive(page_empty_addr + (page_k_stage) * 8); + } + if ((tile_3 & 1) == 0) { + mbarrier_wait(page_full_addr + (load_page_stage) * 8, _phase_page_full); + page_v_stage = load_page_stage; + page_v_addr = smem_page_offsets_addr + page_v_stage * 1024; + load_page_stage += 1; + if (load_page_stage == 6) { load_page_stage = 0; _phase_page_full ^= 1; } + } + #pragma unroll + for (int kv_stage_4 = 0; kv_stage_4 < 4; kv_stage_4++) { + mbarrier_wait(kv_empty_addr + (kv_stage_4) * 8, 1); + if (warp == 12) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_4) * 8, 32768); + } + } + #pragma unroll + for (int group_4 = 0; group_4 < 8; group_4++) { + int row_offset_4 = load_warp_rank * 4 + group_4 * 16; + int raw_rows_4[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_4[0])), "=r"(*reinterpret_cast(&raw_rows_4[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_4[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_4[(0) + 3])) + : "r"(page_v_addr + (tile_3 & 1) * 128 * 4 + row_offset_4 * 4)); + int row0_4 = raw_rows_4[0]; + int row1_4 = raw_rows_4[1]; + int row2_4 = raw_rows_4[2]; + int row3_4 = raw_rows_4[3]; + int dst_4 = smem_v_addr + (unsigned int)(kv_stage_4 * 32768) + (unsigned int)(row_offset_4 * 64 * 2); + if (elect_sync()) { + { + tma_gather4_gmem2smem(dst_4, (&tmap_compressed_kv), kv_stage_4 * 128, row0_4, row1_4, row2_4, row3_4, kv_full_addr + (kv_stage_4) * 8); + tma_gather4_gmem2smem(dst_4 + 16384, (&tmap_compressed_kv), kv_stage_4 * 128 + 64, row0_4, row1_4, row2_4, row3_4, kv_full_addr + (kv_stage_4) * 8); + } + } + } + } + if ((tile_3 & 1) != 0) { + mbarrier_arrive(page_empty_addr + (page_v_stage) * 8); + } + } + int last_tile_1 = num_loop_steps_4 - 1; + mbarrier_arrive(page_empty_addr + (page_k_stage) * 8); + if ((last_tile_1 & 1) == 0) { + mbarrier_wait(page_full_addr + (load_page_stage) * 8, _phase_page_full); + page_v_stage = load_page_stage; + page_v_addr = smem_page_offsets_addr + page_v_stage * 1024; + load_page_stage += 1; + if (load_page_stage == 6) { load_page_stage = 0; _phase_page_full ^= 1; } + } + if (num_loop_steps_4 == 1) { + #pragma unroll + for (int kv_stage_5 = 0; kv_stage_5 < 4; kv_stage_5++) { + mbarrier_wait(kv_empty_addr + (kv_stage_5) * 8, 0); + if (warp == 12) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_5) * 8, 32768); + } + } + #pragma unroll + for (int group_5 = 0; group_5 < 8; group_5++) { + int row_offset_5 = load_warp_rank * 4 + group_5 * 16; + int raw_rows_5[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_5[0])), "=r"(*reinterpret_cast(&raw_rows_5[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_5[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_5[(0) + 3])) + : "r"(page_v_addr + row_offset_5 * 4)); + int row0_5 = raw_rows_5[0]; + int row1_5 = raw_rows_5[1]; + int row2_5 = raw_rows_5[2]; + int row3_5 = raw_rows_5[3]; + int dst_5 = smem_v_addr + (unsigned int)(kv_stage_5 * 32768) + (unsigned int)(row_offset_5 * 64 * 2); + if (elect_sync()) { + { + tma_gather4_gmem2smem(dst_5, (&tmap_swa_kv), kv_stage_5 * 128, row0_5, row1_5, row2_5, row3_5, kv_full_addr + (kv_stage_5) * 8); + tma_gather4_gmem2smem(dst_5 + 16384, (&tmap_swa_kv), kv_stage_5 * 128 + 64, row0_5, row1_5, row2_5, row3_5, kv_full_addr + (kv_stage_5) * 8); + } + } + } + } + } else { + #pragma unroll + for (int kv_stage_6 = 0; kv_stage_6 < 4; kv_stage_6++) { + mbarrier_wait(kv_empty_addr + (kv_stage_6) * 8, 0); + if (warp == 12) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_6) * 8, 32768); + } + } + #pragma unroll + for (int group_6 = 0; group_6 < 8; group_6++) { + int row_offset_6 = load_warp_rank * 4 + group_6 * 16; + int raw_rows_6[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_6[0])), "=r"(*reinterpret_cast(&raw_rows_6[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_6[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_6[(0) + 3])) + : "r"(page_v_addr + (last_tile_1 & 1) * 128 * 4 + row_offset_6 * 4)); + int row0_6 = raw_rows_6[0]; + int row1_6 = raw_rows_6[1]; + int row2_6 = raw_rows_6[2]; + int row3_6 = raw_rows_6[3]; + int dst_6 = smem_v_addr + (unsigned int)(kv_stage_6 * 32768) + (unsigned int)(row_offset_6 * 64 * 2); + if (elect_sync()) { + { + tma_gather4_gmem2smem(dst_6, (&tmap_compressed_kv), kv_stage_6 * 128, row0_6, row1_6, row2_6, row3_6, kv_full_addr + (kv_stage_6) * 8); + tma_gather4_gmem2smem(dst_6 + 16384, (&tmap_compressed_kv), kv_stage_6 * 128 + 64, row0_6, row1_6, row2_6, row3_6, kv_full_addr + (kv_stage_6) * 8); + } + } + } + } + } + mbarrier_arrive(page_empty_addr + (page_v_stage) * 8); + } + mbarrier_wait(work_full_addr + (work_stage_6) * 8, _phase_work_full_6); + uint32_t _clc_valid_0 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_0) + : "r"(work_response_addr + work_stage_6 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_0 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_0) + : "r"(work_response_addr + work_stage_6 * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(work_empty_addr + (work_stage_6) * 8); + work_stage_6 += 1; + if (work_stage_6 == 2) { work_stage_6 = 0; _phase_work_full_6 ^= 1; } + unsigned int valid_6 = _clc_valid_0; + query_idx_6 = _clc_ctaid_0; + if (valid_6 == 0) { + break; + } + } + } + } + + // Cleanup +} + +} // extern "C" diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_21a93d40caefcf5b252d_binding.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_21a93d40caefcf5b252d_binding.cu new file mode 100644 index 00000000000..5accf360b11 --- /dev/null +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_21a93d40caefcf5b252d_binding.cu @@ -0,0 +1,446 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +// Generated by the Cake source exporter — do not edit. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_21a93d40caefcf5b252d'. +#include +#include +#include +#include + +#include "tvm_ffi_utils.h" + +#include +#include +#include +#include +#include +#include +#include + +extern "C" __global__ void kernel_cake_dsv4_21a93d40caefcf5b252d(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int max_q_len, int batch_size, int has_sinks); + + +namespace cake_host_shim_443c43e36fe3b626 { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t CakeDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +// 4D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + int64_t d3 = t.size(t.ndim() - 3); + TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + TVM_FFI_CHECK(d1 % 128 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 128; + uint64_t global_dim[4] = {(uint64_t)(128), (uint64_t)(d2), (uint64_t)((d1 / 128)), (uint64_t)(d3)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (128, 8, 1, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 128; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + (uint64_t)((carrier_stride_1 * 8) / 8), + (uint64_t)((carrier_stride_2 * 8) / 8), + }; + uint32_t box_dim[4] = {128u, 8u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_swa_kv' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_swa_kv' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; + TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (128, 1) exceeds resolved global dims for 'tmap_swa_kv'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + }; + uint32_t box_dim[2] = {128u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_compressed_kv' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_compressed_kv(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_compressed_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_compressed_kv' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_compressed_kv' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_compressed_kv"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_compressed_kv' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved a non-positive global dim"; + TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (128, 1) exceeds resolved global dims for 'tmap_compressed_kv'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + }; + uint32_t box_dim[2] = {128u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_kv') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_compressed_kv, TensorView arg_O, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_seq_lens, TensorView arg_cum_seq_lens_q, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_max_q_len, int64_t arg_batch_size, int64_t arg_has_sinks, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 1, 8, 1); + CheckContiguous(arg_tmap_q, "tmap_q"); + CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); + CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 1, 8, 1); + CheckCudaTensor(arg_tmap_compressed_kv, "tmap_compressed_kv"); + CheckDtype(arg_tmap_compressed_kv, "tmap_compressed_kv", 1, 8, 1); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_swa_indices, "swa_indices"); + CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); + CheckContiguous(arg_swa_indices, "swa_indices"); + CheckCudaTensor(arg_compressed_indices, "compressed_indices"); + CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); + CheckContiguous(arg_compressed_indices, "compressed_indices"); + CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); + CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckCudaTensor(arg_seq_lens, "seq_lens"); + CheckDtype(arg_seq_lens, "seq_lens", 0, 32, 1); + CheckContiguous(arg_seq_lens, "seq_lens"); + CheckCudaTensor(arg_cum_seq_lens_q, "cum_seq_lens_q"); + CheckDtype(arg_cum_seq_lens_q, "cum_seq_lens_q", 0, 32, 1); + CheckContiguous(arg_cum_seq_lens_q, "cum_seq_lens_q"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); + CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); + CheckContiguous(arg_bmm1_scale, "bmm1_scale"); + CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); + CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); + CheckContiguous(arg_bmm2_scale, "bmm2_scale"); + TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) + << "scalar 'swa_index_stride' value " << arg_swa_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) + << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) + << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) + << "scalar 'sparse_topk' value " << arg_sparse_topk + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_q_len >= -2147483648LL && arg_max_q_len <= 2147483647LL, ValueError) + << "scalar 'max_q_len' value " << arg_max_q_len + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_batch_size >= -2147483648LL && arg_batch_size <= 2147483647LL, ValueError) + << "scalar 'batch_size' value " << arg_batch_size + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) + << "scalar 'has_sinks' value " << arg_has_sinks + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); + CheckSameCudaDevice(arg_tmap_compressed_kv, arg_tmap_q, "tmap_compressed_kv", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); + CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); + CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); + CheckSameCudaDevice(arg_seq_lens, arg_tmap_q, "seq_lens", "tmap_q"); + CheckSameCudaDevice(arg_cum_seq_lens_q, arg_tmap_q, "cum_seq_lens_q", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); + CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + + cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); + CUtensorMap p_tmap_compressed_kv = EncodeTma_tmap_compressed_kv(arg_tmap_compressed_kv); + __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); + int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); + int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); + int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); + int* p_seq_lens = static_cast(arg_seq_lens.data_ptr()); + int* p_cum_seq_lens_q = static_cast(arg_cum_seq_lens_q.data_ptr()); + float* p_sinks = static_cast(arg_sinks.data_ptr()); + float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); + float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); + int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; + int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; + int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; + int32_t v_sparse_topk = (int32_t)arg_sparse_topk; + int32_t v_max_q_len = (int32_t)arg_max_q_len; + int32_t v_batch_size = (int32_t)arg_batch_size; + int32_t v_has_sinks = (int32_t)arg_has_sinks; + void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_tmap_compressed_kv, &p_O, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_seq_lens, &p_cum_seq_lens_q, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_max_q_len, &v_batch_size, &v_has_sinks}; + + static std::mutex smem_mu; + static auto* smem_status_by_device = + new std::unordered_map(); + cudaError_t smem_status = cudaSuccess; + { + std::lock_guard lock(smem_mu); + auto it = smem_status_by_device->find(dev.device_id); + if (it == smem_status_by_device->end()) { + int current_device = -1; + smem_status = cudaGetDevice(¤t_device); + if (smem_status == cudaSuccess && current_device != dev.device_id) { + smem_status = cudaErrorInvalidDevice; + } + if (smem_status == cudaSuccess) { + smem_status = cudaFuncSetAttribute( + reinterpret_cast(kernel_cake_dsv4_21a93d40caefcf5b252d), + cudaFuncAttributeMaxDynamicSharedMemorySize, + 147072); + } + smem_status_by_device->emplace(dev.device_id, smem_status); + } else { + smem_status = it->second; + } + } + TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) + << "cudaFuncSetAttribute for kernel_cake_dsv4_21a93d40caefcf5b252d failed: " + << cudaGetErrorString(smem_status); + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(640u, 1u, 1u); + + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes = 147072u; + config.stream = stream; + cudaError_t launch_status = cudaLaunchKernelExC( + &config, reinterpret_cast(kernel_cake_dsv4_21a93d40caefcf5b252d), kargs); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernelExC for kernel_cake_dsv4_21a93d40caefcf5b252d failed: " + << cudaGetErrorString(launch_status); +} + +} // namespace cake_host_shim_443c43e36fe3b626 + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_443c43e36fe3b626::Run); diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_21a93d40caefcf5b252d_kernel.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_21a93d40caefcf5b252d_kernel.cu new file mode 100644 index 00000000000..bc691ca97b1 --- /dev/null +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_21a93d40caefcf5b252d_kernel.cu @@ -0,0 +1,2646 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +#define CAKE_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SOURCE_EXACT_SCRATCH_OFFSET 0 +#define NUM_Q_PIPE_STAGES 1 +#define NUM_KV_PIPE_STAGES 8 +#define NUM_PAGE_PIPE_STAGES 6 +#define NUM_SCORE_PIPE_STAGES 2 +#define NUM_STATS_PIPE_STAGES 2 +#define NUM_P_PIPE_STAGES 2 +#define NUM_O_PIPE_STAGES 1 +#define NUM_WORK_PIPE_STAGES 2 +#define NUM_THROTTLE_PIPE_STAGES 2 +#define SMEM_SMEM_Q_OFF 1024 +#define SMEM_SMEM_Q_STAGE_BYTES 1024 +#define SMEM_SMEM_Q_STRIDE 1024 +#define SMEM_SMEM_KV_OFF 5120 +#define SMEM_SMEM_KV_STAGE_BYTES 16384 +#define SMEM_SMEM_KV_STRIDE 16384 +#define SMEM_SMEM_V_OFF 5120 +#define SMEM_SMEM_V_STAGE_BYTES 16384 +#define SMEM_SMEM_V_STRIDE 16384 +#define SMEM_SMEM_P_OFF 136192 +#define SMEM_SMEM_P_STAGE_BYTES 1024 +#define SMEM_SMEM_P_STRIDE 1024 +#define SMEM_SMEM_PAGE_OFFSETS_OFF 138240 +#define SMEM_SMEM_PAGE_OFFSETS_STAGE_BYTES 1024 +#define SMEM_SMEM_PAGE_OFFSETS_STRIDE 1024 +#define SMEM_SMEM_O_OFF 144384 +#define SMEM_SMEM_O_STAGE_BYTES 2048 +#define SMEM_SMEM_O_STRIDE 2048 +#define SMEM_SMEM_SOFTMAX_OFF 146432 +#define SMEM_SMEM_SOFTMAX_STAGE_BYTES 128 +#define SMEM_SMEM_SOFTMAX_STRIDE 128 +#define SMEM_SMEM_SOFTMAX_U32_OFF 146432 +#define SMEM_SMEM_SOFTMAX_U32_STAGE_BYTES 128 +#define SMEM_SMEM_SOFTMAX_U32_STRIDE 128 +#define SMEM_SMEM_SOFTMAX_UNUSED_OFF 146560 +#define SMEM_SMEM_SOFTMAX_UNUSED_STAGE_BYTES 128 +#define SMEM_SMEM_SOFTMAX_UNUSED_STRIDE 128 +#define SMEM_SMEM_CORR_OFF 146688 +#define SMEM_SMEM_CORR_STAGE_BYTES 128 +#define SMEM_SMEM_CORR_STRIDE 128 +#define SMEM_SMEM_STATS_OFF 146816 +#define SMEM_SMEM_STATS_STAGE_BYTES 96 +#define SMEM_SMEM_STATS_STRIDE 96 +#define SMEM_WORK_RESPONSE_OFF 147008 +#define SMEM_WORK_RESPONSE_STAGE_BYTES 16 +#define SMEM_WORK_RESPONSE_STRIDE 16 +#define SMEM_TOTAL 147072 +#define THREADS 640 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count) : "memory"); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f8f6f4( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x4(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x4.b32" + " {%0, %1, %2, %3}, [%4];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_2d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3}], [%4];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr) { + // Canonical .shared::cta form for non-multicast gather4, matching + // trtllm-gen / cuda_ptx and the PTX ISA qualifier order + // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, + // .shared::cluster is reserved for the multicast variant (ctaMask). + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x4_wait(float* dst, int addr) { + tmem_ld_x4(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(640, 1) void +kernel_cake_dsv4_21a93d40caefcf5b252d(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int max_q_len, int batch_size, int has_sinks) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); + + const int mbar_base = smem; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define kv_full_addr (mbar_base + 16) + #define kv_empty_addr (mbar_base + 80) + #define page_full_addr (mbar_base + 144) + #define page_empty_addr (mbar_base + 192) + #define score_full_addr (mbar_base + 240) + #define score_empty_addr (mbar_base + 256) + #define stats_full_addr (mbar_base + 272) + #define stats_empty_addr (mbar_base + 288) + #define p_full_addr (mbar_base + 304) + #define p_empty_addr (mbar_base + 320) + #define o_full_addr (mbar_base + 336) + #define o_empty_addr (mbar_base + 344) + #define work_full_addr (mbar_base + 352) + #define work_empty_addr (mbar_base + 368) + #define throttle_full_addr (mbar_base + 384) + #define throttle_empty_addr (mbar_base + 400) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + + // Kernel setup ops + uint8_t* smem_q = reinterpret_cast(smem_raw + 1024); + const int smem_q_addr = smem + 1024; + uint8_t* smem_kv = reinterpret_cast(smem_raw + 5120); + const int smem_kv_addr = smem + 5120; + uint8_t* smem_v = reinterpret_cast(smem_raw + 5120); + const int smem_v_addr = smem + 5120; + uint8_t* smem_p = reinterpret_cast(smem_raw + 136192); + const int smem_p_addr = smem + 136192; + int* smem_page_offsets = reinterpret_cast(smem_raw + 138240); + const int smem_page_offsets_addr = smem + 138240; + __nv_bfloat16* smem_o = reinterpret_cast<__nv_bfloat16*>(smem_raw + 144384); + const int smem_o_addr = smem + 144384; + float* smem_softmax = reinterpret_cast(smem_raw + 146432); + const int smem_softmax_addr = smem + 146432; + unsigned int* smem_softmax_u32 = reinterpret_cast(smem_raw + 146432); + const int smem_softmax_u32_addr = smem + 146432; + float* smem_softmax_unused = reinterpret_cast(smem_raw + 146560); + const int smem_softmax_unused_addr = smem + 146560; + float* smem_corr = reinterpret_cast(smem_raw + 146688); + const int smem_corr_addr = smem + 146688; + float* smem_stats = reinterpret_cast(smem_raw + 146816); + const int smem_stats_addr = smem + 146816; + unsigned int* work_response = reinterpret_cast(smem_raw + 147008); + const int work_response_addr = smem + 147008; + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_q))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_swa_kv))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_compressed_kv))) : "memory"); + + // Mbarrier init (18 pipeline groups, 0 ordered-sequence groups, 52 barriers) + // Mbarriers at smem_raw[0..416) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // --- pipeline 'q_pipe' --- + // q_full: 1 barriers, init_count=1 + mbarrier_init(smem + 0, 1); + // q_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 8, 1); + // --- pipeline 'kv_pipe' --- + // kv_full: 8 barriers, init_count=1 + mbarrier_init(smem + 16, 1); + mbarrier_init(smem + 24, 1); + mbarrier_init(smem + 32, 1); + mbarrier_init(smem + 40, 1); + mbarrier_init(smem + 48, 1); + mbarrier_init(smem + 56, 1); + mbarrier_init(smem + 64, 1); + mbarrier_init(smem + 72, 1); + // kv_empty: 8 barriers, init_count=1 + mbarrier_init(smem + 80, 1); + mbarrier_init(smem + 88, 1); + mbarrier_init(smem + 96, 1); + mbarrier_init(smem + 104, 1); + mbarrier_init(smem + 112, 1); + mbarrier_init(smem + 120, 1); + mbarrier_init(smem + 128, 1); + mbarrier_init(smem + 136, 1); + // --- pipeline 'page_pipe' --- + // page_full: 6 barriers, init_count=32 + mbarrier_init(smem + 144, 32); + mbarrier_init(smem + 152, 32); + mbarrier_init(smem + 160, 32); + mbarrier_init(smem + 168, 32); + mbarrier_init(smem + 176, 32); + mbarrier_init(smem + 184, 32); + // page_empty: 6 barriers, init_count=256 + mbarrier_init(smem + 192, 256); + mbarrier_init(smem + 200, 256); + mbarrier_init(smem + 208, 256); + mbarrier_init(smem + 216, 256); + mbarrier_init(smem + 224, 256); + mbarrier_init(smem + 232, 256); + // --- pipeline 'score_pipe' --- + // score_full: 2 barriers, init_count=1 + mbarrier_init(smem + 240, 1); + mbarrier_init(smem + 248, 1); + // score_empty: 2 barriers, init_count=128 + mbarrier_init(smem + 256, 128); + mbarrier_init(smem + 264, 128); + // --- pipeline 'stats_pipe' --- + // stats_full: 2 barriers, init_count=128 + mbarrier_init(smem + 272, 128); + mbarrier_init(smem + 280, 128); + // stats_empty: 2 barriers, init_count=128 + mbarrier_init(smem + 288, 128); + mbarrier_init(smem + 296, 128); + // --- pipeline 'p_pipe' --- + // p_full: 2 barriers, init_count=128 + mbarrier_init(smem + 304, 128); + mbarrier_init(smem + 312, 128); + // p_empty: 2 barriers, init_count=1 + mbarrier_init(smem + 320, 1); + mbarrier_init(smem + 328, 1); + // --- pipeline 'o_pipe' --- + // o_full: 1 barriers, init_count=1 + mbarrier_init(smem + 336, 1); + // o_empty: 1 barriers, init_count=128 + mbarrier_init(smem + 344, 128); + // --- pipeline 'work_pipe' --- + // work_full: 2 barriers, init_count=1 + mbarrier_init(smem + 352, 1); + mbarrier_init(smem + 360, 1); + // work_empty: 2 barriers, init_count=640 + mbarrier_init(smem + 368, 640); + mbarrier_init(smem + 376, 640); + // --- pipeline 'throttle_pipe' --- + // throttle_full: 2 barriers, init_count=256 + mbarrier_init(smem + 384, 256); + mbarrier_init(smem + 392, 256); + // throttle_empty: 2 barriers, init_count=32 + mbarrier_init(smem + 400, 32); + mbarrier_init(smem + 408, 32); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncwarp(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 416); + if (warp == 0) { + int _tmem_hold = smem + 416; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_source_exact_scratch = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 48;"); + } + + // ---- Role: softmax ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 152;"); + { // softmax_main + unsigned int work_stage = 0; + unsigned int score_stage = 0; + unsigned int score_phase = 0; + unsigned int stats_stage = 0; + unsigned int stats_phase = 1; + unsigned int p_stage = 0; + unsigned int p_phase = 1; + unsigned int work_x = blockIdx.x; + unsigned int work_y = blockIdx.y; + unsigned int work_z = blockIdx.z; + const int warp_in_wg = warp; + const int wg_tid = warp_in_wg * 32 + lane; + const int col_pair_base = wg_tid % 4 * 2; + const int ldtm_row_base = warp_in_wg * 32 + lane / 4; + float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; + unsigned int _phase_work_full = 0; + #pragma unroll 1 + for (unsigned int _work = 0; _work < 120; _work++) { + int q_begin = cum_seq_lens_q[work_z]; + int q_end = cum_seq_lens_q[work_z + 1]; + int q_len = q_end - q_begin; + int query_idx = (unsigned int)q_begin + work_x; + int active_topk = 0; + if (work_x < (unsigned int)q_len) { + int _max_4 = ((sparse_topk_lens[query_idx] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx] + sparse_topk_lens_offset) : (0)); + int _min_7 = ((_max_4) < (sparse_topk) ? (_max_4) : (sparse_topk)); + active_topk = _min_7; + } + int num_steps = (active_topk + 128 - 1) / 128; + float row_max_pair[2]; + float row_sum_pair[2]; + row_max_pair[0] = -CAKE_INF; + row_max_pair[1] = -CAKE_INF; + row_sum_pair[0] = 0.0f; + row_sum_pair[1] = 0.0f; + asm volatile("barrier.sync 1, 128;" ::: "memory"); + uint32_t _amf_u_0 = __float_as_uint(-3.4028235e+38f); + uint32_t _amf_mask_0 = -int32_t(_amf_u_0 >> 31) | 0x80000000u; + unsigned int _amf_enc_0 = _amf_u_0 ^ _amf_mask_0; + if (wg_tid < 8) { + smem_softmax_u32[wg_tid] = _amf_enc_0; + } + asm volatile("barrier.sync 1, 128;" ::: "memory"); + #pragma unroll 1 + for (int tile = 0; tile < num_steps; tile++) { + mbarrier_wait(score_full_addr + (score_stage) * 8, score_phase); + int score_col = ((score_stage == 0) ? 0 : 8); + float scores[8]; + float score_lo[4]; + float score_hi[4]; + int tile_end = (tile + 1) * 128; + int is_full_tile = ((active_topk % 128 == 0) ? 1 : 0); + is_full_tile = is_full_tile | ((tile_end < active_topk) ? 1 : 0); + if (is_full_tile != 0) { + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x1.b32" + " {%0, %1, %2, %3}, [%4];" + : "=r"(*reinterpret_cast(&score_lo[0])), "=r"(*reinterpret_cast(&score_lo[1])), "=r"(*reinterpret_cast(&score_lo[2])), "=r"(*reinterpret_cast(&score_lo[3])) + : "r"(taddr + (unsigned int)score_col)); + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x1.b32" + " {%0, %1, %2, %3}, [%4];" + : "=r"(*reinterpret_cast(&score_hi[0])), "=r"(*reinterpret_cast(&score_hi[1])), "=r"(*reinterpret_cast(&score_hi[2])), "=r"(*reinterpret_cast(&score_hi[3])) + : "r"(taddr + (unsigned int)score_col + 1048576)); + #pragma unroll + for (int c = 0; c < 4; c++) { + scores[c] = score_lo[c]; + scores[c + 4] = score_hi[c]; + } + } else { + asm volatile("" ::: "memory"); + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x1.b32" + " {%0, %1, %2, %3}, [%4];" + : "=r"(*reinterpret_cast(&score_lo[0])), "=r"(*reinterpret_cast(&score_lo[1])), "=r"(*reinterpret_cast(&score_lo[2])), "=r"(*reinterpret_cast(&score_lo[3])) + : "r"(taddr + (unsigned int)score_col)); + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x1.b32" + " {%0, %1, %2, %3}, [%4];" + : "=r"(*reinterpret_cast(&score_hi[0])), "=r"(*reinterpret_cast(&score_hi[1])), "=r"(*reinterpret_cast(&score_hi[2])), "=r"(*reinterpret_cast(&score_hi[3])) + : "r"(taddr + (unsigned int)score_col + 1048576)); + #pragma unroll + for (int c_1 = 0; c_1 < 4; c_1++) { + scores[c_1] = score_lo[c_1]; + scores[c_1 + 4] = score_hi[c_1]; + } + if (active_topk <= tile * 128 + ldtm_row_base) { + scores[0] = -CAKE_INF; + scores[1] = -CAKE_INF; + } + if (active_topk <= tile * 128 + ldtm_row_base + 8) { + scores[2] = -CAKE_INF; + scores[3] = -CAKE_INF; + } + if (active_topk <= tile * 128 + ldtm_row_base + 16) { + scores[4] = -CAKE_INF; + scores[5] = -CAKE_INF; + } + if (active_topk <= tile * 128 + ldtm_row_base + 24) { + scores[6] = -CAKE_INF; + scores[7] = -CAKE_INF; + } + asm volatile("" ::: "memory"); + } + float pair_max[2]; + float _max_5 = max_noftz(scores[0], scores[2]); + float _max_6 = max_noftz(scores[4], scores[6]); + float _max_7 = max_noftz(_max_5, _max_6); + pair_max[0] = _max_7; + float _max_8 = max_noftz(scores[1], scores[3]); + float _max_9 = max_noftz(scores[5], scores[7]); + float _max_10 = max_noftz(_max_8, _max_9); + pair_max[1] = _max_10; + #pragma unroll + for (int c_2 = 0; c_2 < 2; c_2++) { + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, pair_max[c_2], 16); + float _max_11 = max_noftz(pair_max[c_2], _shfl_xor_0); + pair_max[c_2] = _max_11; + float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, pair_max[c_2], 8); + float _max_12 = max_noftz(pair_max[c_2], _shfl_xor_1); + pair_max[c_2] = _max_12; + } + float old_max_pair[2]; + float new_max_pair[2]; + #pragma unroll + for (int c_3 = 0; c_3 < 2; c_3++) { + old_max_pair[c_3] = row_max_pair[c_3]; + float _max_13 = max_noftz(row_max_pair[c_3], pair_max[c_3]); + new_max_pair[c_3] = _max_13; + } + if (lane < 8) { + uint32_t _amf_u_1 = __float_as_uint(new_max_pair[0]); + uint32_t _amf_mask_1 = -int32_t(_amf_u_1 >> 31) | 0x80000000u; + unsigned int _amf_enc_1 = _amf_u_1 ^ _amf_mask_1; + uint32_t _amf_u_2 = __float_as_uint(new_max_pair[1]); + uint32_t _amf_mask_2 = -int32_t(_amf_u_2 >> 31) | 0x80000000u; + unsigned int _amf_enc_2 = _amf_u_2 ^ _amf_mask_2; + atomicMax(&smem_softmax_u32[col_pair_base], _amf_enc_1); + atomicMax(&smem_softmax_u32[col_pair_base + 1], _amf_enc_2); + } + asm volatile("barrier.sync 1, 128;" ::: "memory"); + uint32_t _amf_u_3 = smem_softmax_u32[col_pair_base]; + uint32_t _amf_mask_3 = ((_amf_u_3 >> 31) - 1u) | 0x80000000u; + float _amf_dec_0 = __uint_as_float(_amf_u_3 ^ _amf_mask_3); + new_max_pair[0] = _amf_dec_0; + uint32_t _amf_u_4 = smem_softmax_u32[col_pair_base + 1]; + uint32_t _amf_mask_4 = ((_amf_u_4 >> 31) - 1u) | 0x80000000u; + float _amf_dec_1 = __uint_as_float(_amf_u_4 ^ _amf_mask_4); + new_max_pair[1] = _amf_dec_1; + mbarrier_wait(stats_empty_addr + (stats_stage) * 8, stats_phase); + float stats_pair[4]; + stats_pair[0] = old_max_pair[0]; + stats_pair[1] = old_max_pair[1]; + stats_pair[2] = new_max_pair[0]; + stats_pair[3] = new_max_pair[1]; + int stats_col = ((stats_stage == 0) ? 16 : 48); + tmem_st_x4_f32(taddr + (unsigned int)stats_col, stats_pair); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(stats_full_addr + (stats_stage) * 8); + stats_stage += 1; + if (stats_stage == 2) { stats_stage = 0; stats_phase ^= 1; } + float acc_scale_pair[2]; + #pragma unroll + for (int c_4 = 0; c_4 < 2; c_4++) { + float delta = softmax_scale_log2 * (old_max_pair[c_4] - new_max_pair[c_4]); + float _exp2_0 = approx_exp2(delta); + acc_scale_pair[c_4] = ((old_max_pair[c_4] > -CAKE_INF) ? _exp2_0 : 1.0f); + row_max_pair[c_4] = new_max_pair[c_4]; + } + asm volatile("" ::: "memory"); + float exp_values[8]; + #pragma unroll + for (int c_5 = 0; c_5 < 8; c_5++) { + float safe_max = ((new_max_pair[c_5 % 2] == -CAKE_INF) ? 0.0f : new_max_pair[c_5 % 2]); + float _exp2_1 = approx_exp2(scores[c_5] * softmax_scale_log2 - safe_max * softmax_scale_log2); + exp_values[c_5] = _exp2_1; + } + mbarrier_wait(p_empty_addr + (p_stage) * 8, p_phase); + float p_values[8]; + #pragma unroll + for (int c_6 = 0; c_6 < 8; c_6++) { + p_values[c_6] = exp_values[c_6] * 448.0f; + } + unsigned int regs_p[2]; + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(p_values[0]), "f"(p_values[1]), + "f"(p_values[2]), "f"(p_values[3])); + regs_p[0] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(p_values[4]), "f"(p_values[5]), + "f"(p_values[6]), "f"(p_values[7])); + regs_p[1] = _packed; + } + int mtx_idx = lane / 8; + int thr_row_idx = lane % 8; + int seg_col_idx = warp_in_wg * 2 + mtx_idx ^ thr_row_idx; + int stsm_offset = thr_row_idx * 128 + seg_col_idx * 16; + const void* _stmatrix_b8_ptr_5 = reinterpret_cast(reinterpret_cast(smem_p) + (p_stage * 1024 + (unsigned int)stsm_offset)); + uint64_t _stmatrix_b8_addr64_5; + asm volatile("cvta.to.shared.u64 %0, %1;" : "=l"(_stmatrix_b8_addr64_5) : "l"(_stmatrix_b8_ptr_5)); + uint32_t _stmatrix_b8_addr_5; + asm volatile("cvt.u32.u64 %0, %1;" : "=r"(_stmatrix_b8_addr_5) : "l"(_stmatrix_b8_addr64_5)); + asm volatile("stmatrix.sync.aligned.m16n8.x2.trans.shared.b8 [%0], {%1, %2};\n" + :: "r"(_stmatrix_b8_addr_5), "r"(regs_p[0]), "r"(regs_p[1]) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(p_full_addr + (p_stage) * 8); + p_stage += 1; + if (p_stage == 2) { p_stage = 0; p_phase ^= 1; } + float2 _f2_0 = make_float2(exp_values[0], exp_values[1]); + float2 _f2_1 = make_float2(exp_values[2], exp_values[3]); + float2 _f2_2 = make_float2(exp_values[4], exp_values[5]); + float2 _f2_3 = make_float2(exp_values[6], exp_values[7]); + float2 _f2_4 = make_float2(row_sum_pair[0], row_sum_pair[1]); + float2 _f2_5 = make_float2(acc_scale_pair[0], acc_scale_pair[1]); + float2 sum_p0_pair = fma_f32x2_rn_ftz(_f2_4, _f2_5, _f2_0); + float2 sum_p1_pair = add_f32x2(sum_p0_pair, _f2_1); + float2 sum_p2_pair = add_f32x2(sum_p1_pair, _f2_2); + float2 next_sum_pair = add_f32x2(sum_p2_pair, _f2_3); + row_sum_pair[0] = next_sum_pair.x; + row_sum_pair[1] = next_sum_pair.y; + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + mbarrier_arrive(score_empty_addr + (score_stage) * 8); + score_stage += 1; + if (score_stage == 2) { score_stage = 0; score_phase ^= 1; } + } + if (num_steps > 0) { + mbarrier_wait(stats_empty_addr + (stats_stage) * 8, stats_phase); + float final_stats_pair[4]; + final_stats_pair[0] = row_sum_pair[0]; + final_stats_pair[1] = row_sum_pair[1]; + final_stats_pair[2] = row_max_pair[0]; + final_stats_pair[3] = row_max_pair[1]; + int final_stats_col = ((stats_stage == 0) ? 16 : 48); + tmem_st_x4_f32(taddr + (unsigned int)final_stats_col, final_stats_pair); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(stats_full_addr + (stats_stage) * 8); + stats_stage += 1; + if (stats_stage == 2) { stats_stage = 0; stats_phase ^= 1; } + } + mbarrier_wait(work_full_addr + (work_stage) * 8, _phase_work_full); + uint32_t _clc_valid_3 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_3) + : "r"(work_response_addr + work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_9 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_9) + : "r"(work_response_addr + work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_10 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_10) + : "r"(work_response_addr + work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_11 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_11) + : "r"(work_response_addr + work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + work_stage * 8), "r"(0) : "memory"); + work_stage += 1; + if (work_stage == 2) { work_stage = 0; _phase_work_full ^= 1; } + if (_clc_valid_3 == 0) { + break; + } + work_x = _clc_ctaid_9; + work_y = _clc_ctaid_10; + work_z = _clc_ctaid_11; + } + } + } + // ---- Role: correction ---- + if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 136;"); + { // correction_main + unsigned int work_stage_1 = 0; + unsigned int stats_stage_1 = 0; + unsigned int stats_phase_1 = 0; + unsigned int work_x_1 = blockIdx.x; + unsigned int work_y_1 = blockIdx.y; + unsigned int work_z_1 = blockIdx.z; + const int corr_rank = warp - 4; + const int corr_tid = corr_rank * 32 + lane; + const int col_pair_base_1 = corr_tid % 4 * 2; + float softmax_scale_log2_1 = bmm1_scale[0] * 1.4426950408889634f; + float output_scale = bmm2_scale[0]; + unsigned int _phase_o_full_0 = 0; + unsigned int _phase_work_full_1 = 0; + #pragma unroll 1 + for (unsigned int _work_1 = 0; _work_1 < 120; _work_1++) { + int q_begin_1 = cum_seq_lens_q[work_z_1]; + int q_end_1 = cum_seq_lens_q[work_z_1 + 1]; + int q_len_1 = q_end_1 - q_begin_1; + int query_idx_1 = (unsigned int)q_begin_1 + work_x_1; + int active_topk_1 = 0; + if (work_x_1 < (unsigned int)q_len_1) { + int _max_14 = ((sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) : (0)); + int _min_8 = ((_max_14) < (sparse_topk) ? (_max_14) : (sparse_topk)); + active_topk_1 = _min_8; + } + int num_steps_1 = (active_topk_1 + 128 - 1) / 128; + int head_group = work_y_1 / 4; + int value_quarter = work_y_1 % 4; + if (num_steps_1 > 0) { + mbarrier_wait(stats_full_addr + (stats_stage_1) * 8, stats_phase_1); + mbarrier_arrive(stats_empty_addr + (stats_stage_1) * 8); + stats_stage_1 += 1; + if (stats_stage_1 == 2) { stats_stage_1 = 0; stats_phase_1 ^= 1; } + } + #pragma unroll 1 + for (int tile_1 = 1; tile_1 < num_steps_1; tile_1++) { + mbarrier_wait(stats_full_addr + (stats_stage_1) * 8, stats_phase_1); + int stats_col_1 = ((stats_stage_1 == 0) ? 16 : 48); + float _tmem_load_0[4]; + tmem_ld_x4(&_tmem_load_0[0], taddr + (unsigned int)stats_col_1); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + float scale_pair[2]; + #pragma unroll + for (int c_7 = 0; c_7 < 2; c_7++) { + float max_diff = _tmem_load_0[c_7] - _tmem_load_0[c_7 + 2]; + float _exp2_2 = approx_exp2(softmax_scale_log2_1 * max_diff); + scale_pair[c_7] = ((max_diff != 0.0f) ? _exp2_2 : 1.0f); + } + mbarrier_arrive(stats_empty_addr + (stats_stage_1) * 8); + stats_stage_1 += 1; + if (stats_stage_1 == 2) { stats_stage_1 = 0; stats_phase_1 ^= 1; } + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int needs_scale = ((scale_pair[0] != 1.0f) ? 1 : 0); + needs_scale = needs_scale | ((scale_pair[1] != 1.0f) ? 1 : 0); + int _vote_0 = __any_sync(0xFFFFFFFF, needs_scale != 0); + needs_scale = _vote_0; + if (needs_scale != 0) { + float values_lo[4]; + float values_hi[4]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x1.b32" + " {%0, %1, %2, %3}, [%4];" + : "=r"(*reinterpret_cast(&values_lo[0])), "=r"(*reinterpret_cast(&values_lo[1])), "=r"(*reinterpret_cast(&values_lo[2])), "=r"(*reinterpret_cast(&values_lo[3])) + : "r"(taddr + 80)); + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x1.b32" + " {%0, %1, %2, %3}, [%4];" + : "=r"(*reinterpret_cast(&values_hi[0])), "=r"(*reinterpret_cast(&values_hi[1])), "=r"(*reinterpret_cast(&values_hi[2])), "=r"(*reinterpret_cast(&values_hi[3])) + : "r"(taddr + 80 + 1048576)); + #pragma unroll + for (int pair = 0; pair < 4; pair++) { + asm volatile("" ::: "memory"); + values_lo[pair] = values_lo[pair] * scale_pair[pair % 2]; + values_hi[pair] = values_hi[pair] * scale_pair[pair % 2]; + } + asm volatile( + "tcgen05.st.sync.aligned.16x256b.x1.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(taddr + 80), "r"(*reinterpret_cast(&values_lo[0])), "r"(*reinterpret_cast(&values_lo[1])), "r"(*reinterpret_cast(&values_lo[2])), "r"(*reinterpret_cast(&values_lo[3]))); + asm volatile( + "tcgen05.st.sync.aligned.16x256b.x1.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(taddr + 80 + 1048576), "r"(*reinterpret_cast(&values_hi[0])), "r"(*reinterpret_cast(&values_hi[1])), "r"(*reinterpret_cast(&values_hi[2])), "r"(*reinterpret_cast(&values_hi[3]))); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + } + mbarrier_arrive(o_empty_addr); + } + if (num_steps_1 > 0) { + mbarrier_wait(stats_full_addr + (stats_stage_1) * 8, stats_phase_1); + int final_stats_col_1 = ((stats_stage_1 == 0) ? 16 : 48); + float _tmem_load_1[4]; + tmem_ld_x4(&_tmem_load_1[0], taddr + (unsigned int)final_stats_col_1); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + mbarrier_arrive(stats_empty_addr + (stats_stage_1) * 8); + stats_stage_1 += 1; + if (stats_stage_1 == 2) { stats_stage_1 = 0; stats_phase_1 ^= 1; } + float pair_sum[2]; + pair_sum[0] = _tmem_load_1[0]; + pair_sum[1] = _tmem_load_1[1]; + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, pair_sum[0], 16); + pair_sum[0] = pair_sum[0] + _shfl_xor_2; + float _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, pair_sum[1], 16); + pair_sum[1] = pair_sum[1] + _shfl_xor_3; + float _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, pair_sum[0], 8); + pair_sum[0] = pair_sum[0] + _shfl_xor_4; + float _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, pair_sum[1], 8); + pair_sum[1] = pair_sum[1] + _shfl_xor_5; + float _shfl_xor_6 = __shfl_xor_sync(0xFFFFFFFF, pair_sum[0], 4); + pair_sum[0] = pair_sum[0] + _shfl_xor_6; + float _shfl_xor_7 = __shfl_xor_sync(0xFFFFFFFF, pair_sum[1], 4); + pair_sum[1] = pair_sum[1] + _shfl_xor_7; + if (lane < 4) { + smem_corr[corr_rank * 8 + col_pair_base_1] = pair_sum[0]; + smem_corr[corr_rank * 8 + col_pair_base_1 + 1] = pair_sum[1]; + } + asm volatile("barrier.sync 2, 128;" ::: "memory"); + float sum0 = smem_corr[col_pair_base_1] + smem_corr[8 + col_pair_base_1] + smem_corr[16 + col_pair_base_1] + smem_corr[24 + col_pair_base_1]; + float sum1 = smem_corr[col_pair_base_1 + 1] + smem_corr[8 + col_pair_base_1 + 1] + smem_corr[16 + col_pair_base_1 + 1] + smem_corr[24 + col_pair_base_1 + 1]; + if (has_sinks != 0) { + float _exp2_3 = approx_exp2(sinks[head_group * 8 + col_pair_base_1] * 1.4426950408889634f - _tmem_load_1[2] * softmax_scale_log2_1); + sum0 = sum0 + _exp2_3; + float _exp2_4 = approx_exp2(sinks[head_group * 8 + col_pair_base_1 + 1] * 1.4426950408889634f - _tmem_load_1[3] * softmax_scale_log2_1); + sum1 = sum1 + _exp2_4; + } + float _rcp_0 = approx_rcp(sum0); + float scale0 = ((sum0 > 0.0f) ? output_scale * 0.002232142857142857f * _rcp_0 : 0.0f); + float _rcp_1 = approx_rcp(sum1); + float scale1 = ((sum1 > 0.0f) ? output_scale * 0.002232142857142857f * _rcp_1 : 0.0f); + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + float output_lo[4]; + float output_hi[4]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x1.b32" + " {%0, %1, %2, %3}, [%4];" + : "=r"(*reinterpret_cast(&output_lo[0])), "=r"(*reinterpret_cast(&output_lo[1])), "=r"(*reinterpret_cast(&output_lo[2])), "=r"(*reinterpret_cast(&output_lo[3])) + : "r"(taddr + 80)); + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x1.b32" + " {%0, %1, %2, %3}, [%4];" + : "=r"(*reinterpret_cast(&output_hi[0])), "=r"(*reinterpret_cast(&output_hi[1])), "=r"(*reinterpret_cast(&output_hi[2])), "=r"(*reinterpret_cast(&output_hi[3])) + : "r"(taddr + 80 + 1048576)); + unsigned int packed_o[4]; + #pragma unroll + for (int pair_1 = 0; pair_1 < 4; pair_1++) { + asm volatile("" ::: "memory"); + output_lo[pair_1] = output_lo[pair_1] * ((pair_1 % 2 == 0) ? scale0 : scale1); + output_hi[pair_1] = output_hi[pair_1] * ((pair_1 % 2 == 0) ? scale0 : scale1); + } + float output_values[8]; + #pragma unroll + for (int c_8 = 0; c_8 < 4; c_8++) { + output_values[c_8] = output_lo[c_8]; + output_values[c_8 + 4] = output_hi[c_8]; + } + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(output_values[_lp*2 + 0], output_values[_lp*2+1 + 0])); + packed_o[_lp] = *(uint32_t*)&_bf2; + } + int slice_idx = corr_rank / 2; + int warp_idx_in_slice = corr_rank % 2; + int mtx_idx_1 = lane / 8; + int thr_row_idx_1 = lane % 8; + int seg_col_idx_1 = warp_idx_in_slice * 4 + mtx_idx_1 ^ thr_row_idx_1; + int stsm_offset_1 = slice_idx * 8 * 128 + thr_row_idx_1 * 128 + seg_col_idx_1 * 16; + asm volatile("stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 [%0], {%1, %2, %3, %4};\n" + :: "l"(__cvta_generic_to_shared(reinterpret_cast(smem_o) + stsm_offset_1)), "r"(*reinterpret_cast(&packed_o[0])), "r"(*reinterpret_cast(&packed_o[1])), "r"(*reinterpret_cast(&packed_o[2])), "r"(*reinterpret_cast(&packed_o[3])) + : "memory"); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + int copy_base = corr_tid * 16; + int smem_row = copy_base / 128; + int copy_smem_offset = copy_base ^ smem_row % 8 * 16; + int copy_row = smem_row % 8; + int copy_col = (smem_row / 8 * 128 + copy_base % 128) / 2; + unsigned int copy_vec[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(©_vec[0])), "=r"(*reinterpret_cast(©_vec[(0) + 1])), "=r"(*reinterpret_cast(©_vec[(0) + 2])), "=r"(*reinterpret_cast(©_vec[(0) + 3])) + : "r"(smem_o_addr + (unsigned int)copy_smem_offset)); + int head_groups = gridDim.y / 4; + int output_base = query_idx_1 * head_groups * 8 * 512 + head_group * 8 * 512 + value_quarter * 128; + reinterpret_cast(O + (output_base + copy_row * 512 + copy_col))[0] = reinterpret_cast(copy_vec)[0]; + mbarrier_arrive(o_empty_addr); + } + mbarrier_wait(work_full_addr + (work_stage_1) * 8, _phase_work_full_1); + uint32_t _clc_valid_4 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_4) + : "r"(work_response_addr + work_stage_1 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_12 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_12) + : "r"(work_response_addr + work_stage_1 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_13 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_13) + : "r"(work_response_addr + work_stage_1 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_14 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_14) + : "r"(work_response_addr + work_stage_1 * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + work_stage_1 * 8), "r"(0) : "memory"); + work_stage_1 += 1; + if (work_stage_1 == 2) { work_stage_1 = 0; _phase_work_full_1 ^= 1; } + if (_clc_valid_4 == 0) { + break; + } + work_x_1 = _clc_ctaid_12; + work_y_1 = _clc_ctaid_13; + work_z_1 = _clc_ctaid_14; + } + asm volatile("barrier.sync 11, 128;" ::: "memory"); + if (corr_rank == 0) { + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + } + } + // ---- Role: mma ---- + if (warp == 8) { + { // mma_main + unsigned int work_stage_2 = 0; + unsigned int q_stage = 0; + unsigned int q_phase = 0; + unsigned int kv_stage = 0; + unsigned int kv_phase = 0; + unsigned int score_stage_1 = 0; + unsigned int score_phase_1 = 1; + unsigned int p_stage_1 = 0; + unsigned int p_phase_1 = 0; + unsigned int work_x_2 = blockIdx.x; + unsigned int work_y_2 = blockIdx.y; + unsigned int work_z_2 = blockIdx.z; + unsigned int _phase_o_empty_0 = 1; + unsigned int _phase_work_full_2 = 0; + #pragma unroll 1 + for (unsigned int _work_2 = 0; _work_2 < 120; _work_2++) { + int q_begin_2 = cum_seq_lens_q[work_z_2]; + int q_end_2 = cum_seq_lens_q[work_z_2 + 1]; + int q_len_2 = q_end_2 - q_begin_2; + int query_idx_2 = (unsigned int)q_begin_2 + work_x_2; + int active_topk_2 = 0; + if (work_x_2 < (unsigned int)q_len_2) { + int _max_3 = ((sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) : (0)); + int _min_6 = ((_max_3) < (sparse_topk) ? (_max_3) : (sparse_topk)); + active_topk_2 = _min_6; + } + int num_steps_2 = (active_topk_2 + 128 - 1) / 128; + int value_quarter_1 = work_y_2 % 4; + if (num_steps_2 > 0) { + mbarrier_wait(q_full_addr + (q_stage) * 8, q_phase); + int first_output = 1; + mbarrier_wait(score_empty_addr + (score_stage_1) * 8, score_phase_1); + int score_col_1 = ((score_stage_1 == 0) ? 0 : 8); + #pragma unroll + for (int hs = 0; hs < 3; hs++) { + int hs_sel = ((value_quarter_1 > hs) ? hs : hs + 1); + mbarrier_wait(kv_full_addr + (kv_stage) * 8, kv_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_0 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage) * 1024); + int _mma_b_lo_0 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (hs_sel) * 64); + { + uint64_t _mma_ss_a_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_0); + uint64_t _mma_ss_b_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_0); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134348816, ((hs == 0) ? 0 : 1)); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134348816, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134348816, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134348816, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage) * 8); + kv_stage += 1; + if (kv_stage == 8) { kv_stage = 0; kv_phase ^= 1; } + } + mbarrier_wait(kv_full_addr + (kv_stage) * 8, kv_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_1 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage) * 1024); + int _mma_b_lo_1 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (value_quarter_1) * 64); + { + uint64_t _mma_ss_a_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_1); + uint64_t _mma_ss_b_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_1); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134348816, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134348816, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134348816, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134348816, 1); + } + } + kv_stage += 1; + if (kv_stage == 8) { kv_stage = 0; kv_phase ^= 1; } + elect_commit(score_full_addr + (score_stage_1) * 8); + score_stage_1 += 1; + if (score_stage_1 == 2) { score_stage_1 = 0; score_phase_1 ^= 1; } + #pragma unroll 1 + for (int _tile = 0; _tile < num_steps_2 - 1; _tile++) { + mbarrier_wait(score_empty_addr + (score_stage_1) * 8, score_phase_1); + int score_col_0 = ((score_stage_1 == 0) ? 0 : 8); + #pragma unroll + for (int hs_1 = 0; hs_1 < 3; hs_1++) { + int hs_sel_1 = ((value_quarter_1 > hs_1) ? hs_1 : hs_1 + 1); + mbarrier_wait(kv_full_addr + (kv_stage) * 8, kv_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_2 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage) * 1024); + int _mma_b_lo_2 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (hs_sel_1) * 64); + { + uint64_t _mma_ss_a_desc_2 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_2); + uint64_t _mma_ss_b_desc_2 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_2); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134348816, ((hs_1 == 0) ? 0 : 1)); + } + incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134348816, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134348816, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134348816, 1); + } + } + elect_commit(kv_empty_addr + (kv_stage) * 8); + kv_stage += 1; + if (kv_stage == 8) { kv_stage = 0; kv_phase ^= 1; } + } + mbarrier_wait(kv_full_addr + (kv_stage) * 8, kv_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_3 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage) * 1024); + int _mma_b_lo_3 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (value_quarter_1) * 64); + { + uint64_t _mma_ss_a_desc_3 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_3); + uint64_t _mma_ss_b_desc_3 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_3); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134348816, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134348816, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134348816, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134348816, 1); + } + } + kv_stage += 1; + if (kv_stage == 8) { kv_stage = 0; kv_phase ^= 1; } + elect_commit(score_full_addr + (score_stage_1) * 8); + score_stage_1 += 1; + if (score_stage_1 == 2) { score_stage_1 = 0; score_phase_1 ^= 1; } + unsigned int v_cur = (kv_stage + 3) % 8; + mbarrier_wait(p_full_addr + (p_stage_1) * 8, p_phase_1); + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_4 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (v_cur) * 1024); + int _mma_b_lo_4 = make_warp_uniform(((((smem_p_addr) >> 4) & 0x3FFF) | 0x400000) + (p_stage_1) * 64); + { + uint64_t _mma_ss_a_desc_4 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_4); + uint64_t _mma_ss_b_desc_4 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_4); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134381584, ((first_output) ? 0 : 1)); + } + incr_smem_desc_lo(_mma_ss_a_desc_4, 256U); + incr_smem_desc_lo(_mma_ss_b_desc_4, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134381584, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_4, 256U); + incr_smem_desc_lo(_mma_ss_b_desc_4, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134381584, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_4, 256U); + incr_smem_desc_lo(_mma_ss_b_desc_4, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134381584, 1); + } + } + elect_commit(o_full_addr); + elect_commit(kv_empty_addr + (v_cur) * 8); + elect_commit(p_empty_addr + (p_stage_1) * 8); + p_stage_1 += 1; + if (p_stage_1 == 2) { p_stage_1 = 0; p_phase_1 ^= 1; } + first_output = 0; + } + unsigned int v_tail = (kv_stage + 7) % 8; + mbarrier_wait(p_full_addr + (p_stage_1) * 8, p_phase_1); + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_5 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (v_tail) * 1024); + int _mma_b_lo_5 = make_warp_uniform(((((smem_p_addr) >> 4) & 0x3FFF) | 0x400000) + (p_stage_1) * 64); + { + uint64_t _mma_ss_a_desc_5 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_5); + uint64_t _mma_ss_b_desc_5 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_5); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134381584, ((first_output) ? 0 : 1)); + } + incr_smem_desc_lo(_mma_ss_a_desc_5, 256U); + incr_smem_desc_lo(_mma_ss_b_desc_5, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134381584, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_5, 256U); + incr_smem_desc_lo(_mma_ss_b_desc_5, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134381584, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_5, 256U); + incr_smem_desc_lo(_mma_ss_b_desc_5, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134381584, 1); + } + } + elect_commit(o_full_addr); + elect_commit(kv_empty_addr + (v_tail) * 8); + elect_commit(p_empty_addr + (p_stage_1) * 8); + p_stage_1 += 1; + if (p_stage_1 == 2) { p_stage_1 = 0; p_phase_1 ^= 1; } + elect_commit(q_empty_addr + (q_stage) * 8); + q_phase ^= 1; + } + mbarrier_wait(work_full_addr + (work_stage_2) * 8, _phase_work_full_2); + uint32_t _clc_valid_2 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_2) + : "r"(work_response_addr + work_stage_2 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_6 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_6) + : "r"(work_response_addr + work_stage_2 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_7 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_7) + : "r"(work_response_addr + work_stage_2 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_8 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_8) + : "r"(work_response_addr + work_stage_2 * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + work_stage_2 * 8), "r"(0) : "memory"); + work_stage_2 += 1; + if (work_stage_2 == 2) { work_stage_2 = 0; _phase_work_full_2 ^= 1; } + if (_clc_valid_2 == 0) { + break; + } + work_x_2 = _clc_ctaid_6; + work_y_2 = _clc_ctaid_7; + work_z_2 = _clc_ctaid_8; + } + #pragma unroll + for (int tail = 0; tail < 2; tail++) { + elect_commit(score_full_addr + (tail) * 8); + } + } + } + // ---- Role: page_loader ---- + if (warp == 9) { + { // page_loader_main + unsigned int work_stage_3 = 0; + unsigned int page_stage = 0; + unsigned int page_phase = 1; + unsigned int work_x_3 = blockIdx.x; + unsigned int work_y_3 = blockIdx.y; + unsigned int work_z_3 = blockIdx.z; + unsigned int _phase_work_full_3 = 0; + #pragma unroll 1 + for (unsigned int _work_3 = 0; _work_3 < 120; _work_3++) { + int q_begin_3 = cum_seq_lens_q[work_z_3]; + int q_end_3 = cum_seq_lens_q[work_z_3 + 1]; + int q_len_3 = q_end_3 - q_begin_3; + int query_idx_3 = (unsigned int)q_begin_3 + work_x_3; + int active_topk_3 = 0; + if (work_x_3 < (unsigned int)q_len_3) { + int _max_0 = ((sparse_topk_lens[query_idx_3] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_3] + sparse_topk_lens_offset) : (0)); + int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); + active_topk_3 = _min_0; + } + int num_steps_3 = (active_topk_3 + 128 - 1) / 128; + int num_pairs = (num_steps_3 + 1) / 2; + int _max_1 = ((active_topk_3 - 1 & -4) > (0) ? (active_topk_3 - 1 & -4) : (0)); + int last_aligned = _max_1; + int* swa_row = swa_indices + (query_idx_3 * swa_index_stride); + int* compressed_row = compressed_indices + (query_idx_3 * compressed_index_stride); + #pragma unroll 1 + for (int pair_2 = 0; pair_2 < num_pairs; pair_2++) { + int lane_base = lane * 4; + #pragma unroll + for (int duplicate = 0; duplicate < 2; duplicate++) { + mbarrier_wait(page_empty_addr + (page_stage) * 8, page_phase); + int page0 = pair_2 * 2 * 128 + lane_base; + int page1 = page0 + 128; + int safe0 = ((page0 <= last_aligned) ? page0 : last_aligned); + int safe1 = ((page1 <= last_aligned) ? page1 : last_aligned); + int* src0 = ((safe0 < 128) ? swa_row + safe0 : compressed_row + (safe0 - 128)); + int* src1 = ((safe1 < 128) ? swa_row + safe1 : compressed_row + (safe1 - 128)); + asm volatile("cp.async.ca.shared::cta.global [%0], [%1], 16;" + :: "r"(smem_page_offsets_addr + page_stage * 1024 + (unsigned int)(lane_base * 4)), "l"(src0)); + asm volatile("cp.async.ca.shared::cta.global [%0], [%1], 16;" + :: "r"(smem_page_offsets_addr + page_stage * 1024 + (unsigned int)((128 + lane_base) * 4)), "l"(src1)); + asm volatile( + "{\n\t" + "cp.async.mbarrier.arrive.noinc.shared::cta.b64 [%0];\n\t" + "}" + :: "r"(page_full_addr + (page_stage) * 8) : "memory"); + page_stage += 1; + if (page_stage == 6) { page_stage = 0; page_phase ^= 1; } + } + } + mbarrier_wait(work_full_addr + (work_stage_3) * 8, _phase_work_full_3); + uint32_t _clc_valid_0 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_0) + : "r"(work_response_addr + work_stage_3 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_0 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_0) + : "r"(work_response_addr + work_stage_3 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_1 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_1) + : "r"(work_response_addr + work_stage_3 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_2 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_2) + : "r"(work_response_addr + work_stage_3 * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + work_stage_3 * 8), "r"(0) : "memory"); + work_stage_3 += 1; + if (work_stage_3 == 2) { work_stage_3 = 0; _phase_work_full_3 ^= 1; } + if (_clc_valid_0 == 0) { + break; + } + work_x_3 = _clc_ctaid_0; + work_y_3 = _clc_ctaid_1; + work_z_3 = _clc_ctaid_2; + } + } + } + // ---- Role: scheduler ---- + if (warp == 10) { + { // scheduler_main + asm volatile("griddepcontrol.wait;" ::: "memory"); + unsigned int work_stage_4 = 0; + unsigned int throttle_stage = 0; + unsigned int throttle_phase = 0; + unsigned int _phase_work_empty = 1; + unsigned int _phase_work_full_4 = 0; + #pragma unroll 1 + for (unsigned int _work_4 = 0; _work_4 < 120; _work_4++) { + mbarrier_wait(throttle_full_addr + (throttle_stage) * 8, throttle_phase); + mbarrier_arrive(throttle_empty_addr + (throttle_stage) * 8); + throttle_stage += 1; + if (throttle_stage == 2) { throttle_stage = 0; throttle_phase ^= 1; } + mbarrier_wait(work_empty_addr + (work_stage_4) * 8, _phase_work_empty); + if (elect_sync()) { + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [remAddr32], %2;\n\t" + "}" + :: "r"(work_full_addr + work_stage_4 * 8), "r"(0), "r"((uint32_t)(16)) : "memory"); + asm volatile( + "fence.proxy.async.shared::cta;\n\t" + "clusterlaunchcontrol.try_cancel.async.shared::cta" + ".mbarrier::complete_tx::bytes.b128" + " [%0], [%1];" + :: "r"(work_response_addr + work_stage_4 * 16 + 0 * 16), "r"(work_full_addr + work_stage_4 * 8) + : "memory"); + } + mbarrier_wait(work_full_addr + (work_stage_4) * 8, _phase_work_full_4); + uint32_t _clc_valid_5 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_5) + : "r"(work_response_addr + work_stage_4 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_15 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_15) + : "r"(work_response_addr + work_stage_4 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_16 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_16) + : "r"(work_response_addr + work_stage_4 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_17 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_17) + : "r"(work_response_addr + work_stage_4 * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + work_stage_4 * 8), "r"(0) : "memory"); + work_stage_4 += 1; + if (work_stage_4 == 2) { work_stage_4 = 0; _phase_work_empty ^= 1; _phase_work_full_4 ^= 1; } + if (_clc_valid_5 == 0) { + break; + } + } + } + } + // ---- Role: padding ---- + if (warp == 11) { + { // padding_main + unsigned int work_stage_5 = 0; + unsigned int work_x_4 = blockIdx.x; + unsigned int work_y_4 = blockIdx.y; + unsigned int work_z_4 = blockIdx.z; + unsigned int _phase_work_full_5 = 0; + #pragma unroll 1 + for (unsigned int _work_5 = 0; _work_5 < 120; _work_5++) { + mbarrier_wait(work_full_addr + (work_stage_5) * 8, _phase_work_full_5); + uint32_t _clc_valid_6 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_6) + : "r"(work_response_addr + work_stage_5 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_18 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_18) + : "r"(work_response_addr + work_stage_5 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_19 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_19) + : "r"(work_response_addr + work_stage_5 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_20 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_20) + : "r"(work_response_addr + work_stage_5 * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + work_stage_5 * 8), "r"(0) : "memory"); + work_stage_5 += 1; + if (work_stage_5 == 2) { work_stage_5 = 0; _phase_work_full_5 ^= 1; } + if (_clc_valid_6 == 0) { + break; + } + work_x_4 = _clc_ctaid_18; + work_y_4 = _clc_ctaid_19; + work_z_4 = _clc_ctaid_20; + } + } + } + // ---- Role: loader ---- + if (warp >= 12 && warp <= 19) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 72;"); + { // loader_main + asm volatile("griddepcontrol.wait;" ::: "memory"); + const int load_rank = warp - 12; + unsigned int work_stage_6 = 0; + unsigned int throttle_stage_1 = 0; + unsigned int throttle_phase_1 = 1; + unsigned int q_stage_1 = 0; + unsigned int q_phase_1 = 1; + unsigned int kv_stage_1 = 0; + unsigned int kv_phase_1 = 1; + unsigned int page_k_stage = 0; + unsigned int page_k_phase = 0; + unsigned int page_v_stage = 1; + unsigned int page_v_phase = 0; + unsigned int work_x_5 = blockIdx.x; + unsigned int work_y_5 = blockIdx.y; + unsigned int work_z_5 = blockIdx.z; + unsigned int _phase_work_full_6 = 0; + #pragma unroll 1 + for (unsigned int _work_6 = 0; _work_6 < 120; _work_6++) { + int q_begin_4 = cum_seq_lens_q[work_z_5]; + int q_end_4 = cum_seq_lens_q[work_z_5 + 1]; + int q_len_4 = q_end_4 - q_begin_4; + int query_idx_4 = (unsigned int)q_begin_4 + work_x_5; + int active_topk_4 = 0; + if (work_x_5 < (unsigned int)q_len_4) { + int _max_2 = ((sparse_topk_lens[query_idx_4] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_4] + sparse_topk_lens_offset) : (0)); + int _min_1 = ((_max_2) < (sparse_topk) ? (_max_2) : (sparse_topk)); + active_topk_4 = _min_1; + } + int num_steps_4 = (active_topk_4 + 128 - 1) / 128; + int head_group_1 = work_y_5 / 4; + int value_quarter_2 = work_y_5 % 4; + mbarrier_wait(throttle_empty_addr + (throttle_stage_1) * 8, throttle_phase_1); + mbarrier_arrive(throttle_full_addr + (throttle_stage_1) * 8); + throttle_stage_1 += 1; + if (throttle_stage_1 == 2) { throttle_stage_1 = 0; throttle_phase_1 ^= 1; } + if (num_steps_4 > 0) { + mbarrier_wait(q_empty_addr + (q_stage_1) * 8, q_phase_1); + if (load_rank == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(q_full_addr + (q_stage_1) * 8, 4096); + #pragma unroll + for (int hs_2 = 0; hs_2 < 4; hs_2++) { + tma_4d_gmem2smem(smem_q_addr + (unsigned int)(hs_2 * 1024), (&tmap_q), 0, head_group_1 * 8, hs_2, query_idx_4, q_full_addr + (q_stage_1) * 8); + } + } + } + mbarrier_wait(page_full_addr + (page_k_stage) * 8, page_k_phase); + int page_base = 0; + int warps_tile_idx = 0; + int _min_2 = ((active_topk_4) < (128) ? (active_topk_4) : (128)); + int warps_valid_rows = _min_2; + int _min_3 = (((warps_valid_rows + 3) / 4) < (8) ? ((warps_valid_rows + 3) / 4) : (8)); + int warps_tile = _min_3; + int warps_ge1 = 8; + int warps_ge2 = 8; + if (warps_tile_idx >= 1) { + warps_ge1 = warps_tile; + } + if (warps_tile_idx >= 2) { + warps_ge2 = warps_tile; + } + #pragma unroll + for (int hs_3 = 0; hs_3 < 4; hs_3++) { + int hs_sel_2 = ((hs_3 == 3) ? value_quarter_2 : ((value_quarter_2 > hs_3) ? hs_3 : hs_3 + 1)); + mbarrier_wait(kv_empty_addr + (kv_stage_1) * 8, kv_phase_1); + if (load_rank == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_1) * 8, ((hs_3 >= 4) ? warps_ge1 : warps_ge2) * 2048); + } + } + asm volatile("barrier.sync 9, 256;" ::: "memory"); + if (load_rank < ((hs_3 >= 4) ? warps_ge1 : warps_ge2)) { + int rows_all[16]; + #pragma unroll + for (int group = 0; group < 4; group++) { + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&rows_all[group * 4])), "=r"(*reinterpret_cast(&rows_all[(group * 4) + 1])), "=r"(*reinterpret_cast(&rows_all[(group * 4) + 2])), "=r"(*reinterpret_cast(&rows_all[(group * 4) + 3])) + : "r"(smem_page_offsets_addr + page_k_stage * 1024 + (unsigned int)((page_base + load_rank * 4 + group * 32) * 4))); + } + if (elect_sync()) { + { + #pragma unroll + for (int group_1 = 0; group_1 < 4; group_1++) { + const int row_base = load_rank * 4 + group_1 * 32; + int dst_k = smem_kv_addr + kv_stage_1 * 16384 + (unsigned int)(row_base * 128); + tma_gather4_gmem2smem(dst_k, (&tmap_swa_kv), hs_sel_2 * 128, rows_all[group_1 * 4], rows_all[group_1 * 4 + 1], rows_all[group_1 * 4 + 2], rows_all[group_1 * 4 + 3], kv_full_addr + (kv_stage_1) * 8); + } + } + } + } + kv_stage_1 += 1; + if (kv_stage_1 == 8) { kv_stage_1 = 0; kv_phase_1 ^= 1; } + } + #pragma unroll 1 + for (int tile_2 = 0; tile_2 < num_steps_4 - 1; tile_2++) { + int next_tile = tile_2 + 1; + if ((next_tile & 1) == 0) { + mbarrier_wait(page_full_addr + (page_k_stage) * 8, page_k_phase); + } + int page_base_0 = (next_tile & 1) * 128; + int warps_tile_idx_1 = next_tile; + int _min_4 = ((active_topk_4 - next_tile * 128) < (128) ? (active_topk_4 - next_tile * 128) : (128)); + int warps_valid_rows_2 = _min_4; + int _min_5 = (((warps_valid_rows_2 + 3) / 4) < (8) ? ((warps_valid_rows_2 + 3) / 4) : (8)); + int warps_tile_3 = _min_5; + int warps_ge1_4 = 8; + int warps_ge2_5 = 8; + if (warps_tile_idx_1 >= 1) { + warps_ge1_4 = warps_tile_3; + } + if (warps_tile_idx_1 >= 2) { + warps_ge2_5 = warps_tile_3; + } + #pragma unroll + for (int hs_4 = 0; hs_4 < 4; hs_4++) { + int hs_sel_3 = ((hs_4 == 3) ? value_quarter_2 : ((value_quarter_2 > hs_4) ? hs_4 : hs_4 + 1)); + mbarrier_wait(kv_empty_addr + (kv_stage_1) * 8, kv_phase_1); + if (load_rank == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_1) * 8, ((hs_4 >= 4) ? warps_ge1_4 : warps_ge2_5) * 2048); + } + } + asm volatile("barrier.sync 9, 256;" ::: "memory"); + if (load_rank < ((hs_4 >= 4) ? warps_ge1_4 : warps_ge2_5)) { + int rows_all_1[16]; + #pragma unroll + for (int group_2 = 0; group_2 < 4; group_2++) { + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&rows_all_1[group_2 * 4])), "=r"(*reinterpret_cast(&rows_all_1[(group_2 * 4) + 1])), "=r"(*reinterpret_cast(&rows_all_1[(group_2 * 4) + 2])), "=r"(*reinterpret_cast(&rows_all_1[(group_2 * 4) + 3])) + : "r"(smem_page_offsets_addr + page_k_stage * 1024 + (unsigned int)((page_base_0 + load_rank * 4 + group_2 * 32) * 4))); + } + if (elect_sync()) { + if (next_tile == 0) { + #pragma unroll + for (int group_3 = 0; group_3 < 4; group_3++) { + const int row_base_1 = load_rank * 4 + group_3 * 32; + int dst_k_1 = smem_kv_addr + kv_stage_1 * 16384 + (unsigned int)(row_base_1 * 128); + tma_gather4_gmem2smem(dst_k_1, (&tmap_swa_kv), hs_sel_3 * 128, rows_all_1[group_3 * 4], rows_all_1[group_3 * 4 + 1], rows_all_1[group_3 * 4 + 2], rows_all_1[group_3 * 4 + 3], kv_full_addr + (kv_stage_1) * 8); + } + } else { + #pragma unroll + for (int group_4 = 0; group_4 < 4; group_4++) { + const int row_base_2 = load_rank * 4 + group_4 * 32; + int dst_k_2 = smem_kv_addr + kv_stage_1 * 16384 + (unsigned int)(row_base_2 * 128); + tma_gather4_gmem2smem(dst_k_2, (&tmap_compressed_kv), hs_sel_3 * 128, rows_all_1[group_4 * 4], rows_all_1[group_4 * 4 + 1], rows_all_1[group_4 * 4 + 2], rows_all_1[group_4 * 4 + 3], kv_full_addr + (kv_stage_1) * 8); + } + } + } + } + kv_stage_1 += 1; + if (kv_stage_1 == 8) { kv_stage_1 = 0; kv_phase_1 ^= 1; } + } + if ((next_tile & 1) != 0 && next_tile < num_steps_4 - 1) { + mbarrier_arrive(page_empty_addr + (page_k_stage) * 8); + page_k_stage += 1; + if (page_k_stage == 6) { page_k_stage = 0; page_k_phase ^= 1; } + page_k_stage += 1; + if (page_k_stage == 6) { page_k_stage = 0; page_k_phase ^= 1; } + } + if ((tile_2 & 1) == 0) { + mbarrier_wait(page_full_addr + (page_v_stage) * 8, page_v_phase); + } + if ((tile_2 & 1) != 0) { + mbarrier_arrive(page_empty_addr + (page_v_stage) * 8); + page_v_stage += 1; + if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } + page_v_stage += 1; + if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } + } + } + int last_tile = num_steps_4 - 1; + mbarrier_arrive(page_empty_addr + (page_k_stage) * 8); + page_k_stage += 1; + if (page_k_stage == 6) { page_k_stage = 0; page_k_phase ^= 1; } + page_k_stage += 1; + if (page_k_stage == 6) { page_k_stage = 0; page_k_phase ^= 1; } + if ((last_tile & 1) == 0) { + mbarrier_wait(page_full_addr + (page_v_stage) * 8, page_v_phase); + } + mbarrier_arrive(page_empty_addr + (page_v_stage) * 8); + page_v_stage += 1; + if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } + page_v_stage += 1; + if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } + } + mbarrier_wait(work_full_addr + (work_stage_6) * 8, _phase_work_full_6); + uint32_t _clc_valid_1 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_1) + : "r"(work_response_addr + work_stage_6 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_3 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_3) + : "r"(work_response_addr + work_stage_6 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_4 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_4) + : "r"(work_response_addr + work_stage_6 * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_5 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_5) + : "r"(work_response_addr + work_stage_6 * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + work_stage_6 * 8), "r"(0) : "memory"); + work_stage_6 += 1; + if (work_stage_6 == 2) { work_stage_6 = 0; _phase_work_full_6 ^= 1; } + if (_clc_valid_1 == 0) { + break; + } + work_x_5 = _clc_ctaid_3; + work_y_5 = _clc_ctaid_4; + work_z_5 = _clc_ctaid_5; + } + } + } + + // Cleanup +} + +} // extern "C" diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_42fa8e98da5fa46d1a94_binding.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_42fa8e98da5fa46d1a94_binding.cu new file mode 100644 index 00000000000..214a8b9199b --- /dev/null +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_42fa8e98da5fa46d1a94_binding.cu @@ -0,0 +1,606 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +// Generated by the Cake source exporter — do not edit. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_42fa8e98da5fa46d1a94'. +#include +#include +#include +#include + +#include "tvm_ffi_utils.h" + +#include +#include +#include +#include +#include +#include +#include + +extern "C" __global__ void kernel_cake_dsv4_42fa8e98da5fa46d1a94(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, const __grid_constant__ CUtensorMap O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); + + +namespace cake_host_shim_c0e0335c8d570682 { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t CakeDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +// 4D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + int64_t d3 = t.size(t.ndim() - 3); + TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(d3)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 64u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_swa_k' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_swa_k(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_swa_k' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_swa_k' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_swa_k' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_swa_k"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_swa_k' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_swa_k' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_k'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_swa_k') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_compressed_k' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_compressed_k(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_compressed_k' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_compressed_k' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_compressed_k' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_compressed_k"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_compressed_k' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_compressed_k' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_k'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_k') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_swa_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_swa_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_swa_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_swa_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_swa_v' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_swa_v"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_swa_v' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_swa_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_v'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_swa_v') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_compressed_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_compressed_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_compressed_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_compressed_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_compressed_v' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_compressed_v"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_compressed_v' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_compressed_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_v'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_v') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'O' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_O(const TensorView& t, int64_t tp_num_heads, int64_t tp_num_query_tokens) { + TVM_FFI_CHECK(t.ndim() >= 1, ValueError) + << "TMA source 'O' must have at least 1 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'O' must have unit innermost stride, got " << t.stride(-1); + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(8), (uint64_t)(4), (uint64_t)((tp_num_query_tokens * tp_num_heads))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'O' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1] && 1u <= global_dim[2] && 32u <= global_dim[3], ValueError) + << "TMA box (64, 1, 1, 32) exceeds resolved global dims for 'O'"; + int64_t carrier_stride_0 = 64; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 512; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = 2048; + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 1u, 1u, 32u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'O') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_compressed_k, TensorView arg_tmap_swa_v, TensorView arg_tmap_compressed_v, TensorView arg_O, TensorView arg_partial_lse, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_num_query_tokens, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckContiguous(arg_tmap_q, "tmap_q"); + CheckCudaTensor(arg_tmap_swa_k, "tmap_swa_k"); + CheckDtype(arg_tmap_swa_k, "tmap_swa_k", 4, 16, 1); + CheckCudaTensor(arg_tmap_compressed_k, "tmap_compressed_k"); + CheckDtype(arg_tmap_compressed_k, "tmap_compressed_k", 4, 16, 1); + CheckCudaTensor(arg_tmap_swa_v, "tmap_swa_v"); + CheckDtype(arg_tmap_swa_v, "tmap_swa_v", 4, 16, 1); + CheckCudaTensor(arg_tmap_compressed_v, "tmap_compressed_v"); + CheckDtype(arg_tmap_compressed_v, "tmap_compressed_v", 4, 16, 1); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_partial_lse, "partial_lse"); + CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); + CheckContiguous(arg_partial_lse, "partial_lse"); + CheckCudaTensor(arg_swa_indices, "swa_indices"); + CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); + CheckContiguous(arg_swa_indices, "swa_indices"); + CheckCudaTensor(arg_compressed_indices, "compressed_indices"); + CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); + CheckContiguous(arg_compressed_indices, "compressed_indices"); + CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); + CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); + CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); + CheckContiguous(arg_bmm1_scale, "bmm1_scale"); + CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); + CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); + CheckContiguous(arg_bmm2_scale, "bmm2_scale"); + TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) + << "scalar 'num_heads' value " << arg_num_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_query_tokens >= -2147483648LL && arg_num_query_tokens <= 2147483647LL, ValueError) + << "scalar 'num_query_tokens' value " << arg_num_query_tokens + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) + << "scalar 'swa_index_stride' value " << arg_swa_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) + << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) + << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) + << "scalar 'sparse_topk' value " << arg_sparse_topk + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) + << "scalar 'has_sinks' value " << arg_has_sinks + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_swa_k, arg_tmap_q, "tmap_swa_k", "tmap_q"); + CheckSameCudaDevice(arg_tmap_compressed_k, arg_tmap_q, "tmap_compressed_k", "tmap_q"); + CheckSameCudaDevice(arg_tmap_swa_v, arg_tmap_q, "tmap_swa_v", "tmap_q"); + CheckSameCudaDevice(arg_tmap_compressed_v, arg_tmap_q, "tmap_compressed_v", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_partial_lse, arg_tmap_q, "partial_lse", "tmap_q"); + CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); + CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); + CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); + CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") must be divisible by cluster dims (2, 1, 1)"; + + cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_swa_k = EncodeTma_tmap_swa_k(arg_tmap_swa_k); + CUtensorMap p_tmap_compressed_k = EncodeTma_tmap_compressed_k(arg_tmap_compressed_k); + CUtensorMap p_tmap_swa_v = EncodeTma_tmap_swa_v(arg_tmap_swa_v); + CUtensorMap p_tmap_compressed_v = EncodeTma_tmap_compressed_v(arg_tmap_compressed_v); + CUtensorMap p_O = EncodeTma_O(arg_O, (int64_t)arg_num_heads, (int64_t)arg_num_query_tokens); + float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); + int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); + int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); + int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); + float* p_sinks = static_cast(arg_sinks.data_ptr()); + float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); + float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); + int32_t v_num_heads = (int32_t)arg_num_heads; + int32_t v_num_query_tokens = (int32_t)arg_num_query_tokens; + int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; + int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; + int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; + int32_t v_sparse_topk = (int32_t)arg_sparse_topk; + int32_t v_has_sinks = (int32_t)arg_has_sinks; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + void* kargs[] = {&p_tmap_q, &p_tmap_swa_k, &p_tmap_compressed_k, &p_tmap_swa_v, &p_tmap_compressed_v, &p_O, &p_partial_lse, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_num_query_tokens, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_has_sinks, &v_total_work_items}; + + static std::mutex smem_mu; + static auto* smem_status_by_device = + new std::unordered_map(); + cudaError_t smem_status = cudaSuccess; + { + std::lock_guard lock(smem_mu); + auto it = smem_status_by_device->find(dev.device_id); + if (it == smem_status_by_device->end()) { + int current_device = -1; + smem_status = cudaGetDevice(¤t_device); + if (smem_status == cudaSuccess && current_device != dev.device_id) { + smem_status = cudaErrorInvalidDevice; + } + if (smem_status == cudaSuccess) { + smem_status = cudaFuncSetAttribute( + reinterpret_cast(kernel_cake_dsv4_42fa8e98da5fa46d1a94), + cudaFuncAttributeMaxDynamicSharedMemorySize, + 220032); + } + smem_status_by_device->emplace(dev.device_id, smem_status); + } else { + smem_status = it->second; + } + } + TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) + << "cudaFuncSetAttribute for kernel_cake_dsv4_42fa8e98da5fa46d1a94 failed: " + << cudaGetErrorString(smem_status); + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(512u, 1u, 1u); + + // Extended direct-source launch for a linked CUDA kernel symbol. + cudaLaunchAttribute attrs[2]{}; + int n = 0; + attrs[n].id = cudaLaunchAttributeClusterDimension; + attrs[n].val.clusterDim.x = 2u; + attrs[n].val.clusterDim.y = 1u; + attrs[n].val.clusterDim.z = 1u; + ++n; + attrs[n].id = cudaLaunchAttributeClusterSchedulingPolicyPreference; + attrs[n].val.clusterSchedulingPolicyPreference = cudaClusterSchedulingPolicySpread; + ++n; + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes = 220032u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; + cudaError_t launch_status = cudaLaunchKernelExC( + &config, reinterpret_cast(kernel_cake_dsv4_42fa8e98da5fa46d1a94), kargs); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernelExC for kernel_cake_dsv4_42fa8e98da5fa46d1a94 failed: " + << cudaGetErrorString(launch_status); +} + +} // namespace cake_host_shim_c0e0335c8d570682 + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_c0e0335c8d570682::Run); diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_42fa8e98da5fa46d1a94_kernel.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_42fa8e98da5fa46d1a94_kernel.cu new file mode 100644 index 00000000000..a1879057902 --- /dev/null +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_42fa8e98da5fa46d1a94_kernel.cu @@ -0,0 +1,2364 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define CAKE_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_SCRATCH_OFFSET 0 +#define NUM_K_PIPE_STAGES 3 +#define NUM_V_PIPE_STAGES 2 +#define SMEM_SMEM_Q_OFF 1024 +#define SMEM_SMEM_Q_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_STRIDE 8192 +#define SMEM_SMEM_K_OFF 66560 +#define SMEM_SMEM_K_STAGE_BYTES 16384 +#define SMEM_SMEM_K_STRIDE 16384 +#define SMEM_SMEM_V_LO_OFF 115712 +#define SMEM_SMEM_V_LO_STAGE_BYTES 16384 +#define SMEM_SMEM_V_LO_STRIDE 32768 +#define SMEM_SMEM_V_HI_OFF 132096 +#define SMEM_SMEM_V_HI_STAGE_BYTES 16384 +#define SMEM_SMEM_V_HI_STRIDE 32768 +#define SMEM_SMEM_P_OFF 181248 +#define SMEM_SMEM_P_STAGE_BYTES 8192 +#define SMEM_SMEM_P_STRIDE 8192 +#define SMEM_SMEM_INDICES_OFF 214016 +#define SMEM_SMEM_INDICES_STAGE_BYTES 2048 +#define SMEM_SMEM_INDICES_STRIDE 2048 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_OFF 218624 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STRIDE 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_OFF 219136 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STRIDE 512 +#define SMEM_SMEM_INDEX_VALID_OFF 219648 +#define SMEM_SMEM_INDEX_VALID_STAGE_BYTES 288 +#define SMEM_SMEM_INDEX_VALID_STRIDE 288 +#define SMEM_TOTAL 220032 +#define THREADS 512 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count) : "memory"); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16_cta2( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" + "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" + "}\n" + :: "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x4(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x4.b32" + " {%0, %1, %2, %3}, [%4];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); +} + + +__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ float2 ex2_emulation_f32x2_value(float2 value) { + const float c0 = 1.0f, c1 = 0.695146143436431884765625f; + const float c2 = 0.227564394474029541015625f, c3 = 0.077119089663028717041015625f; + const float magic = 12582912.0f; + float x0 = max_noftz(value.x, -127.0f), x1 = max_noftz(value.y, -127.0f); + float2 xc2 = make_float2(x0, x1), magic2 = make_float2(magic, magic); + float2 xr2; + asm("add.rm.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xr2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&magic2)); + float2 c3_2 = make_float2(c3, c3), c2_2 = make_float2(c2, c2); + float2 c1_2 = make_float2(c1, c1), c0_2 = make_float2(c0, c0); + float2 xrb2, xfrac2; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xrb2) + : "l"(*(unsigned long long*)&xr2), "l"(*(unsigned long long*)&magic2)); + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xfrac2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&xrb2)); + float2 poly2; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&c3_2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c2_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c1_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c0_2)); + int x0r_i, x1r_i, p0_i, p1_i; + asm("mov.b64 {%0, %1}, %2;" : "=r"(x0r_i), "=r"(x1r_i) : "l"(*(unsigned long long*)&xr2)); + asm("mov.b64 {%0, %1}, %2;" : "=r"(p0_i), "=r"(p1_i) : "l"(*(unsigned long long*)&poly2)); + float r0, r1; + asm("mov.b32 %0, %1;" : "=f"(r0) : "r"((x0r_i << 23) + p0_i)); + asm("mov.b32 %0, %1;" : "=f"(r1) : "r"((x1r_i << 23) + p1_i)); + return make_float2(r0, r1); +} + +__device__ __forceinline__ void ex2_emulation_f32x2(float* x0_ptr, float* x1_ptr) { + float2 result = ex2_emulation_f32x2_value(make_float2(*x0_ptr, *x1_ptr)); + *x0_ptr = result.x; *x1_ptr = result.y; +} + +__device__ __forceinline__ void softmax_frag_exp2_cast( + float* sv, uint32_t* pv, int use_emu) +{ + #pragma unroll + for (int j = 0; j < 16; j++) { + if (use_emu && j >= 12) + ex2_emulation_f32x2(&sv[j*2], &sv[j*2+1]); + else { + sv[j*2] = approx_exp2(sv[j*2]); + sv[j*2+1] = approx_exp2(sv[j*2+1]); + } + } + #pragma unroll + for (int j = 0; j < 16; j++) { + __nv_bfloat162 bf = __float22bfloat162_rn({sv[j*2], sv[j*2+1]}); + pv[j] = reinterpret_cast(bf); + } +} + + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_2d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3}], [%4];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr) { + // Canonical .shared::cta form for non-multicast gather4, matching + // trtllm-gen / cuda_ptx and the PTX ISA qualifier order + // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, + // .shared::cluster is reserved for the multicast variant (ctaMask). + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem_cta2( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr) { + // Canonical .shared::cta form; see tma_gather4_gmem2smem above. + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem_mc( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr, unsigned short cta_mask) { + // Multicast variant: the PTX grammar ties the .shared::cluster + // destination to .multicast::cluster + ctaMask (cf. cuda_ptx / + // SM100_TMA_LOAD_MULTICAST_2D_GATHER4 in CUTLASS). + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" + ".mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem_mc_cta2( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr, unsigned short cta_mask) { + // Multicast + cta_group::2 variant; see tma_gather4_gmem2smem_mc. + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" + ".mbarrier::complete_tx::bytes.multicast::cluster.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void tma_store_4d( + const void *tmap, int x, int y, int z, int w, unsigned smem_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.global.shared::cta.tile.bulk_group" + " [%0, {%1, %2, %3, %4}], [%5];" + :: "l"(tmap), "r"(x), "r"(y), "r"(z), "r"(w), "r"(smem_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %1;\n\t" + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" + "}\n" + :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); +} + +extern "C" { + +__global__ __launch_bounds__(512, 1) __cluster_dims__(2,1,1) void +kernel_cake_dsv4_42fa8e98da5fa46d1a94(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, const __grid_constant__ CUtensorMap O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); + + const int mbar_base = smem; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define k_full_addr (mbar_base + 16) + #define k_empty_addr (mbar_base + 40) + #define v_full_addr (mbar_base + 64) + #define v_empty_addr (mbar_base + 80) + #define s_full_addr (mbar_base + 96) + #define s_empty_addr (mbar_base + 112) + #define p_full_addr (mbar_base + 128) + #define p_empty_addr (mbar_base + 144) + #define o_empty_addr (mbar_base + 160) + #define stats_addr (mbar_base + 168) + #define stats_empty_addr (mbar_base + 184) + #define o_full_addr (mbar_base + 200) + #define epi_q_free_addr (mbar_base + 208) + #define o_first_slice_seeded_addr (mbar_base + 216) + #define tmem_dealloc_peer_addr (mbar_base + 224) + #define index_valid_full_addr (mbar_base + 232) + #define index_valid_empty_addr (mbar_base + 248) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 2; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + + // Kernel setup ops + __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_q_addr = smem + 1024; + __nv_bfloat16* smem_k = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_k_addr = smem + 66560; + __nv_bfloat16* smem_v_lo = reinterpret_cast<__nv_bfloat16*>(smem_raw + 115712); + const int smem_v_lo_addr = smem + 115712; + __nv_bfloat16* smem_v_hi = reinterpret_cast<__nv_bfloat16*>(smem_raw + 132096); + const int smem_v_hi_addr = smem + 132096; + __nv_bfloat16* smem_p = reinterpret_cast<__nv_bfloat16*>(smem_raw + 181248); + const int smem_p_addr = smem + 181248; + int* smem_indices = reinterpret_cast(smem_raw + 214016); + const int smem_indices_addr = smem + 214016; + float* smem_softmax_exchange = reinterpret_cast(smem_raw + 218624); + const int smem_softmax_exchange_addr = smem + 218624; + float* smem_epilogue_exchange = reinterpret_cast(smem_raw + 219136); + const int smem_epilogue_exchange_addr = smem + 219136; + unsigned int* smem_index_valid = reinterpret_cast(smem_raw + 219648); + const int smem_index_valid_addr = smem + 219648; + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // Mbarrier init (19 pipeline groups, 0 ordered-sequence groups, 33 barriers) + // Mbarriers at smem_raw[0..264) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // q_full: 1 barriers, init_count=1 + mbarrier_init(smem + 0, 1); + // q_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 8, 1); + // --- pipeline 'k_pipe' --- + // k_full: 3 barriers, init_count=1 + mbarrier_init(smem + 16, 1); + mbarrier_init(smem + 24, 1); + mbarrier_init(smem + 32, 1); + // k_empty: 3 barriers, init_count=1 + mbarrier_init(smem + 40, 1); + mbarrier_init(smem + 48, 1); + mbarrier_init(smem + 56, 1); + // --- pipeline 'v_pipe' --- + // v_full: 2 barriers, init_count=1 + mbarrier_init(smem + 64, 1); + mbarrier_init(smem + 72, 1); + // v_empty: 2 barriers, init_count=1 + mbarrier_init(smem + 80, 1); + mbarrier_init(smem + 88, 1); + // s_full: 2 barriers, init_count=1 + mbarrier_init(smem + 96, 1); + mbarrier_init(smem + 104, 1); + // s_empty: 2 barriers, init_count=256 + mbarrier_init(smem + 112, 256); + mbarrier_init(smem + 120, 256); + // p_full: 2 barriers, init_count=256 + mbarrier_init(smem + 128, 256); + mbarrier_init(smem + 136, 256); + // p_empty: 2 barriers, init_count=1 + mbarrier_init(smem + 144, 1); + mbarrier_init(smem + 152, 1); + // o_empty: 1 barriers, init_count=256 + mbarrier_init(smem + 160, 256); + // stats: 2 barriers, init_count=128 + mbarrier_init(smem + 168, 128); + mbarrier_init(smem + 176, 128); + // stats_empty: 2 barriers, init_count=128 + mbarrier_init(smem + 184, 128); + mbarrier_init(smem + 192, 128); + // o_full: 1 barriers, init_count=1 + mbarrier_init(smem + 200, 1); + // epi_q_free: 1 barriers, init_count=4 + mbarrier_init(smem + 208, 4); + // o_first_slice_seeded: 1 barriers, init_count=256 + mbarrier_init(smem + 216, 256); + // tmem_dealloc_peer: 1 barriers, init_count=32 + mbarrier_init(smem + 224, 32); + // index_valid_full: 2 barriers, init_count=32 + mbarrier_init(smem + 232, 32); + mbarrier_init(smem + 240, 32); + // index_valid_empty: 2 barriers, init_count=128 + mbarrier_init(smem + 248, 128); + mbarrier_init(smem + 256, 128); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 264); + if (warp == 0) { + int _tmem_hold = smem + 264; + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); + } + + __syncthreads(); + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_tmem_scratch = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 15) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 56;"); + } + + // ---- Role: softmax_wg ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); + { // softmax_wg_main + const int softmax_dummy = 0; + float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; + const int head_row_base = warp % 2 * 32; + const int score_row_base = warp % 4 * 32; + const int n_half = warp % 4 / 2; + const int my_row = head_row_base + lane; + const int exchange_idx = n_half * 64 + my_row; + int tile_cursor = 0; + int valid_iter = 0; + float seed_zero[4]; + #pragma unroll + for (int seed_i = 0; seed_i < 4; seed_i++) { + seed_zero[seed_i] = 0.0f; + } + #pragma unroll + for (int seed_stage = 0; seed_stage < 2; seed_stage++) { + #pragma unroll + for (int seed_col = 128 + seed_stage * 128; seed_col < 128 + (seed_stage + 1) * 128; seed_col += 4) { + int seed_addr = taddr + (unsigned int)seed_col + (unsigned int)(score_row_base << 16); + tmem_st_x4_f32(seed_addr, seed_zero); + } + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(o_first_slice_seeded_addr); + int peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_first_slice_seeded_addr), "r"(peer_rank) : "memory"); + #pragma unroll 1 + for (unsigned int work_idx = cluster_id; work_idx < total_work_items; work_idx += num_clusters) { + int split_idx = work_idx % (unsigned int)(((1) ? 4 : 1)); + int query_idx = work_idx / (unsigned int)(((1) ? 4 : 1)); + int _max_0 = ((sparse_topk_lens[query_idx] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx] + sparse_topk_lens_offset) : (0)); + int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); + int active_topk = _min_0; + int num_kv_tiles = ((sparse_topk + 128 - 1) / 128 + ((1) ? 4 : 1) - 1) / ((1) ? 4 : 1); + int first_kv_tile = split_idx * num_kv_tiles; + int head_idx = cta_rank * 64 + my_row; + float row_max_val = -CAKE_INF; + float row_sum_val = 0.0f; + if (has_sinks != 0 && head_idx < num_heads && split_idx == 0) { + row_max_val = sinks[head_idx] * 1.4426950408889634f / softmax_scale_log2; + if (n_half == 0) { + row_sum_val = 1.0f; + } + } + int valid_slot = valid_iter & 1; + int valid_full_phase = valid_iter >> 1 & 1; + mbarrier_wait(index_valid_full_addr + (valid_slot) * 8, valid_full_phase); + int valid_word_base = valid_slot * 36 + n_half * 2; + #pragma unroll 1 + for (int tile = 0; tile < num_kv_tiles; tile++) { + int pipeline_tile = tile_cursor + tile; + int score_phase = pipeline_tile & 1; + int s_full_phase = pipeline_tile >> 1 & 1; + mbarrier_wait(s_full_addr + (score_phase) * 8, s_full_phase); + int score_offset = ((score_phase != 0) ? 64 : 0); + int score_base = taddr + (unsigned int)score_offset + (unsigned int)(score_row_base << 16); + float _tmem_load_0[64]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]) + : "r"(score_base)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]) + : "r"(score_base + 32)); + int valid_cols = active_topk - (first_kv_tile + tile) * 128 - n_half * 64; + if (valid_cols < 0) { + valid_cols = 0; + } + if (valid_cols > 64) { + valid_cols = 64; + } + unsigned int keep_lo = smem_index_valid[valid_word_base + tile * 4]; + unsigned int keep_hi = smem_index_valid[valid_word_base + tile * 4 + 1]; + uint32_t _thresh_mask_0; + { + int _lim_0 = valid_cols; + if (_lim_0 <= 0) { _thresh_mask_0 = 0u; } + else if (_lim_0 >= 32) { _thresh_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_thresh_mask_0) : "r"(_lim_0)); + } + } + uint32_t _thresh_mask_1; + { + int _lim_1 = valid_cols - 32; + if (_lim_1 <= 0) { _thresh_mask_1 = 0u; } + else if (_lim_1 >= 32) { _thresh_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_thresh_mask_1) : "r"(_lim_1)); + } + } + if (!(_thresh_mask_0 & keep_lo & (1u << 0))) _tmem_load_0[0] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 1))) _tmem_load_0[1] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 2))) _tmem_load_0[2] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 3))) _tmem_load_0[3] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 4))) _tmem_load_0[4] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 5))) _tmem_load_0[5] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 6))) _tmem_load_0[6] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 7))) _tmem_load_0[7] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 8))) _tmem_load_0[8] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 9))) _tmem_load_0[9] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 10))) _tmem_load_0[10] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 11))) _tmem_load_0[11] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 12))) _tmem_load_0[12] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 13))) _tmem_load_0[13] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 14))) _tmem_load_0[14] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 15))) _tmem_load_0[15] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 16))) _tmem_load_0[16] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 17))) _tmem_load_0[17] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 18))) _tmem_load_0[18] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 19))) _tmem_load_0[19] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 20))) _tmem_load_0[20] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 21))) _tmem_load_0[21] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 22))) _tmem_load_0[22] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 23))) _tmem_load_0[23] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 24))) _tmem_load_0[24] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 25))) _tmem_load_0[25] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 26))) _tmem_load_0[26] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 27))) _tmem_load_0[27] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 28))) _tmem_load_0[28] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 29))) _tmem_load_0[29] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 30))) _tmem_load_0[30] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 31))) _tmem_load_0[31] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 0))) _tmem_load_0[32] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 1))) _tmem_load_0[33] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 2))) _tmem_load_0[34] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 3))) _tmem_load_0[35] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 4))) _tmem_load_0[36] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 5))) _tmem_load_0[37] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 6))) _tmem_load_0[38] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 7))) _tmem_load_0[39] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 8))) _tmem_load_0[40] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 9))) _tmem_load_0[41] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 10))) _tmem_load_0[42] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 11))) _tmem_load_0[43] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 12))) _tmem_load_0[44] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 13))) _tmem_load_0[45] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 14))) _tmem_load_0[46] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 15))) _tmem_load_0[47] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 16))) _tmem_load_0[48] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 17))) _tmem_load_0[49] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 18))) _tmem_load_0[50] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 19))) _tmem_load_0[51] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 20))) _tmem_load_0[52] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 21))) _tmem_load_0[53] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 22))) _tmem_load_0[54] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 23))) _tmem_load_0[55] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 24))) _tmem_load_0[56] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 25))) _tmem_load_0[57] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 26))) _tmem_load_0[58] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 27))) _tmem_load_0[59] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 28))) _tmem_load_0[60] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 29))) _tmem_load_0[61] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 30))) _tmem_load_0[62] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 31))) _tmem_load_0[63] = -CAKE_INF; + float2 _reg_reduce_max2_2 = {-CAKE_INF, -CAKE_INF}; + row_max_x32_accum(&_tmem_load_0[0], _reg_reduce_max2_2); + row_max_x32_accum(&_tmem_load_0[32], _reg_reduce_max2_2); + float _tmem_load_0_max = row_max_reduce(_reg_reduce_max2_2); + float _max_1 = max_noftz(_tmem_load_0_max, row_max_val); + float new_max = _max_1; + smem_softmax_exchange[exchange_idx] = new_max; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + float _max_2 = max_noftz(new_max, smem_softmax_exchange[exchange_idx ^ 64]); + new_max = _max_2; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + float no_correction = (((new_max - row_max_val) * softmax_scale_log2 <= 0.0f) ? 1.0f : 0.0f); + float delta = softmax_scale_log2 * (row_max_val - new_max); + float _exp2_0 = approx_exp2(delta); + float exp_delta = _exp2_0; + float acc_scale = ((row_max_val > -CAKE_INF) ? exp_delta : 1.0f); + row_max_val = new_max; + float safe_max = ((new_max == -CAKE_INF) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + const float2 _fma_b2_3 = {softmax_scale_log2, softmax_scale_log2}; + const float2 _fma_c2_4 = {-max_scaled, -max_scaled}; + #pragma unroll + for (int _lf = 0; _lf < 32; _lf++) + fma_f32x2_inplace(&reinterpret_cast(_tmem_load_0)[_lf], _fma_b2_3, _fma_c2_4); + #pragma unroll + for (int _le = 0; _le < 64; _le++) { + _tmem_load_0[_le] = approx_exp2(_tmem_load_0[_le]); + } + float2 _reg_reduce_sum2_5 = make_float2(0.0f, 0.0f); + softmax_block_sum(&_tmem_load_0[0], &_reg_reduce_sum2_5); + softmax_block_sum(&_tmem_load_0[32], &_reg_reduce_sum2_5); + float _tmem_load_0_sum = _reg_reduce_sum2_5.x + _reg_reduce_sum2_5.y; + row_sum_val = row_sum_val * acc_scale + _tmem_load_0_sum; + int stats_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + mbarrier_wait(stats_empty_addr + (score_phase) * 8, stats_empty_phase); + float meta[4]; + meta[0] = row_sum_val; + meta[1] = row_max_val; + meta[2] = acc_scale; + meta[3] = no_correction; + int meta_addr = taddr + 384 + (unsigned int)(score_phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(head_row_base << 16); + tmem_st_x4_f32(meta_addr, meta); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (score_phase) * 8); + int p_stage = pipeline_tile % 2; + int p_empty_phase = pipeline_tile / 2 & 1 ^ 1; + mbarrier_wait(p_empty_addr + (p_stage) * 8, p_empty_phase); + uint32_t _tmem_load_0_bf16[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_0[_lp*2 + 0], _tmem_load_0[_lp*2+1 + 0])); + _tmem_load_0_bf16[_lp] = *(uint32_t*)&_bf2; + } + int smem_p_stage = p_stage * 2 + n_half; + int p_base_smem = smem_p_addr + (unsigned int)(smem_p_stage * 8192); + #pragma unroll + for (int vec = 0; vec < 8; vec++) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_base_smem + (my_row * 128 + vec * 16 ^ (my_row * 128 + vec * 16 >> 7 & 7) << 4))), "r"(_tmem_load_0_bf16[vec * 4]), "r"(_tmem_load_0_bf16[vec * 4 + 1]), "r"(_tmem_load_0_bf16[vec * 4 + 2]), "r"(_tmem_load_0_bf16[vec * 4 + 3]) : "memory"); + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(p_full_addr + (unsigned int)(p_stage * 8)), "r"(0) : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(s_empty_addr + (unsigned int)(score_phase * 8)), "r"(0) : "memory"); + } + mbarrier_arrive(index_valid_empty_addr + (valid_slot) * 8); + valid_iter = valid_iter + 1; + tile_cursor = tile_cursor + num_kv_tiles; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + } + } + // ---- Role: correction_wg ---- + if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); + { // correction_wg_main + const int correction_dummy = 0; + float softmax_scale_log2_1 = bmm1_scale[0] * 1.4426950408889634f; + float output_scale = bmm2_scale[0]; + const int head_row_base_1 = warp % 2 * 32; + const int output_row_base = warp % 4 * 32; + const int n_half_1 = warp % 4 / 2; + const int my_row_1 = head_row_base_1 + lane; + const int corr_row = head_row_base_1 << 16; + int tile_cursor_1 = 0; + unsigned int _phase_o_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = cluster_id; work_idx_1 < total_work_items; work_idx_1 += num_clusters) { + int split_idx_1 = work_idx_1 % (unsigned int)(((1) ? 4 : 1)); + int query_idx_1 = work_idx_1 / (unsigned int)(((1) ? 4 : 1)); + int num_kv_tiles_1 = ((sparse_topk + 128 - 1) / 128 + ((1) ? 4 : 1) - 1) / ((1) ? 4 : 1); + float final_sum = 0.0f; + float final_max = -CAKE_INF; + #pragma unroll 1 + for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { + int pipeline_tile_1 = tile_cursor_1 + tile_1; + int score_phase_1 = pipeline_tile_1 & 1; + int stats_phase = pipeline_tile_1 >> 1 & 1; + mbarrier_wait(stats_addr + (score_phase_1) * 8, stats_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int meta_addr_1 = taddr + 384 + (unsigned int)(score_phase_1 * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row; + float _tmem_load_1[4]; + tmem_ld_x4(&_tmem_load_1[0], meta_addr_1); + final_sum = _tmem_load_1[0]; + final_max = _tmem_load_1[1]; + float acc_scale_1 = _tmem_load_1[2]; + float no_correction_1 = _tmem_load_1[3]; + if (tile_1 > 0) { + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _vote_0 = __all_sync(0xFFFFFFFF, no_correction_1 == 1.0f); + int skip_correction = _vote_0; + if (skip_correction == 0) { + #pragma unroll + for (int local_slice = 0; local_slice < 2; local_slice++) { + #pragma unroll + for (int acc_stage = 0; acc_stage < 2; acc_stage++) { + int o_base = taddr + 128 + (unsigned int)(acc_stage * 128) + (unsigned int)(local_slice * 64) + (unsigned int)(output_row_base << 16); + #pragma unroll + for (int c = 0; c < 64; c += 16) { + float _tmem_load_2[16]; + tmem_ld_x16(&_tmem_load_2[0], o_base + c); + const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + mul_f32x2_inplace(&reinterpret_cast(_tmem_load_2)[_ls], _scale2_0); + tmem_st_x16_f32(o_base + c, _tmem_load_2); + } + } + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + } + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + } + mbarrier_arrive(stats_empty_addr + (score_phase_1) * 8); + } + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + const int exchange_idx_1 = n_half_1 * 64 + my_row_1; + smem_epilogue_exchange[exchange_idx_1] = final_sum; + asm volatile("barrier.sync 4, 128;" ::: "memory"); + final_sum = final_sum + smem_epilogue_exchange[exchange_idx_1 ^ 64]; + asm volatile("barrier.sync 4, 128;" ::: "memory"); + float _rcp_0 = approx_rcp(final_sum); + float inv_sum = ((final_sum > 0.0f) ? _rcp_0 : 0.0f); + int head_idx_1 = cta_rank * 64 + my_row_1; + int output_offset = ((1) ? ((query_idx_1 * num_heads + head_idx_1) * ((1) ? 4 : 1) + split_idx_1) * 512 : (query_idx_1 * num_heads + head_idx_1) * 512); + if (n_half_1 == 0 && head_idx_1 < num_heads) { + int stat_offset = (query_idx_1 * num_heads + head_idx_1) * ((1) ? 4 : 1) + split_idx_1; + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(final_sum)); + partial_lse[stat_offset] = ((final_sum > 0.0f) ? final_max * softmax_scale_log2_1 + _log2_0 : -CAKE_INF); + } + float epi_scale = inv_sum * output_scale; + int epi_box_base = smem_q_addr + (unsigned int)(warp % 4 * 16384); + int epi_row0 = query_idx_1 * num_heads + cta_rank * 64 + head_row_base_1; + #pragma unroll + for (int acc_stage_1 = 0; acc_stage_1 < 2; acc_stage_1++) { + #pragma unroll + for (int local_slice_1 = 0; local_slice_1 < 2; local_slice_1++) { + int logical_slice = acc_stage_1 * 2 * 2 + n_half_1 * 2 + local_slice_1; + int epi_box = epi_box_base + (acc_stage_1 * 2 + local_slice_1) * 4096; + int o_base_1 = taddr + 128 + (unsigned int)(acc_stage_1 * 128) + (unsigned int)(local_slice_1 * 64) + (unsigned int)(output_row_base << 16); + #pragma unroll + for (int c_1 = 0; c_1 < 64; c_1 += 32) { + float _tmem_load_3[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]) + : "r"(o_base_1 + c_1)); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + float epi_scaled[32]; + #pragma unroll + for (int e = 0; e < 32; e++) { + epi_scaled[e] = _tmem_load_3[e] * epi_scale; + } + unsigned int epi_packed[16]; + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(epi_scaled[_lp*2 + 0], epi_scaled[_lp*2+1 + 0])); + epi_packed[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((epi_box + (lane * 128 + c_1 * 2 ^ (lane * 128 + c_1 * 2 >> 7 & 7) << 4))), "r"(epi_packed[0]), "r"(epi_packed[1]), "r"(epi_packed[2]), "r"(epi_packed[3]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((epi_box + (lane * 128 + (c_1 * 2 + 16) ^ (lane * 128 + (c_1 * 2 + 16) >> 7 & 7) << 4))), "r"(epi_packed[4]), "r"(epi_packed[5]), "r"(epi_packed[6]), "r"(epi_packed[7]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((epi_box + (lane * 128 + (c_1 * 2 + 32) ^ (lane * 128 + (c_1 * 2 + 32) >> 7 & 7) << 4))), "r"(epi_packed[8]), "r"(epi_packed[9]), "r"(epi_packed[10]), "r"(epi_packed[11]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((epi_box + (lane * 128 + (c_1 * 2 + 48) ^ (lane * 128 + (c_1 * 2 + 48) >> 7 & 7) << 4))), "r"(epi_packed[12]), "r"(epi_packed[13]), "r"(epi_packed[14]), "r"(epi_packed[15]) : "memory"); + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + __syncwarp(); + if (elect_sync()) { + tma_store_4d((&O), 0, logical_slice, split_idx_1, epi_row0, epi_box); + asm volatile("cp.async.bulk.commit_group;"); + } + } + } + if (elect_sync()) { + asm volatile("cp.async.bulk.wait_group.read 0;"); + mbarrier_arrive(epi_q_free_addr); + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + tile_cursor_1 = tile_cursor_1 + num_kv_tiles_1; + } + if (elect_sync()) { + asm volatile("cp.async.bulk.wait_group 0;"); + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + } + } + // ---- Role: mma_warp ---- + if (warp == 8) { + { // mma_warp_main + const int mma_dummy = 0; + unsigned int k_stage = 0; + unsigned int v_stage = 0; + int tile_cursor_2 = 0; + unsigned int _phase_o_first_slice_seeded_0 = 0; + mbarrier_wait(o_first_slice_seeded_addr, _phase_o_first_slice_seeded_0); + _phase_o_first_slice_seeded_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + unsigned int _phase_q_full_0 = 0; + unsigned int _phase_k_full = 0; + unsigned int _phase_o_empty_0 = 1; + unsigned int _phase_v_full = 0; + #pragma unroll 1 + for (unsigned int work_idx_2 = cluster_id; work_idx_2 < total_work_items; work_idx_2 += num_clusters) { + int num_kv_tiles_2 = ((sparse_topk + 128 - 1) / 128 + ((1) ? 4 : 1) - 1) / ((1) ? 4 : 1); + if (cta_rank == 0) { + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + int first_pv = 1; + #pragma unroll 1 + for (int tile_2 = 0; tile_2 < num_kv_tiles_2; tile_2++) { + int pipeline_tile_2 = tile_cursor_2 + tile_2; + int score_phase_2 = pipeline_tile_2 & 1; + int s_empty_phase = pipeline_tile_2 >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (score_phase_2) * 8, s_empty_phase); + int score_col = ((score_phase_2 != 0) ? 64 : 0); + #pragma unroll + for (int qk_stage = 0; qk_stage < 4; qk_stage++) { + mbarrier_wait(k_full_addr + (k_stage) * 8, _phase_k_full); + int _mma_a_lo_0 = (((smem_q_addr) >> 4) & 0x3FFF) + (qk_stage * 2) * 512; + int _mma_b_lo_0 = (((smem_k_addr) >> 4) & 0x3FFF) + (k_stage) * 1024; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (score_col))), "r"(((qk_stage == 0) ? 0 : 1))); + int _mma_a_lo_1 = (((smem_q_addr) >> 4) & 0x3FFF) + (qk_stage * 2 + 1) * 512; + int _mma_b_lo_1 = (((smem_k_addr + 8192) >> 4) & 0x3FFF) + (k_stage) * 1024; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + elect_commit_cg2_multicast(k_empty_addr + (k_stage) * 8, (uint16_t)(3)); + k_stage += 1; + if (k_stage == 3) { k_stage = 0; _phase_k_full ^= 1; } + } + elect_commit_cg2_multicast(s_full_addr + (score_phase_2) * 8, (uint16_t)(3)); + if (tile_2 > 0) { + int prev_pipeline_tile = pipeline_tile_2 - 1; + int p_stage_1 = prev_pipeline_tile % 2; + int p_full_phase = prev_pipeline_tile / 2 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (p_stage_1) * 8, p_full_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps = 0; ps < 2; ps++) { + int smem_p_stage_1 = p_stage_1 * 2 + ps; + mbarrier_wait(v_full_addr + (v_stage) * 8, _phase_v_full); + #pragma unroll + for (int acc_stage_2 = 0; acc_stage_2 < 2; acc_stage_2++) { + int output_col = 128 + acc_stage_2 * 128; + if (acc_stage_2 == 0) { + int _mma_a_lo_2 = (((smem_p_addr) >> 4) & 0x3FFF) + (smem_p_stage_1) * 512; + int _mma_b_lo_2 = ((((smem_v_lo_addr) >> 4) & 0x3FFF) | 0x2000000) + (v_stage) * 2048; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (output_col))), "r"(((((ps == 0) ? first_pv : 0)) ? 0 : 1))); + } else { + int _mma_a_lo_3 = (((smem_p_addr) >> 4) & 0x3FFF) + (smem_p_stage_1) * 512; + int _mma_b_lo_3 = ((((smem_v_hi_addr) >> 4) & 0x3FFF) | 0x2000000) + (v_stage) * 2048; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"((tmem_tmem_scratch + (output_col))), "r"(((((ps == 0) ? first_pv : 0)) ? 0 : 1))); + } + } + elect_commit_cg2_multicast(v_empty_addr + (v_stage) * 8, (uint16_t)(3)); + v_stage += 1; + if (v_stage == 2) { v_stage = 0; _phase_v_full ^= 1; } + } + first_pv = 0; + elect_commit_cg2_multicast(p_empty_addr + (p_stage_1) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + } + } + int last_pipeline_tile = tile_cursor_2 + num_kv_tiles_2 - 1; + int last_p_stage = last_pipeline_tile % 2; + int last_p_phase = last_pipeline_tile / 2 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (last_p_stage) * 8, last_p_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_1 = 0; ps_1 < 2; ps_1++) { + int smem_p_stage_2 = last_p_stage * 2 + ps_1; + mbarrier_wait(v_full_addr + (v_stage) * 8, _phase_v_full); + #pragma unroll + for (int acc_stage_3 = 0; acc_stage_3 < 2; acc_stage_3++) { + int output_col_1 = 128 + acc_stage_3 * 128; + if (acc_stage_3 == 0) { + int _mma_a_lo_4 = (((smem_p_addr) >> 4) & 0x3FFF) + (smem_p_stage_2) * 512; + int _mma_b_lo_4 = ((((smem_v_lo_addr) >> 4) & 0x3FFF) | 0x2000000) + (v_stage) * 2048; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"((tmem_tmem_scratch + (output_col_1))), "r"(((((ps_1 == 0) ? first_pv : 0)) ? 0 : 1))); + } else { + int _mma_a_lo_5 = (((smem_p_addr) >> 4) & 0x3FFF) + (smem_p_stage_2) * 512; + int _mma_b_lo_5 = ((((smem_v_hi_addr) >> 4) & 0x3FFF) | 0x2000000) + (v_stage) * 2048; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_5), "r"(_mma_b_lo_5), "r"((tmem_tmem_scratch + (output_col_1))), "r"(((((ps_1 == 0) ? first_pv : 0)) ? 0 : 1))); + } + } + elect_commit_cg2_multicast(v_empty_addr + (v_stage) * 8, (uint16_t)(3)); + v_stage += 1; + if (v_stage == 2) { v_stage = 0; _phase_v_full ^= 1; } + } + elect_commit_cg2_multicast(p_empty_addr + (last_p_stage) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + } + tile_cursor_2 = tile_cursor_2 + num_kv_tiles_2; + } + if (cta_rank == 0) { + #pragma unroll + for (int tail_offset = 0; tail_offset < 2; tail_offset++) { + int tail_tile = tile_cursor_2 + tail_offset; + int tail_stage = tail_tile & 1; + int tail_phase = tail_tile >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (tail_stage) * 8, tail_phase); + } + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + } + asm volatile("tcgen05.fence::before_thread_sync;"); + int peer_rank_1 = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_dealloc_peer_addr), "r"(peer_rank_1) : "memory"); + mbarrier_wait(tmem_dealloc_peer_addr, 0); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + } + // ---- Role: load_wg ---- + if (warp >= 9 && warp <= 12) { + { // load_wg_main + const int load_dummy = 0; + const int load_warp_rank = warp - 9; + unsigned int k_stage_1 = 0; + unsigned int v_stage_1 = 0; + unsigned int _phase_q_empty_0 = 1; + unsigned int _phase_epi_q_free_0 = 1; + unsigned int _phase_k_empty = 1; + unsigned int _phase_v_empty = 1; + #pragma unroll 1 + for (unsigned int work_idx_3 = cluster_id; work_idx_3 < total_work_items; work_idx_3 += num_clusters) { + int split_idx_2 = work_idx_3 % (unsigned int)(((1) ? 4 : 1)); + int query_idx_2 = work_idx_3 / (unsigned int)(((1) ? 4 : 1)); + int num_kv_tiles_3 = ((sparse_topk + 128 - 1) / 128 + ((1) ? 4 : 1) - 1) / ((1) ? 4 : 1); + int first_kv_tile_1 = split_idx_2 * num_kv_tiles_3; + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + mbarrier_wait(epi_q_free_addr, _phase_epi_q_free_0); + _phase_epi_q_free_0 ^= 1; + if (load_warp_rank == 0) { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(131072)) : "memory"); + } + } + if (elect_sync()) { + #pragma unroll + for (int q_stage = 0; q_stage < 8; q_stage++) { + tma_4d_gmem2smem_cta2(smem_q_addr + (unsigned int)(q_stage * 8192), (&tmap_q), 0, cta_rank * 64, q_stage, query_idx_2, ((q_full_addr) & 0xFEFFFFFF)); + } + } + } + asm volatile("barrier.sync 2, 160;" ::: "memory"); + #pragma unroll + for (int qk_stage_1 = 0; qk_stage_1 < 4; qk_stage_1++) { + mbarrier_wait(k_empty_addr + (k_stage_1) * 8, _phase_k_empty); + if (load_warp_rank == 0) { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + } + } + } + int group = load_warp_rank * 4 + lane; + if (lane < 4) { + int index_offset = cta_rank * 64 + group * 4; + int raw_rows[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows[0])), "=r"(*reinterpret_cast(&raw_rows[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows[(0) + 3])) + : "r"(smem_indices_addr + (unsigned int)(index_offset * 4))); + if (first_kv_tile_1 == 0) { + tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group * 512), (&tmap_swa_k), qk_stage_1 * 128, ((raw_rows[0] >= 0) ? raw_rows[0] : ((1) ? group * 4 : 0)), ((raw_rows[1] >= 0) ? raw_rows[1] : ((1) ? group * 4 + 1 : 0)), ((raw_rows[2] >= 0) ? raw_rows[2] : ((1) ? group * 4 + 2 : 0)), ((raw_rows[3] >= 0) ? raw_rows[3] : ((1) ? group * 4 + 3 : 0)), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group * 512) + 8192, (&tmap_swa_k), qk_stage_1 * 128 + 64, ((raw_rows[0] >= 0) ? raw_rows[0] : ((1) ? group * 4 : 0)), ((raw_rows[1] >= 0) ? raw_rows[1] : ((1) ? group * 4 + 1 : 0)), ((raw_rows[2] >= 0) ? raw_rows[2] : ((1) ? group * 4 + 2 : 0)), ((raw_rows[3] >= 0) ? raw_rows[3] : ((1) ? group * 4 + 3 : 0)), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + } else { + tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group * 512), (&tmap_compressed_k), qk_stage_1 * 128, ((raw_rows[0] >= 0) ? raw_rows[0] : ((1) ? group * 4 : 0)), ((raw_rows[1] >= 0) ? raw_rows[1] : ((1) ? group * 4 + 1 : 0)), ((raw_rows[2] >= 0) ? raw_rows[2] : ((1) ? group * 4 + 2 : 0)), ((raw_rows[3] >= 0) ? raw_rows[3] : ((1) ? group * 4 + 3 : 0)), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group * 512) + 8192, (&tmap_compressed_k), qk_stage_1 * 128 + 64, ((raw_rows[0] >= 0) ? raw_rows[0] : ((1) ? group * 4 : 0)), ((raw_rows[1] >= 0) ? raw_rows[1] : ((1) ? group * 4 + 1 : 0)), ((raw_rows[2] >= 0) ? raw_rows[2] : ((1) ? group * 4 + 2 : 0)), ((raw_rows[3] >= 0) ? raw_rows[3] : ((1) ? group * 4 + 3 : 0)), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + } + } + k_stage_1 += 1; + if (k_stage_1 == 3) { k_stage_1 = 0; _phase_k_empty ^= 1; } + } + #pragma unroll 1 + for (int tile_3 = 1; tile_3 < num_kv_tiles_3; tile_3++) { + #pragma unroll + for (int qk_stage_2 = 0; qk_stage_2 < 4; qk_stage_2++) { + mbarrier_wait(k_empty_addr + (k_stage_1) * 8, _phase_k_empty); + if (load_warp_rank == 0) { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + } + } + } + int group_1 = load_warp_rank * 4 + lane; + if (lane < 4) { + int index_offset_1 = tile_3 * 128 + cta_rank * 64 + group_1 * 4; + int raw_rows_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_1[0])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 3])) + : "r"(smem_indices_addr + (unsigned int)(index_offset_1 * 4))); + tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group_1 * 512), (&tmap_compressed_k), qk_stage_2 * 128, ((raw_rows_1[0] >= 0) ? raw_rows_1[0] : ((1) ? group_1 * 4 : 0)), ((raw_rows_1[1] >= 0) ? raw_rows_1[1] : ((1) ? group_1 * 4 + 1 : 0)), ((raw_rows_1[2] >= 0) ? raw_rows_1[2] : ((1) ? group_1 * 4 + 2 : 0)), ((raw_rows_1[3] >= 0) ? raw_rows_1[3] : ((1) ? group_1 * 4 + 3 : 0)), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group_1 * 512) + 8192, (&tmap_compressed_k), qk_stage_2 * 128 + 64, ((raw_rows_1[0] >= 0) ? raw_rows_1[0] : ((1) ? group_1 * 4 : 0)), ((raw_rows_1[1] >= 0) ? raw_rows_1[1] : ((1) ? group_1 * 4 + 1 : 0)), ((raw_rows_1[2] >= 0) ? raw_rows_1[2] : ((1) ? group_1 * 4 + 2 : 0)), ((raw_rows_1[3] >= 0) ? raw_rows_1[3] : ((1) ? group_1 * 4 + 3 : 0)), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + } + k_stage_1 += 1; + if (k_stage_1 == 3) { k_stage_1 = 0; _phase_k_empty ^= 1; } + } + int prev_tile = tile_3 - 1; + #pragma unroll + for (int ps_2 = 0; ps_2 < 2; ps_2++) { + mbarrier_wait(v_empty_addr + (v_stage_1) * 8, _phase_v_empty); + if (load_warp_rank == 0) { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(65536)) : "memory"); + } + } + } + int group_2 = load_warp_rank * 4 + lane; + if (lane < 4) { + int index_offset_2 = prev_tile * 128 + ps_2 * 64 + group_2 * 4; + int raw_rows_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_2[0])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 3])) + : "r"(smem_indices_addr + (unsigned int)(index_offset_2 * 4))); + #pragma unroll + for (int acc_stage_4 = 0; acc_stage_4 < 2; acc_stage_4++) { + int dst_v = ((acc_stage_4 == 0) ? smem_v_lo_addr + v_stage_1 * 32768 : smem_v_hi_addr + v_stage_1 * 32768); + if (first_kv_tile_1 + prev_tile == 0) { + tma_gather4_gmem2smem_mc_cta2(dst_v + group_2 * 512, (&tmap_swa_v), acc_stage_4 * 128 * 2 + cta_rank * 128, ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : ((1) ? group_2 * 4 : 0)), ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : ((1) ? group_2 * 4 + 1 : 0)), ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : ((1) ? group_2 * 4 + 2 : 0)), ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : ((1) ? group_2 * 4 + 3 : 0)), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + tma_gather4_gmem2smem_mc_cta2(dst_v + group_2 * 512 + 8192, (&tmap_swa_v), acc_stage_4 * 128 * 2 + cta_rank * 128 + 64, ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : ((1) ? group_2 * 4 : 0)), ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : ((1) ? group_2 * 4 + 1 : 0)), ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : ((1) ? group_2 * 4 + 2 : 0)), ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : ((1) ? group_2 * 4 + 3 : 0)), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + } else { + tma_gather4_gmem2smem_mc_cta2(dst_v + group_2 * 512, (&tmap_compressed_v), acc_stage_4 * 128 * 2 + cta_rank * 128, ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : ((1) ? group_2 * 4 : 0)), ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : ((1) ? group_2 * 4 + 1 : 0)), ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : ((1) ? group_2 * 4 + 2 : 0)), ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : ((1) ? group_2 * 4 + 3 : 0)), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + tma_gather4_gmem2smem_mc_cta2(dst_v + group_2 * 512 + 8192, (&tmap_compressed_v), acc_stage_4 * 128 * 2 + cta_rank * 128 + 64, ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : ((1) ? group_2 * 4 : 0)), ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : ((1) ? group_2 * 4 + 1 : 0)), ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : ((1) ? group_2 * 4 + 2 : 0)), ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : ((1) ? group_2 * 4 + 3 : 0)), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + } + } + } + v_stage_1 += 1; + if (v_stage_1 == 2) { v_stage_1 = 0; _phase_v_empty ^= 1; } + } + } + int last_tile = num_kv_tiles_3 - 1; + #pragma unroll + for (int ps_3 = 0; ps_3 < 2; ps_3++) { + mbarrier_wait(v_empty_addr + (v_stage_1) * 8, _phase_v_empty); + if (load_warp_rank == 0) { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(65536)) : "memory"); + } + } + } + int group_3 = load_warp_rank * 4 + lane; + if (lane < 4) { + int index_offset_3 = last_tile * 128 + ps_3 * 64 + group_3 * 4; + int raw_rows_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_3[0])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 3])) + : "r"(smem_indices_addr + (unsigned int)(index_offset_3 * 4))); + #pragma unroll + for (int acc_stage_5 = 0; acc_stage_5 < 2; acc_stage_5++) { + int dst_v_1 = ((acc_stage_5 == 0) ? smem_v_lo_addr + v_stage_1 * 32768 : smem_v_hi_addr + v_stage_1 * 32768); + if (first_kv_tile_1 + last_tile == 0) { + tma_gather4_gmem2smem_mc_cta2(dst_v_1 + group_3 * 512, (&tmap_swa_v), acc_stage_5 * 128 * 2 + cta_rank * 128, ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : ((1) ? group_3 * 4 : 0)), ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : ((1) ? group_3 * 4 + 1 : 0)), ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : ((1) ? group_3 * 4 + 2 : 0)), ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : ((1) ? group_3 * 4 + 3 : 0)), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + tma_gather4_gmem2smem_mc_cta2(dst_v_1 + group_3 * 512 + 8192, (&tmap_swa_v), acc_stage_5 * 128 * 2 + cta_rank * 128 + 64, ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : ((1) ? group_3 * 4 : 0)), ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : ((1) ? group_3 * 4 + 1 : 0)), ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : ((1) ? group_3 * 4 + 2 : 0)), ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : ((1) ? group_3 * 4 + 3 : 0)), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + } else { + tma_gather4_gmem2smem_mc_cta2(dst_v_1 + group_3 * 512, (&tmap_compressed_v), acc_stage_5 * 128 * 2 + cta_rank * 128, ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : ((1) ? group_3 * 4 : 0)), ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : ((1) ? group_3 * 4 + 1 : 0)), ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : ((1) ? group_3 * 4 + 2 : 0)), ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : ((1) ? group_3 * 4 + 3 : 0)), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + tma_gather4_gmem2smem_mc_cta2(dst_v_1 + group_3 * 512 + 8192, (&tmap_compressed_v), acc_stage_5 * 128 * 2 + cta_rank * 128 + 64, ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : ((1) ? group_3 * 4 : 0)), ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : ((1) ? group_3 * 4 + 1 : 0)), ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : ((1) ? group_3 * 4 + 2 : 0)), ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : ((1) ? group_3 * 4 + 3 : 0)), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + } + } + } + v_stage_1 += 1; + if (v_stage_1 == 2) { v_stage_1 = 0; _phase_v_empty ^= 1; } + } + asm volatile("barrier.sync 2, 160;" ::: "memory"); + } + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + } + } + // ---- Role: index_warp ---- + if (warp == 13) { + { // index_warp_main + const int index_dummy = 0; + int valid_iter_1 = 0; + #pragma unroll 1 + for (unsigned int work_idx_4 = cluster_id; work_idx_4 < total_work_items; work_idx_4 += num_clusters) { + int split_idx_3 = work_idx_4 % (unsigned int)(((1) ? 4 : 1)); + int query_idx_3 = work_idx_4 / (unsigned int)(((1) ? 4 : 1)); + int owner_col_base = split_idx_3 * (4 / ((1) ? 4 : 1)) * 128; + int* swa_row = swa_indices + (query_idx_3 * swa_index_stride); + int* compressed_row = compressed_indices + (query_idx_3 * compressed_index_stride); + int valid_slot_1 = valid_iter_1 & 1; + int valid_empty_phase = valid_iter_1 >> 1 & 1 ^ 1; + mbarrier_wait(index_valid_empty_addr + (valid_slot_1) * 8, valid_empty_phase); + #pragma unroll + for (int index_pass = 0; index_pass < 4 / ((1) ? 4 : 1); index_pass++) { + int index_offset_4 = index_pass * 128 + lane * 4; + int combined_col = owner_col_base + index_offset_4; + int* index_src = ((combined_col < 128) ? swa_row + combined_col : compressed_row + (combined_col - 128)); + int values[4]; + #pragma unroll + for (int i = 0; i < 4; i++) { + values[i] = -1; + } + if (combined_col + 4 <= sparse_topk) { + int _vec_load_0[4]; + { + const int4* _ivptr_0 = reinterpret_cast(index_src + 0); + int4 _ivld_0; + _ivld_0 = *_ivptr_0; + _vec_load_0[0 + 0] = _ivld_0.x; + _vec_load_0[0 + 1] = _ivld_0.y; + _vec_load_0[0 + 2] = _ivld_0.z; + _vec_load_0[0 + 3] = _ivld_0.w; + } + #pragma unroll + for (int i_1 = 0; i_1 < 4; i_1++) { + values[i_1] = _vec_load_0[i_1]; + } + } + if (combined_col < sparse_topk && combined_col + 4 > sparse_topk) { + #pragma unroll + for (int i_2 = 0; i_2 < 4; i_2++) { + if (combined_col + i_2 < sparse_topk) { + values[i_2] = index_src[i_2]; + } + } + } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(smem_indices_addr + (unsigned int)(index_offset_4 * 4)), "r"(values[0]), "r"(values[1]), "r"(values[2]), "r"(values[3]) : "memory"); + unsigned int valid_word = (((values[0] >= 0) ? 1 : 0) | ((values[1] >= 0) ? 1 : 0) << 1 | ((values[2] >= 0) ? 1 : 0) << 2 | ((values[3] >= 0) ? 1 : 0) << 3) << (lane & 7) * 4; + unsigned int _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, valid_word, 1); + valid_word = valid_word | _shfl_xor_0; + unsigned int _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, valid_word, 2); + valid_word = valid_word | _shfl_xor_1; + unsigned int _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, valid_word, 4); + valid_word = valid_word | _shfl_xor_2; + if ((lane & 7) == 0) { + smem_index_valid[valid_slot_1 * 36 + index_pass * 4 + (lane >> 3)] = valid_word; + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(index_valid_full_addr + (valid_slot_1) * 8); + asm volatile("barrier.sync 2, 160;" ::: "memory"); + asm volatile("barrier.sync 2, 160;" ::: "memory"); + valid_iter_1 = valid_iter_1 + 1; + } + } + } + // ---- Role: empty_warps ---- + if (warp >= 14 && warp <= 15) { + { // empty_warps_main + const int empty_dummy = 0; + #pragma unroll 1 + for (unsigned int work_idx_5 = cluster_id; work_idx_5 < total_work_items; work_idx_5 += num_clusters) { + asm volatile("tcgen05.fence::after_thread_sync;"); + } + } + } + + // Cleanup +} + +} // extern "C" diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_530fe7e6cb8c7d09744d_binding.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_530fe7e6cb8c7d09744d_binding.cu new file mode 100644 index 00000000000..51916b01598 --- /dev/null +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_530fe7e6cb8c7d09744d_binding.cu @@ -0,0 +1,445 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +// Generated by the Cake source exporter — do not edit. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_530fe7e6cb8c7d09744d'. +#include +#include +#include +#include + +#include "tvm_ffi_utils.h" + +#include +#include +#include +#include +#include +#include +#include + +extern "C" __global__ void kernel_cake_dsv4_530fe7e6cb8c7d09744d(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items); + + +namespace cake_host_shim_ed72cc11693a1ade { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t CakeDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +// 4D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + int64_t d3 = t.size(t.ndim() - 3); + TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + TVM_FFI_CHECK(d1 % 128 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 128; + uint64_t global_dim[4] = {(uint64_t)(128), (uint64_t)(d2), (uint64_t)((d1 / 128)), (uint64_t)(d3)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (128, 128, 1, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 128; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + (uint64_t)((carrier_stride_1 * 8) / 8), + (uint64_t)((carrier_stride_2 * 8) / 8), + }; + uint32_t box_dim[4] = {128u, 128u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_swa_kv' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_swa_kv' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; + TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (128, 1) exceeds resolved global dims for 'tmap_swa_kv'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + }; + uint32_t box_dim[2] = {128u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_compressed_kv' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_compressed_kv(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_compressed_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_compressed_kv' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_compressed_kv' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_compressed_kv"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_compressed_kv' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved a non-positive global dim"; + TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (128, 1) exceeds resolved global dims for 'tmap_compressed_kv'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + }; + uint32_t box_dim[2] = {128u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_kv') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_compressed_kv, TensorView arg_O, TensorView arg_partial_lse, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_num_query_tokens, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 1, 8, 1); + CheckContiguous(arg_tmap_q, "tmap_q"); + CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); + CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 1, 8, 1); + CheckCudaTensor(arg_tmap_compressed_kv, "tmap_compressed_kv"); + CheckDtype(arg_tmap_compressed_kv, "tmap_compressed_kv", 1, 8, 1); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_partial_lse, "partial_lse"); + CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); + CheckContiguous(arg_partial_lse, "partial_lse"); + CheckCudaTensor(arg_swa_indices, "swa_indices"); + CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); + CheckContiguous(arg_swa_indices, "swa_indices"); + CheckCudaTensor(arg_compressed_indices, "compressed_indices"); + CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); + CheckContiguous(arg_compressed_indices, "compressed_indices"); + CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); + CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); + CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); + CheckContiguous(arg_bmm1_scale, "bmm1_scale"); + CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); + CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); + CheckContiguous(arg_bmm2_scale, "bmm2_scale"); + TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) + << "scalar 'num_heads' value " << arg_num_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) + << "scalar 'swa_index_stride' value " << arg_swa_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) + << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) + << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_query_tokens >= -2147483648LL && arg_num_query_tokens <= 2147483647LL, ValueError) + << "scalar 'num_query_tokens' value " << arg_num_query_tokens + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) + << "scalar 'sparse_topk' value " << arg_sparse_topk + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) + << "scalar 'has_sinks' value " << arg_has_sinks + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); + CheckSameCudaDevice(arg_tmap_compressed_kv, arg_tmap_q, "tmap_compressed_kv", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_partial_lse, arg_tmap_q, "partial_lse", "tmap_q"); + CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); + CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); + CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); + CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + + cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); + CUtensorMap p_tmap_compressed_kv = EncodeTma_tmap_compressed_kv(arg_tmap_compressed_kv); + __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); + float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); + int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); + int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); + int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); + float* p_sinks = static_cast(arg_sinks.data_ptr()); + float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); + float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); + int32_t v_num_heads = (int32_t)arg_num_heads; + int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; + int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; + int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; + int32_t v_num_query_tokens = (int32_t)arg_num_query_tokens; + int32_t v_sparse_topk = (int32_t)arg_sparse_topk; + int32_t v_has_sinks = (int32_t)arg_has_sinks; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_tmap_compressed_kv, &p_O, &p_partial_lse, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_num_query_tokens, &v_sparse_topk, &v_has_sinks, &v_total_work_items}; + + static std::mutex smem_mu; + static auto* smem_status_by_device = + new std::unordered_map(); + cudaError_t smem_status = cudaSuccess; + { + std::lock_guard lock(smem_mu); + auto it = smem_status_by_device->find(dev.device_id); + if (it == smem_status_by_device->end()) { + int current_device = -1; + smem_status = cudaGetDevice(¤t_device); + if (smem_status == cudaSuccess && current_device != dev.device_id) { + smem_status = cudaErrorInvalidDevice; + } + if (smem_status == cudaSuccess) { + smem_status = cudaFuncSetAttribute( + reinterpret_cast(kernel_cake_dsv4_530fe7e6cb8c7d09744d), + cudaFuncAttributeMaxDynamicSharedMemorySize, + 205824); + } + smem_status_by_device->emplace(dev.device_id, smem_status); + } else { + smem_status = it->second; + } + } + TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) + << "cudaFuncSetAttribute for kernel_cake_dsv4_530fe7e6cb8c7d09744d failed: " + << cudaGetErrorString(smem_status); + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(512u, 1u, 1u); + + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes = 205824u; + config.stream = stream; + cudaError_t launch_status = cudaLaunchKernelExC( + &config, reinterpret_cast(kernel_cake_dsv4_530fe7e6cb8c7d09744d), kargs); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernelExC for kernel_cake_dsv4_530fe7e6cb8c7d09744d failed: " + << cudaGetErrorString(launch_status); +} + +} // namespace cake_host_shim_ed72cc11693a1ade + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_ed72cc11693a1ade::Run); diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_530fe7e6cb8c7d09744d_kernel.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_530fe7e6cb8c7d09744d_kernel.cu new file mode 100644 index 00000000000..37358d2befd --- /dev/null +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_530fe7e6cb8c7d09744d_kernel.cu @@ -0,0 +1,2912 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +#define CAKE_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_SCRATCH_OFFSET 0 +#define NUM_K_PIPE_STAGES 1 +#define NUM_V_PIPE_STAGES 1 +#define NUM_KV_PIPE_STAGES 8 +#define NUM_INDEX_PIPE_STAGES 1 +#define SMEM_SMEM_Q_OFF 1024 +#define SMEM_SMEM_Q_STAGE_BYTES 16384 +#define SMEM_SMEM_Q_STRIDE 16384 +#define SMEM_SMEM_KV_OFF 66560 +#define SMEM_SMEM_KV_STAGE_BYTES 16384 +#define SMEM_SMEM_KV_STRIDE 16384 +#define SMEM_SMEM_V_OFF 66560 +#define SMEM_SMEM_V_STAGE_BYTES 16384 +#define SMEM_SMEM_V_STRIDE 16384 +#define SMEM_SMEM_Q_FULL_OFF 1024 +#define SMEM_SMEM_Q_FULL_STAGE_BYTES 32768 +#define SMEM_SMEM_Q_FULL_STRIDE 32768 +#define SMEM_SMEM_K_FULL_OFF 66560 +#define SMEM_SMEM_K_FULL_STAGE_BYTES 32768 +#define SMEM_SMEM_K_FULL_STRIDE 32768 +#define SMEM_SMEM_V_FULL_OFF 132096 +#define SMEM_SMEM_V_FULL_STAGE_BYTES 32768 +#define SMEM_SMEM_V_FULL_STRIDE 32768 +#define SMEM_SMEM_STATS_MAX_OFF 197632 +#define SMEM_SMEM_STATS_MAX_STAGE_BYTES 1024 +#define SMEM_SMEM_STATS_MAX_STRIDE 1024 +#define SMEM_SMEM_STATS_SUM_OFF 198656 +#define SMEM_SMEM_STATS_SUM_STAGE_BYTES 512 +#define SMEM_SMEM_STATS_SUM_STRIDE 512 +#define SMEM_SMEM_STATS_FINAL_MAX_OFF 199168 +#define SMEM_SMEM_STATS_FINAL_MAX_STAGE_BYTES 512 +#define SMEM_SMEM_STATS_FINAL_MAX_STRIDE 512 +#define SMEM_SMEM_SOFTMAX_WARP_PAIR_EXCHANGE_OFF 199680 +#define SMEM_SMEM_SOFTMAX_WARP_PAIR_EXCHANGE_STAGE_BYTES 1024 +#define SMEM_SMEM_SOFTMAX_WARP_PAIR_EXCHANGE_STRIDE 1024 +#define SMEM_SMEM_CORR_WARP_PAIR_EXCHANGE_OFF 200704 +#define SMEM_SMEM_CORR_WARP_PAIR_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_CORR_WARP_PAIR_EXCHANGE_STRIDE 512 +#define SMEM_SMEM_SPARSE_INDICES_OFF 201216 +#define SMEM_SMEM_SPARSE_INDICES_STAGE_BYTES 4608 +#define SMEM_SMEM_SPARSE_INDICES_STRIDE 4608 +#define SMEM_SMEM_P_FP8_OFF 1024 +#define SMEM_SMEM_P_FP8_STAGE_BYTES 8192 +#define SMEM_SMEM_P_FP8_STRIDE 8192 +#define SMEM_TOTAL 205824 +#define THREADS 512 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count) : "memory"); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f8f6f4( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x4(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x4.b32" + " {%0, %1, %2, %3}, [%4];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); +} + + +__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ float2 ex2_emulation_f32x2_value(float2 value) { + const float c0 = 1.0f, c1 = 0.695146143436431884765625f; + const float c2 = 0.227564394474029541015625f, c3 = 0.077119089663028717041015625f; + const float magic = 12582912.0f; + float x0 = max_noftz(value.x, -127.0f), x1 = max_noftz(value.y, -127.0f); + float2 xc2 = make_float2(x0, x1), magic2 = make_float2(magic, magic); + float2 xr2; + asm("add.rm.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xr2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&magic2)); + float2 c3_2 = make_float2(c3, c3), c2_2 = make_float2(c2, c2); + float2 c1_2 = make_float2(c1, c1), c0_2 = make_float2(c0, c0); + float2 xrb2, xfrac2; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xrb2) + : "l"(*(unsigned long long*)&xr2), "l"(*(unsigned long long*)&magic2)); + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xfrac2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&xrb2)); + float2 poly2; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&c3_2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c2_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c1_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c0_2)); + int x0r_i, x1r_i, p0_i, p1_i; + asm("mov.b64 {%0, %1}, %2;" : "=r"(x0r_i), "=r"(x1r_i) : "l"(*(unsigned long long*)&xr2)); + asm("mov.b64 {%0, %1}, %2;" : "=r"(p0_i), "=r"(p1_i) : "l"(*(unsigned long long*)&poly2)); + float r0, r1; + asm("mov.b32 %0, %1;" : "=f"(r0) : "r"((x0r_i << 23) + p0_i)); + asm("mov.b32 %0, %1;" : "=f"(r1) : "r"((x1r_i << 23) + p1_i)); + return make_float2(r0, r1); +} + +__device__ __forceinline__ void ex2_emulation_f32x2(float* x0_ptr, float* x1_ptr) { + float2 result = ex2_emulation_f32x2_value(make_float2(*x0_ptr, *x1_ptr)); + *x0_ptr = result.x; *x1_ptr = result.y; +} + +__device__ __forceinline__ void softmax_frag_exp2_cast( + float* sv, uint32_t* pv, int use_emu) +{ + #pragma unroll + for (int j = 0; j < 16; j++) { + if (use_emu && j >= 12) + ex2_emulation_f32x2(&sv[j*2], &sv[j*2+1]); + else { + sv[j*2] = approx_exp2(sv[j*2]); + sv[j*2+1] = approx_exp2(sv[j*2+1]); + } + } + #pragma unroll + for (int j = 0; j < 16; j++) { + __nv_bfloat162 bf = __float22bfloat162_rn({sv[j*2], sv[j*2+1]}); + pv[j] = reinterpret_cast(bf); + } +} + + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_2d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3}], [%4];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr) { + // Canonical .shared::cta form for non-multicast gather4, matching + // trtllm-gen / cuda_ptx and the PTX ISA qualifier order + // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, + // .shared::cluster is reserved for the multicast variant (ctaMask). + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem_cta2( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr) { + // Canonical .shared::cta form; see tma_gather4_gmem2smem above. + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem_mc( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr, unsigned short cta_mask) { + // Multicast variant: the PTX grammar ties the .shared::cluster + // destination to .multicast::cluster + ctaMask (cf. cuda_ptx / + // SM100_TMA_LOAD_MULTICAST_2D_GATHER4 in CUTLASS). + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" + ".mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem_mc_cta2( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr, unsigned short cta_mask) { + // Multicast + cta_group::2 variant; see tma_gather4_gmem2smem_mc. + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" + ".mbarrier::complete_tx::bytes.multicast::cluster.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_st_x8_u32(int addr, const uint32_t* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x8.b32" + " [%0], {%1,%2,%3,%4,%5,%6,%7,%8};" + :: "r"(addr), + "r"(src[0]), "r"(src[1]), "r"(src[2]), "r"(src[3]), + "r"(src[4]), "r"(src[5]), "r"(src[6]), "r"(src[7])); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(512, 1) void +kernel_cake_dsv4_530fe7e6cb8c7d09744d(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); + + const int mbar_base = smem; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define k_full_addr (mbar_base + 16) + #define k_empty_addr (mbar_base + 24) + #define v_full_addr (mbar_base + 32) + #define v_empty_addr (mbar_base + 40) + #define kv_full_addr (mbar_base + 48) + #define kv_empty_addr (mbar_base + 112) + #define index_full_addr (mbar_base + 176) + #define index_empty_addr (mbar_base + 184) + #define s_full_addr (mbar_base + 192) + #define p_full_addr (mbar_base + 208) + #define corr_done_addr (mbar_base + 224) + #define stats_addr (mbar_base + 240) + #define sum_ready_addr (mbar_base + 256) + #define o_done_addr (mbar_base + 264) + #define pv_done_addr (mbar_base + 272) + #define s_seeded_addr (mbar_base + 280) + #define q_pair_ready_addr (mbar_base + 288) + #define kv_pair_ready_addr (mbar_base + 296) + #define pv_pair_ready_addr (mbar_base + 360) + #define tmem_dealloc_addr (mbar_base + 376) + #define tmem_dealloc_peer_addr (mbar_base + 384) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + + // Kernel setup ops + uint8_t* smem_q = reinterpret_cast(smem_raw + 1024); + const int smem_q_addr = smem + 1024; + uint8_t* smem_kv = reinterpret_cast(smem_raw + 66560); + const int smem_kv_addr = smem + 66560; + uint8_t* smem_v = reinterpret_cast(smem_raw + 66560); + const int smem_v_addr = smem + 66560; + uint8_t* smem_q_full = reinterpret_cast(smem_raw + 1024); + const int smem_q_full_addr = smem + 1024; + uint8_t* smem_k_full = reinterpret_cast(smem_raw + 66560); + const int smem_k_full_addr = smem + 66560; + uint8_t* smem_v_full = reinterpret_cast(smem_raw + 132096); + const int smem_v_full_addr = smem + 132096; + float* smem_stats_max = reinterpret_cast(smem_raw + 197632); + const int smem_stats_max_addr = smem + 197632; + float* smem_stats_sum = reinterpret_cast(smem_raw + 198656); + const int smem_stats_sum_addr = smem + 198656; + float* smem_stats_final_max = reinterpret_cast(smem_raw + 199168); + const int smem_stats_final_max_addr = smem + 199168; + float* smem_softmax_warp_pair_exchange = reinterpret_cast(smem_raw + 199680); + const int smem_softmax_warp_pair_exchange_addr = smem + 199680; + float* smem_corr_warp_pair_exchange = reinterpret_cast(smem_raw + 200704); + const int smem_corr_warp_pair_exchange_addr = smem + 200704; + int* smem_sparse_indices = reinterpret_cast(smem_raw + 201216); + const int smem_sparse_indices_addr = smem + 201216; + uint8_t* smem_p_fp8 = reinterpret_cast(smem_raw + 1024); + const int smem_p_fp8_addr = smem + 1024; + + // Mbarrier init (23 pipeline groups, 0 ordered-sequence groups, 49 barriers) + // Mbarriers at smem_raw[0..392) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // q_full: 1 barriers, init_count=1 + mbarrier_init(smem + 0, 1); + // q_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 8, 1); + // --- pipeline 'k_pipe' --- + // k_full: 1 barriers, init_count=1 + mbarrier_init(smem + 16, 1); + // k_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 24, 1); + // --- pipeline 'v_pipe' --- + // v_full: 1 barriers, init_count=1 + mbarrier_init(smem + 32, 1); + // v_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 40, 1); + // --- pipeline 'kv_pipe' --- + // kv_full: 8 barriers, init_count=1 + mbarrier_init(smem + 48, 1); + mbarrier_init(smem + 56, 1); + mbarrier_init(smem + 64, 1); + mbarrier_init(smem + 72, 1); + mbarrier_init(smem + 80, 1); + mbarrier_init(smem + 88, 1); + mbarrier_init(smem + 96, 1); + mbarrier_init(smem + 104, 1); + // kv_empty: 8 barriers, init_count=1 + mbarrier_init(smem + 112, 1); + mbarrier_init(smem + 120, 1); + mbarrier_init(smem + 128, 1); + mbarrier_init(smem + 136, 1); + mbarrier_init(smem + 144, 1); + mbarrier_init(smem + 152, 1); + mbarrier_init(smem + 160, 1); + mbarrier_init(smem + 168, 1); + // --- pipeline 'index_pipe' --- + // index_full: 1 barriers, init_count=1 + mbarrier_init(smem + 176, 1); + // index_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 184, 1); + // s_full: 2 barriers, init_count=1 + mbarrier_init(smem + 192, 1); + mbarrier_init(smem + 200, 1); + // p_full: 2 barriers, init_count=128 + mbarrier_init(smem + 208, 128); + mbarrier_init(smem + 216, 128); + // corr_done: 2 barriers, init_count=128 + mbarrier_init(smem + 224, 128); + mbarrier_init(smem + 232, 128); + // stats: 2 barriers, init_count=128 + mbarrier_init(smem + 240, 128); + mbarrier_init(smem + 248, 128); + // sum_ready: 1 barriers, init_count=128 + mbarrier_init(smem + 256, 128); + // o_done: 1 barriers, init_count=1 + mbarrier_init(smem + 264, 1); + // pv_done: 1 barriers, init_count=1 + mbarrier_init(smem + 272, 1); + // s_seeded: 1 barriers, init_count=128 + mbarrier_init(smem + 280, 128); + // q_pair_ready: 1 barriers, init_count=64 + mbarrier_init(smem + 288, 64); + // --- pipeline 'kv_pipe' --- + // kv_pair_ready: 8 barriers, init_count=64 + mbarrier_init(smem + 296, 64); + mbarrier_init(smem + 304, 64); + mbarrier_init(smem + 312, 64); + mbarrier_init(smem + 320, 64); + mbarrier_init(smem + 328, 64); + mbarrier_init(smem + 336, 64); + mbarrier_init(smem + 344, 64); + mbarrier_init(smem + 352, 64); + // pv_pair_ready: 2 barriers, init_count=64 + mbarrier_init(smem + 360, 64); + mbarrier_init(smem + 368, 64); + // tmem_dealloc: 1 barriers, init_count=416 + mbarrier_init(smem + 376, 416); + // tmem_dealloc_peer: 1 barriers, init_count=32 + mbarrier_init(smem + 384, 32); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncwarp(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 392); + if (warp == 0) { + int _tmem_hold = smem + 392; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_tmem_scratch = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 15) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 64;"); + } + + // ---- Role: index_warp ---- + if (warp == 13) { + { // index_warp_main + const int index_dummy = 0; + unsigned int _phase_index_empty = 1; + unsigned int _phase_q_full_0 = 0; + { + #pragma unroll 1 + for (unsigned int work_idx = bid; work_idx < total_work_items; work_idx += num_bids) { + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + } + } + } + } + // ---- Role: load_warp ---- + if (warp >= 9 && warp <= 12) { + { // load_warp_main + const int wg2_dummy = 0; + const int load_warp_rank = warp - ((0) ? 12 : 9); + int all_num_kv_tiles = (sparse_topk + 128 - 1) / 128; + unsigned int load_k_stage = 0; + unsigned int load_v_stage = 0; + unsigned int load_kv_stage = 0; + unsigned int load_k_index_stage = 0; + unsigned int load_v_index_stage = 0; + int k_cta_offset = 0; + unsigned int _phase_q_empty_0 = 1; + unsigned int _phase_index_full = 0; + unsigned int _phase_k_empty = 1; + unsigned int _phase_kv_empty = 1; + unsigned int _phase_v_empty = 1; + #pragma unroll 1 + for (unsigned int work_idx_1 = bid; work_idx_1 < total_work_items; work_idx_1 += num_bids) { + int split_idx = 0; + int query_idx = work_idx_1 >> 1; + int v_chunk = work_idx_1 & 1; + int tiles_per_split = all_num_kv_tiles; + int first_tile = split_idx * tiles_per_split; + int num_kv_tiles = tiles_per_split; + int sparse_extent = num_kv_tiles * 128; + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + if (load_warp_rank == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(q_full_addr, 65536); + #pragma unroll + for (int q_stage = 0; q_stage < 4; q_stage++) { + tma_4d_gmem2smem(smem_q_addr + (unsigned int)(q_stage * 16384), (&tmap_q), 0, 0, q_stage, query_idx, q_full_addr); + } + } + } + { + int num_index_passes = sparse_extent / 128; + #pragma unroll 1 + for (int index_pass = load_warp_rank; index_pass < num_index_passes; index_pass += 4) { + int index_offset = index_pass * 128 + lane * 4; + int tile_col_base = (first_tile + index_pass) * 128; + int* row_tile_ptr = ((tile_col_base < 128) ? swa_indices + (query_idx * swa_index_stride) : compressed_indices + (query_idx * compressed_index_stride + (tile_col_base - 128))); + int lane_col = tile_col_base + lane * 4; + int sparse_rows[4]; + #pragma unroll + for (int row_i = 0; row_i < 4; row_i++) { + sparse_rows[row_i] = -1; + } + if (lane_col + 4 <= sparse_topk) { + int _vec_load_0[4]; + { + const int4* _ivptr_0 = reinterpret_cast(row_tile_ptr + (lane * 4) + 0); + int4 _ivld_0; + _ivld_0 = *_ivptr_0; + _vec_load_0[0 + 0] = _ivld_0.x; + _vec_load_0[0 + 1] = _ivld_0.y; + _vec_load_0[0 + 2] = _ivld_0.z; + _vec_load_0[0 + 3] = _ivld_0.w; + } + #pragma unroll + for (int row_i_1 = 0; row_i_1 < 4; row_i_1++) { + sparse_rows[row_i_1] = _vec_load_0[row_i_1]; + } + } else { + #pragma unroll + for (int row_i_2 = 0; row_i_2 < 4; row_i_2++) { + if (lane_col + row_i_2 < sparse_topk) { + sparse_rows[row_i_2] = row_tile_ptr[lane * 4 + row_i_2]; + } + } + } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(smem_sparse_indices_addr + (unsigned int)(index_offset * 4)), "r"(sparse_rows[0]), "r"(sparse_rows[1]), "r"(sparse_rows[2]), "r"(sparse_rows[3]) : "memory"); + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("barrier.sync 8, 256;" ::: "memory"); + } + int k_index_stage_base = smem_sparse_indices_addr; + { + #pragma unroll + for (int qk_stage = 0; qk_stage < 4; qk_stage++) { + if (qk_stage == load_warp_rank) { + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (load_kv_stage) * 8, 16384); + } + int k_dst = smem_kv_addr + load_kv_stage * 16384; + int group = lane; + int k_group_limit = ((work_idx_1 == (unsigned int)bid) ? 16 : 32); + if (group < k_group_limit) { + int group_offset = k_cta_offset + group * 4; + int raw_rows[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows[0])), "=r"(*reinterpret_cast(&raw_rows[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows[(0) + 3])) + : "r"(k_index_stage_base + group_offset * 4)); + int row0 = ((raw_rows[0] >= 0) ? raw_rows[0] : 0); + int row1 = ((raw_rows[1] >= 0) ? raw_rows[1] : 0); + int row2 = ((raw_rows[2] >= 0) ? raw_rows[2] : 0); + int row3 = ((raw_rows[3] >= 0) ? raw_rows[3] : 0); + if (first_tile == 0) { + tma_gather4_gmem2smem(k_dst + group * 512, (&tmap_swa_kv), qk_stage * 128, row0, row1, row2, row3, kv_full_addr + (load_kv_stage) * 8); + } else { + tma_gather4_gmem2smem(k_dst + group * 512, (&tmap_compressed_kv), qk_stage * 128, row0, row1, row2, row3, kv_full_addr + (load_kv_stage) * 8); + } + } + } + load_kv_stage += 1; + if (load_kv_stage == 8) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + } + #pragma unroll 1 + for (int tile = 1; tile < num_kv_tiles; tile++) { + if ((0 & (int)((tile & 1) == 0)) != 0) { + mbarrier_wait(index_full_addr + (load_k_index_stage) * 8, _phase_index_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + } + k_index_stage_base = smem_sparse_indices_addr; + { + #pragma unroll + for (int qk_stage_1 = 0; qk_stage_1 < 4; qk_stage_1++) { + if (qk_stage_1 == load_warp_rank) { + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (load_kv_stage) * 8, 16384); + } + int k_dst_1 = smem_kv_addr + load_kv_stage * 16384; + int group_1 = lane; + if (group_1 < 32) { + int group_offset_1 = tile * 128 + k_cta_offset + group_1 * 4; + int raw_rows_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_1[0])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 3])) + : "r"(k_index_stage_base + group_offset_1 * 4)); + int row0_1 = ((raw_rows_1[0] >= 0) ? raw_rows_1[0] : 0); + int row1_1 = ((raw_rows_1[1] >= 0) ? raw_rows_1[1] : 0); + int row2_1 = ((raw_rows_1[2] >= 0) ? raw_rows_1[2] : 0); + int row3_1 = ((raw_rows_1[3] >= 0) ? raw_rows_1[3] : 0); + tma_gather4_gmem2smem(k_dst_1 + group_1 * 512, (&tmap_compressed_kv), qk_stage_1 * 128, row0_1, row1_1, row2_1, row3_1, kv_full_addr + (load_kv_stage) * 8); + } + } + load_kv_stage += 1; + if (load_kv_stage == 8) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + } + int prev_tile = tile - 1; + int global_prev_tile = first_tile + prev_tile; + int v_index_stage_base = smem_sparse_indices_addr; + { + #pragma unroll + for (int pv_stage = 0; pv_stage < 2; pv_stage++) { + if (pv_stage == load_warp_rank) { + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (load_kv_stage) * 8, 16384); + } + int v_dst = smem_kv_addr + load_kv_stage * 16384; + int v_col = v_chunk * 256 + pv_stage * 128; + int group_2 = lane; + int group_offset_2 = prev_tile * 128 + group_2 * 4; + int raw_rows_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_2[0])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 3])) + : "r"(v_index_stage_base + group_offset_2 * 4)); + int row0_2 = ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : 0); + int row1_2 = ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : 0); + int row2_2 = ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : 0); + int row3_2 = ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : 0); + if (global_prev_tile == 0) { + tma_gather4_gmem2smem(v_dst + group_2 * 512, (&tmap_swa_kv), v_col, row0_2, row1_2, row2_2, row3_2, kv_full_addr + (load_kv_stage) * 8); + } else { + tma_gather4_gmem2smem(v_dst + group_2 * 512, (&tmap_compressed_kv), v_col, row0_2, row1_2, row2_2, row3_2, kv_full_addr + (load_kv_stage) * 8); + } + } + load_kv_stage += 1; + if (load_kv_stage == 8) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + } + } + int last_tile = num_kv_tiles - 1; + int global_last_tile = first_tile + last_tile; + { + #pragma unroll + for (int pv_stage_1 = 0; pv_stage_1 < 2; pv_stage_1++) { + if (pv_stage_1 == load_warp_rank) { + mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (load_kv_stage) * 8, 16384); + } + int v_dst_1 = smem_kv_addr + load_kv_stage * 16384; + int v_col_1 = v_chunk * 256 + pv_stage_1 * 128; + int group_3 = lane; + int group_offset_3 = last_tile * 128 + group_3 * 4; + int raw_rows_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_3[0])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 3])) + : "r"(smem_sparse_indices_addr + (unsigned int)(group_offset_3 * 4))); + int row0_3 = ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : 0); + int row1_3 = ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : 0); + int row2_3 = ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : 0); + int row3_3 = ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : 0); + if (global_last_tile == 0) { + tma_gather4_gmem2smem(v_dst_1 + group_3 * 512, (&tmap_swa_kv), v_col_1, row0_3, row1_3, row2_3, row3_3, kv_full_addr + (load_kv_stage) * 8); + } else { + tma_gather4_gmem2smem(v_dst_1 + group_3 * 512, (&tmap_compressed_kv), v_col_1, row0_3, row1_3, row2_3, row3_3, kv_full_addr + (load_kv_stage) * 8); + } + } + load_kv_stage += 1; + if (load_kv_stage == 8) { load_kv_stage = 0; _phase_kv_empty ^= 1; } + } + } + } + mbarrier_arrive(tmem_dealloc_addr); + } + } + // ---- Role: softmax_wg ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 184;"); + { // softmax_wg_main + float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; + const int wg_dummy_inc = 0; + int all_num_kv_tiles_1 = (sparse_topk + 128 - 1) / 128; + const int tmem_row_base = ((0) ? warp % 2 * 32 : warp % 4 * 32); + const int tmem_score_row_base = ((0) ? (int)(warp % 4 * 32) : tmem_row_base); + const int n_half = ((0) ? (int)(warp % 4 / 2) : 0); + const int my_row = tmem_row_base + lane; + const int stats_row = n_half * 64 + my_row; + int softmax_tile_cursor = 0; + unsigned int softmax_index_stage = 0; + unsigned int _phase_q_full_0_1 = 0; + #pragma unroll 1 + for (unsigned int work_idx_2 = bid; work_idx_2 < total_work_items; work_idx_2 += num_bids) { + int split_idx_1 = 0; + int query_idx_1 = work_idx_2 >> 1; + int tiles_per_split_1 = all_num_kv_tiles_1; + int first_tile_1 = split_idx_1 * tiles_per_split_1; + int num_kv_tiles_1 = tiles_per_split_1; + int _max_0 = ((sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) : (0)); + int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); + int active_topk = _min_0; + { + float seed_zero[4]; + #pragma unroll + for (int seed_c4 = 0; seed_c4 < 4; seed_c4++) { + seed_zero[seed_c4] = 0.0f; + } + #pragma unroll + for (int seed_half = 0; seed_half < 2; seed_half++) { + #pragma unroll + for (int seed_c = 0; seed_c < 8; seed_c++) { + int seed_addr = taddr + (unsigned int)(((seed_half == 0) ? 96 : 224)) + (unsigned int)(seed_c * 4) + (unsigned int)(tmem_row_base << 16); + tmem_st_x4_f32(seed_addr, seed_zero); + } + } + { + #pragma unroll + for (int seed_half_1 = 0; seed_half_1 < 2; seed_half_1++) { + #pragma unroll + for (int seed_c_1 = 16; seed_c_1 < 24; seed_c_1++) { + int seed_addr_1 = taddr + (unsigned int)(((seed_half_1 == 0) ? 0 : 128)) + (unsigned int)(seed_c_1 * 4) + (unsigned int)(tmem_row_base << 16); + tmem_st_x4_f32(seed_addr_1, seed_zero); + } + } + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(s_seeded_addr); + } + { + asm volatile("barrier.sync 8, 256;" ::: "memory"); + if (work_idx_2 == (unsigned int)bid) { + const int helper_stage = warp % 4; + int helper_dst = smem_kv_addr + (unsigned int)(helper_stage * 16384); + int helper_group = lane; + if (helper_group >= 16) { + int helper_rows[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&helper_rows[0])), "=r"(*reinterpret_cast(&helper_rows[(0) + 1])), "=r"(*reinterpret_cast(&helper_rows[(0) + 2])), "=r"(*reinterpret_cast(&helper_rows[(0) + 3])) + : "r"(smem_sparse_indices_addr + (unsigned int)(helper_group * 16))); + int hrow0 = ((helper_rows[0] >= 0) ? helper_rows[0] : 0); + int hrow1 = ((helper_rows[1] >= 0) ? helper_rows[1] : 0); + int hrow2 = ((helper_rows[2] >= 0) ? helper_rows[2] : 0); + int hrow3 = ((helper_rows[3] >= 0) ? helper_rows[3] : 0); + if (first_tile_1 == 0) { + tma_gather4_gmem2smem(helper_dst + helper_group * 512, (&tmap_swa_kv), helper_stage * 128, hrow0, hrow1, hrow2, hrow3, kv_full_addr + (helper_stage) * 8); + } else { + tma_gather4_gmem2smem(helper_dst + helper_group * 512, (&tmap_compressed_kv), helper_stage * 128, hrow0, hrow1, hrow2, hrow3, kv_full_addr + (helper_stage) * 8); + } + } + } + } + { + mbarrier_wait(q_full_addr, _phase_q_full_0_1); + _phase_q_full_0_1 ^= 1; + } + float row_max_val = -CAKE_INF; + float row_sum_val = 0.0f; + int sink_head = ((0) ? my_row : my_row); + if (has_sinks != 0 && sink_head < num_heads && split_idx_1 == 0) { + row_max_val = sinks[sink_head] * 1.4426950408889634f / softmax_scale_log2; + row_sum_val = 1.0f; + } + #pragma unroll 1 + for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { + int pipeline_tile = softmax_tile_cursor + tile_1; + int phase = pipeline_tile & 1; + int s_wait_phase = pipeline_tile >> 1 & 1; + mbarrier_wait(s_full_addr + (phase) * 8, s_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int s_off = ((phase != 0) ? 128 : 0); + int s_base = taddr + (unsigned int)s_off + (unsigned int)(tmem_score_row_base << 16); + float new_max = row_max_val; + int valid_sparse_cols = ((active_topk < sparse_topk) ? active_topk : sparse_topk); + valid_sparse_cols = valid_sparse_cols - (first_tile_1 + tile_1) * 128 - n_half * 64; + if (valid_sparse_cols < 0) { + valid_sparse_cols = 0; + } + if (valid_sparse_cols > 128) { + valid_sparse_cols = 128; + } + int index_tile_addr = smem_sparse_indices_addr + (unsigned int)(tile_1 * 512); + int index_lane_addr = index_tile_addr + (n_half * 64 + lane) * 4; + int staged_index[4]; + unsigned int invalid_cols[4]; + unsigned int any_invalid = 0; + #pragma unroll + for (int w = 0; w < 4; w++) { + asm volatile("ld.shared.b32 %0, [%1];" : "=r"(*reinterpret_cast(&staged_index[w])) : "r"(index_lane_addr + w * 128)); + unsigned int _vote_0 = __ballot_sync(0xFFFFFFFF, staged_index[w] < 0); + invalid_cols[w] = _vote_0; + any_invalid = any_invalid | invalid_cols[w]; + } + float _tmem_load_0[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]) + : "r"(s_base)); + float _tmem_load_1[4]; + tmem_ld_x4(&_tmem_load_1[0], s_base); + float _tmem_load_2[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]) + : "r"(s_base + 32)); + float _tmem_load_3[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]) + : "r"(s_base + 64)); + float _tmem_load_4[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_4[0]), "=f"(_tmem_load_4[1]), "=f"(_tmem_load_4[2]), "=f"(_tmem_load_4[3]), "=f"(_tmem_load_4[4]), "=f"(_tmem_load_4[5]), "=f"(_tmem_load_4[6]), "=f"(_tmem_load_4[7]), "=f"(_tmem_load_4[8]), "=f"(_tmem_load_4[9]), "=f"(_tmem_load_4[10]), "=f"(_tmem_load_4[11]), "=f"(_tmem_load_4[12]), "=f"(_tmem_load_4[13]), "=f"(_tmem_load_4[14]), "=f"(_tmem_load_4[15]), "=f"(_tmem_load_4[16]), "=f"(_tmem_load_4[17]), "=f"(_tmem_load_4[18]), "=f"(_tmem_load_4[19]), "=f"(_tmem_load_4[20]), "=f"(_tmem_load_4[21]), "=f"(_tmem_load_4[22]), "=f"(_tmem_load_4[23]), "=f"(_tmem_load_4[24]), "=f"(_tmem_load_4[25]), "=f"(_tmem_load_4[26]), "=f"(_tmem_load_4[27]), "=f"(_tmem_load_4[28]), "=f"(_tmem_load_4[29]), "=f"(_tmem_load_4[30]), "=f"(_tmem_load_4[31]) + : "r"(s_base + 96)); + { + int frag_valid_0 = valid_sparse_cols; + uint32_t _slice_lo_mask_1; + { + int _lim_0 = frag_valid_0; + if (_lim_0 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) _tmem_load_0[0] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) _tmem_load_0[1] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) _tmem_load_0[2] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) _tmem_load_0[3] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) _tmem_load_0[4] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) _tmem_load_0[5] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) _tmem_load_0[6] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) _tmem_load_0[7] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) _tmem_load_0[8] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) _tmem_load_0[9] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) _tmem_load_0[10] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) _tmem_load_0[11] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) _tmem_load_0[12] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) _tmem_load_0[13] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) _tmem_load_0[14] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) _tmem_load_0[15] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) _tmem_load_0[16] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) _tmem_load_0[17] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) _tmem_load_0[18] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) _tmem_load_0[19] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) _tmem_load_0[20] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) _tmem_load_0[21] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) _tmem_load_0[22] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) _tmem_load_0[23] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) _tmem_load_0[24] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) _tmem_load_0[25] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) _tmem_load_0[26] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) _tmem_load_0[27] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) _tmem_load_0[28] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) _tmem_load_0[29] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) _tmem_load_0[30] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) _tmem_load_0[31] = -CAKE_INF; + if (any_invalid != 0) { + unsigned int frag_mask_word_0 = invalid_cols[0]; + _tmem_load_0[0] = (((frag_mask_word_0 & 1) != 0) ? -CAKE_INF : _tmem_load_0[0]); + _tmem_load_0[1] = (((frag_mask_word_0 >> 1 & 1) != 0) ? -CAKE_INF : _tmem_load_0[1]); + _tmem_load_0[2] = (((frag_mask_word_0 >> 2 & 1) != 0) ? -CAKE_INF : _tmem_load_0[2]); + _tmem_load_0[3] = (((frag_mask_word_0 >> 3 & 1) != 0) ? -CAKE_INF : _tmem_load_0[3]); + _tmem_load_0[4] = (((frag_mask_word_0 >> 4 & 1) != 0) ? -CAKE_INF : _tmem_load_0[4]); + _tmem_load_0[5] = (((frag_mask_word_0 >> 5 & 1) != 0) ? -CAKE_INF : _tmem_load_0[5]); + _tmem_load_0[6] = (((frag_mask_word_0 >> 6 & 1) != 0) ? -CAKE_INF : _tmem_load_0[6]); + _tmem_load_0[7] = (((frag_mask_word_0 >> 7 & 1) != 0) ? -CAKE_INF : _tmem_load_0[7]); + _tmem_load_0[8] = (((frag_mask_word_0 >> 8 & 1) != 0) ? -CAKE_INF : _tmem_load_0[8]); + _tmem_load_0[9] = (((frag_mask_word_0 >> 9 & 1) != 0) ? -CAKE_INF : _tmem_load_0[9]); + _tmem_load_0[10] = (((frag_mask_word_0 >> 10 & 1) != 0) ? -CAKE_INF : _tmem_load_0[10]); + _tmem_load_0[11] = (((frag_mask_word_0 >> 11 & 1) != 0) ? -CAKE_INF : _tmem_load_0[11]); + _tmem_load_0[12] = (((frag_mask_word_0 >> 12 & 1) != 0) ? -CAKE_INF : _tmem_load_0[12]); + _tmem_load_0[13] = (((frag_mask_word_0 >> 13 & 1) != 0) ? -CAKE_INF : _tmem_load_0[13]); + _tmem_load_0[14] = (((frag_mask_word_0 >> 14 & 1) != 0) ? -CAKE_INF : _tmem_load_0[14]); + _tmem_load_0[15] = (((frag_mask_word_0 >> 15 & 1) != 0) ? -CAKE_INF : _tmem_load_0[15]); + _tmem_load_0[16] = (((frag_mask_word_0 >> 16 & 1) != 0) ? -CAKE_INF : _tmem_load_0[16]); + _tmem_load_0[17] = (((frag_mask_word_0 >> 17 & 1) != 0) ? -CAKE_INF : _tmem_load_0[17]); + _tmem_load_0[18] = (((frag_mask_word_0 >> 18 & 1) != 0) ? -CAKE_INF : _tmem_load_0[18]); + _tmem_load_0[19] = (((frag_mask_word_0 >> 19 & 1) != 0) ? -CAKE_INF : _tmem_load_0[19]); + _tmem_load_0[20] = (((frag_mask_word_0 >> 20 & 1) != 0) ? -CAKE_INF : _tmem_load_0[20]); + _tmem_load_0[21] = (((frag_mask_word_0 >> 21 & 1) != 0) ? -CAKE_INF : _tmem_load_0[21]); + _tmem_load_0[22] = (((frag_mask_word_0 >> 22 & 1) != 0) ? -CAKE_INF : _tmem_load_0[22]); + _tmem_load_0[23] = (((frag_mask_word_0 >> 23 & 1) != 0) ? -CAKE_INF : _tmem_load_0[23]); + _tmem_load_0[24] = (((frag_mask_word_0 >> 24 & 1) != 0) ? -CAKE_INF : _tmem_load_0[24]); + _tmem_load_0[25] = (((frag_mask_word_0 >> 25 & 1) != 0) ? -CAKE_INF : _tmem_load_0[25]); + _tmem_load_0[26] = (((frag_mask_word_0 >> 26 & 1) != 0) ? -CAKE_INF : _tmem_load_0[26]); + _tmem_load_0[27] = (((frag_mask_word_0 >> 27 & 1) != 0) ? -CAKE_INF : _tmem_load_0[27]); + _tmem_load_0[28] = (((frag_mask_word_0 >> 28 & 1) != 0) ? -CAKE_INF : _tmem_load_0[28]); + _tmem_load_0[29] = (((frag_mask_word_0 >> 29 & 1) != 0) ? -CAKE_INF : _tmem_load_0[29]); + _tmem_load_0[30] = (((frag_mask_word_0 >> 30 & 1) != 0) ? -CAKE_INF : _tmem_load_0[30]); + _tmem_load_0[31] = (((frag_mask_word_0 >> 31 & 1) != 0) ? -CAKE_INF : _tmem_load_0[31]); + } + float2 _reg_reduce_max2_1 = {-CAKE_INF, -CAKE_INF}; + _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_0[0], _tmem_load_0[1])); + _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_0[2], _tmem_load_0[3])); + _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_0[4], _tmem_load_0[5])); + _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_0[6], _tmem_load_0[7])); + _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_0[8], _tmem_load_0[9])); + _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_0[10], _tmem_load_0[11])); + _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_0[12], _tmem_load_0[13])); + _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_0[14], _tmem_load_0[15])); + _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_0[16], _tmem_load_0[17])); + _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_0[18], _tmem_load_0[19])); + _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_0[20], _tmem_load_0[21])); + _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_0[22], _tmem_load_0[23])); + _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_0[24], _tmem_load_0[25])); + _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_0[26], _tmem_load_0[27])); + _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_0[28], _tmem_load_0[29])); + _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_0[30], _tmem_load_0[31])); + float _tmem_load_0_max = row_max_reduce(_reg_reduce_max2_1); + float _max_3 = max_noftz(new_max, _tmem_load_0_max); + new_max = _max_3; + int frag_valid_1 = valid_sparse_cols - 32; + uint32_t _slice_lo_mask_2; + { + int _lim_2 = frag_valid_1; + if (_lim_2 <= 0) { _slice_lo_mask_2 = 0u; } + else if (_lim_2 >= 32) { _slice_lo_mask_2 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_2) : "r"(_lim_2)); + } + } + if (!(_slice_lo_mask_2 & (1u << 0))) _tmem_load_2[0] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 1))) _tmem_load_2[1] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 2))) _tmem_load_2[2] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 3))) _tmem_load_2[3] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 4))) _tmem_load_2[4] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 5))) _tmem_load_2[5] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 6))) _tmem_load_2[6] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 7))) _tmem_load_2[7] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 8))) _tmem_load_2[8] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 9))) _tmem_load_2[9] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 10))) _tmem_load_2[10] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 11))) _tmem_load_2[11] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 12))) _tmem_load_2[12] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 13))) _tmem_load_2[13] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 14))) _tmem_load_2[14] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 15))) _tmem_load_2[15] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 16))) _tmem_load_2[16] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 17))) _tmem_load_2[17] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 18))) _tmem_load_2[18] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 19))) _tmem_load_2[19] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 20))) _tmem_load_2[20] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 21))) _tmem_load_2[21] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 22))) _tmem_load_2[22] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 23))) _tmem_load_2[23] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 24))) _tmem_load_2[24] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 25))) _tmem_load_2[25] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 26))) _tmem_load_2[26] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 27))) _tmem_load_2[27] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 28))) _tmem_load_2[28] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 29))) _tmem_load_2[29] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 30))) _tmem_load_2[30] = -CAKE_INF; + if (!(_slice_lo_mask_2 & (1u << 31))) _tmem_load_2[31] = -CAKE_INF; + if (any_invalid != 0) { + unsigned int frag_mask_word_1 = invalid_cols[1]; + _tmem_load_2[0] = (((frag_mask_word_1 & 1) != 0) ? -CAKE_INF : _tmem_load_2[0]); + _tmem_load_2[1] = (((frag_mask_word_1 >> 1 & 1) != 0) ? -CAKE_INF : _tmem_load_2[1]); + _tmem_load_2[2] = (((frag_mask_word_1 >> 2 & 1) != 0) ? -CAKE_INF : _tmem_load_2[2]); + _tmem_load_2[3] = (((frag_mask_word_1 >> 3 & 1) != 0) ? -CAKE_INF : _tmem_load_2[3]); + _tmem_load_2[4] = (((frag_mask_word_1 >> 4 & 1) != 0) ? -CAKE_INF : _tmem_load_2[4]); + _tmem_load_2[5] = (((frag_mask_word_1 >> 5 & 1) != 0) ? -CAKE_INF : _tmem_load_2[5]); + _tmem_load_2[6] = (((frag_mask_word_1 >> 6 & 1) != 0) ? -CAKE_INF : _tmem_load_2[6]); + _tmem_load_2[7] = (((frag_mask_word_1 >> 7 & 1) != 0) ? -CAKE_INF : _tmem_load_2[7]); + _tmem_load_2[8] = (((frag_mask_word_1 >> 8 & 1) != 0) ? -CAKE_INF : _tmem_load_2[8]); + _tmem_load_2[9] = (((frag_mask_word_1 >> 9 & 1) != 0) ? -CAKE_INF : _tmem_load_2[9]); + _tmem_load_2[10] = (((frag_mask_word_1 >> 10 & 1) != 0) ? -CAKE_INF : _tmem_load_2[10]); + _tmem_load_2[11] = (((frag_mask_word_1 >> 11 & 1) != 0) ? -CAKE_INF : _tmem_load_2[11]); + _tmem_load_2[12] = (((frag_mask_word_1 >> 12 & 1) != 0) ? -CAKE_INF : _tmem_load_2[12]); + _tmem_load_2[13] = (((frag_mask_word_1 >> 13 & 1) != 0) ? -CAKE_INF : _tmem_load_2[13]); + _tmem_load_2[14] = (((frag_mask_word_1 >> 14 & 1) != 0) ? -CAKE_INF : _tmem_load_2[14]); + _tmem_load_2[15] = (((frag_mask_word_1 >> 15 & 1) != 0) ? -CAKE_INF : _tmem_load_2[15]); + _tmem_load_2[16] = (((frag_mask_word_1 >> 16 & 1) != 0) ? -CAKE_INF : _tmem_load_2[16]); + _tmem_load_2[17] = (((frag_mask_word_1 >> 17 & 1) != 0) ? -CAKE_INF : _tmem_load_2[17]); + _tmem_load_2[18] = (((frag_mask_word_1 >> 18 & 1) != 0) ? -CAKE_INF : _tmem_load_2[18]); + _tmem_load_2[19] = (((frag_mask_word_1 >> 19 & 1) != 0) ? -CAKE_INF : _tmem_load_2[19]); + _tmem_load_2[20] = (((frag_mask_word_1 >> 20 & 1) != 0) ? -CAKE_INF : _tmem_load_2[20]); + _tmem_load_2[21] = (((frag_mask_word_1 >> 21 & 1) != 0) ? -CAKE_INF : _tmem_load_2[21]); + _tmem_load_2[22] = (((frag_mask_word_1 >> 22 & 1) != 0) ? -CAKE_INF : _tmem_load_2[22]); + _tmem_load_2[23] = (((frag_mask_word_1 >> 23 & 1) != 0) ? -CAKE_INF : _tmem_load_2[23]); + _tmem_load_2[24] = (((frag_mask_word_1 >> 24 & 1) != 0) ? -CAKE_INF : _tmem_load_2[24]); + _tmem_load_2[25] = (((frag_mask_word_1 >> 25 & 1) != 0) ? -CAKE_INF : _tmem_load_2[25]); + _tmem_load_2[26] = (((frag_mask_word_1 >> 26 & 1) != 0) ? -CAKE_INF : _tmem_load_2[26]); + _tmem_load_2[27] = (((frag_mask_word_1 >> 27 & 1) != 0) ? -CAKE_INF : _tmem_load_2[27]); + _tmem_load_2[28] = (((frag_mask_word_1 >> 28 & 1) != 0) ? -CAKE_INF : _tmem_load_2[28]); + _tmem_load_2[29] = (((frag_mask_word_1 >> 29 & 1) != 0) ? -CAKE_INF : _tmem_load_2[29]); + _tmem_load_2[30] = (((frag_mask_word_1 >> 30 & 1) != 0) ? -CAKE_INF : _tmem_load_2[30]); + _tmem_load_2[31] = (((frag_mask_word_1 >> 31 & 1) != 0) ? -CAKE_INF : _tmem_load_2[31]); + } + float2 _reg_reduce_max2_3 = {-CAKE_INF, -CAKE_INF}; + _reg_reduce_max2_3.x = max_noftz(_reg_reduce_max2_3.x, max_noftz(_tmem_load_2[0], _tmem_load_2[1])); + _reg_reduce_max2_3.y = max_noftz(_reg_reduce_max2_3.y, max_noftz(_tmem_load_2[2], _tmem_load_2[3])); + _reg_reduce_max2_3.x = max_noftz(_reg_reduce_max2_3.x, max_noftz(_tmem_load_2[4], _tmem_load_2[5])); + _reg_reduce_max2_3.y = max_noftz(_reg_reduce_max2_3.y, max_noftz(_tmem_load_2[6], _tmem_load_2[7])); + _reg_reduce_max2_3.x = max_noftz(_reg_reduce_max2_3.x, max_noftz(_tmem_load_2[8], _tmem_load_2[9])); + _reg_reduce_max2_3.y = max_noftz(_reg_reduce_max2_3.y, max_noftz(_tmem_load_2[10], _tmem_load_2[11])); + _reg_reduce_max2_3.x = max_noftz(_reg_reduce_max2_3.x, max_noftz(_tmem_load_2[12], _tmem_load_2[13])); + _reg_reduce_max2_3.y = max_noftz(_reg_reduce_max2_3.y, max_noftz(_tmem_load_2[14], _tmem_load_2[15])); + _reg_reduce_max2_3.x = max_noftz(_reg_reduce_max2_3.x, max_noftz(_tmem_load_2[16], _tmem_load_2[17])); + _reg_reduce_max2_3.y = max_noftz(_reg_reduce_max2_3.y, max_noftz(_tmem_load_2[18], _tmem_load_2[19])); + _reg_reduce_max2_3.x = max_noftz(_reg_reduce_max2_3.x, max_noftz(_tmem_load_2[20], _tmem_load_2[21])); + _reg_reduce_max2_3.y = max_noftz(_reg_reduce_max2_3.y, max_noftz(_tmem_load_2[22], _tmem_load_2[23])); + _reg_reduce_max2_3.x = max_noftz(_reg_reduce_max2_3.x, max_noftz(_tmem_load_2[24], _tmem_load_2[25])); + _reg_reduce_max2_3.y = max_noftz(_reg_reduce_max2_3.y, max_noftz(_tmem_load_2[26], _tmem_load_2[27])); + _reg_reduce_max2_3.x = max_noftz(_reg_reduce_max2_3.x, max_noftz(_tmem_load_2[28], _tmem_load_2[29])); + _reg_reduce_max2_3.y = max_noftz(_reg_reduce_max2_3.y, max_noftz(_tmem_load_2[30], _tmem_load_2[31])); + float _tmem_load_2_max = row_max_reduce(_reg_reduce_max2_3); + float _max_4 = max_noftz(new_max, _tmem_load_2_max); + new_max = _max_4; + int frag_valid_2 = valid_sparse_cols - 64; + uint32_t _slice_lo_mask_3; + { + int _lim_4 = frag_valid_2; + if (_lim_4 <= 0) { _slice_lo_mask_3 = 0u; } + else if (_lim_4 >= 32) { _slice_lo_mask_3 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_3) : "r"(_lim_4)); + } + } + if (!(_slice_lo_mask_3 & (1u << 0))) _tmem_load_3[0] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 1))) _tmem_load_3[1] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 2))) _tmem_load_3[2] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 3))) _tmem_load_3[3] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 4))) _tmem_load_3[4] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 5))) _tmem_load_3[5] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 6))) _tmem_load_3[6] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 7))) _tmem_load_3[7] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 8))) _tmem_load_3[8] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 9))) _tmem_load_3[9] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 10))) _tmem_load_3[10] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 11))) _tmem_load_3[11] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 12))) _tmem_load_3[12] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 13))) _tmem_load_3[13] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 14))) _tmem_load_3[14] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 15))) _tmem_load_3[15] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 16))) _tmem_load_3[16] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 17))) _tmem_load_3[17] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 18))) _tmem_load_3[18] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 19))) _tmem_load_3[19] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 20))) _tmem_load_3[20] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 21))) _tmem_load_3[21] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 22))) _tmem_load_3[22] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 23))) _tmem_load_3[23] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 24))) _tmem_load_3[24] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 25))) _tmem_load_3[25] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 26))) _tmem_load_3[26] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 27))) _tmem_load_3[27] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 28))) _tmem_load_3[28] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 29))) _tmem_load_3[29] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 30))) _tmem_load_3[30] = -CAKE_INF; + if (!(_slice_lo_mask_3 & (1u << 31))) _tmem_load_3[31] = -CAKE_INF; + if (any_invalid != 0) { + unsigned int frag_mask_word_2 = invalid_cols[2]; + _tmem_load_3[0] = (((frag_mask_word_2 & 1) != 0) ? -CAKE_INF : _tmem_load_3[0]); + _tmem_load_3[1] = (((frag_mask_word_2 >> 1 & 1) != 0) ? -CAKE_INF : _tmem_load_3[1]); + _tmem_load_3[2] = (((frag_mask_word_2 >> 2 & 1) != 0) ? -CAKE_INF : _tmem_load_3[2]); + _tmem_load_3[3] = (((frag_mask_word_2 >> 3 & 1) != 0) ? -CAKE_INF : _tmem_load_3[3]); + _tmem_load_3[4] = (((frag_mask_word_2 >> 4 & 1) != 0) ? -CAKE_INF : _tmem_load_3[4]); + _tmem_load_3[5] = (((frag_mask_word_2 >> 5 & 1) != 0) ? -CAKE_INF : _tmem_load_3[5]); + _tmem_load_3[6] = (((frag_mask_word_2 >> 6 & 1) != 0) ? -CAKE_INF : _tmem_load_3[6]); + _tmem_load_3[7] = (((frag_mask_word_2 >> 7 & 1) != 0) ? -CAKE_INF : _tmem_load_3[7]); + _tmem_load_3[8] = (((frag_mask_word_2 >> 8 & 1) != 0) ? -CAKE_INF : _tmem_load_3[8]); + _tmem_load_3[9] = (((frag_mask_word_2 >> 9 & 1) != 0) ? -CAKE_INF : _tmem_load_3[9]); + _tmem_load_3[10] = (((frag_mask_word_2 >> 10 & 1) != 0) ? -CAKE_INF : _tmem_load_3[10]); + _tmem_load_3[11] = (((frag_mask_word_2 >> 11 & 1) != 0) ? -CAKE_INF : _tmem_load_3[11]); + _tmem_load_3[12] = (((frag_mask_word_2 >> 12 & 1) != 0) ? -CAKE_INF : _tmem_load_3[12]); + _tmem_load_3[13] = (((frag_mask_word_2 >> 13 & 1) != 0) ? -CAKE_INF : _tmem_load_3[13]); + _tmem_load_3[14] = (((frag_mask_word_2 >> 14 & 1) != 0) ? -CAKE_INF : _tmem_load_3[14]); + _tmem_load_3[15] = (((frag_mask_word_2 >> 15 & 1) != 0) ? -CAKE_INF : _tmem_load_3[15]); + _tmem_load_3[16] = (((frag_mask_word_2 >> 16 & 1) != 0) ? -CAKE_INF : _tmem_load_3[16]); + _tmem_load_3[17] = (((frag_mask_word_2 >> 17 & 1) != 0) ? -CAKE_INF : _tmem_load_3[17]); + _tmem_load_3[18] = (((frag_mask_word_2 >> 18 & 1) != 0) ? -CAKE_INF : _tmem_load_3[18]); + _tmem_load_3[19] = (((frag_mask_word_2 >> 19 & 1) != 0) ? -CAKE_INF : _tmem_load_3[19]); + _tmem_load_3[20] = (((frag_mask_word_2 >> 20 & 1) != 0) ? -CAKE_INF : _tmem_load_3[20]); + _tmem_load_3[21] = (((frag_mask_word_2 >> 21 & 1) != 0) ? -CAKE_INF : _tmem_load_3[21]); + _tmem_load_3[22] = (((frag_mask_word_2 >> 22 & 1) != 0) ? -CAKE_INF : _tmem_load_3[22]); + _tmem_load_3[23] = (((frag_mask_word_2 >> 23 & 1) != 0) ? -CAKE_INF : _tmem_load_3[23]); + _tmem_load_3[24] = (((frag_mask_word_2 >> 24 & 1) != 0) ? -CAKE_INF : _tmem_load_3[24]); + _tmem_load_3[25] = (((frag_mask_word_2 >> 25 & 1) != 0) ? -CAKE_INF : _tmem_load_3[25]); + _tmem_load_3[26] = (((frag_mask_word_2 >> 26 & 1) != 0) ? -CAKE_INF : _tmem_load_3[26]); + _tmem_load_3[27] = (((frag_mask_word_2 >> 27 & 1) != 0) ? -CAKE_INF : _tmem_load_3[27]); + _tmem_load_3[28] = (((frag_mask_word_2 >> 28 & 1) != 0) ? -CAKE_INF : _tmem_load_3[28]); + _tmem_load_3[29] = (((frag_mask_word_2 >> 29 & 1) != 0) ? -CAKE_INF : _tmem_load_3[29]); + _tmem_load_3[30] = (((frag_mask_word_2 >> 30 & 1) != 0) ? -CAKE_INF : _tmem_load_3[30]); + _tmem_load_3[31] = (((frag_mask_word_2 >> 31 & 1) != 0) ? -CAKE_INF : _tmem_load_3[31]); + } + float2 _reg_reduce_max2_5 = {-CAKE_INF, -CAKE_INF}; + _reg_reduce_max2_5.x = max_noftz(_reg_reduce_max2_5.x, max_noftz(_tmem_load_3[0], _tmem_load_3[1])); + _reg_reduce_max2_5.y = max_noftz(_reg_reduce_max2_5.y, max_noftz(_tmem_load_3[2], _tmem_load_3[3])); + _reg_reduce_max2_5.x = max_noftz(_reg_reduce_max2_5.x, max_noftz(_tmem_load_3[4], _tmem_load_3[5])); + _reg_reduce_max2_5.y = max_noftz(_reg_reduce_max2_5.y, max_noftz(_tmem_load_3[6], _tmem_load_3[7])); + _reg_reduce_max2_5.x = max_noftz(_reg_reduce_max2_5.x, max_noftz(_tmem_load_3[8], _tmem_load_3[9])); + _reg_reduce_max2_5.y = max_noftz(_reg_reduce_max2_5.y, max_noftz(_tmem_load_3[10], _tmem_load_3[11])); + _reg_reduce_max2_5.x = max_noftz(_reg_reduce_max2_5.x, max_noftz(_tmem_load_3[12], _tmem_load_3[13])); + _reg_reduce_max2_5.y = max_noftz(_reg_reduce_max2_5.y, max_noftz(_tmem_load_3[14], _tmem_load_3[15])); + _reg_reduce_max2_5.x = max_noftz(_reg_reduce_max2_5.x, max_noftz(_tmem_load_3[16], _tmem_load_3[17])); + _reg_reduce_max2_5.y = max_noftz(_reg_reduce_max2_5.y, max_noftz(_tmem_load_3[18], _tmem_load_3[19])); + _reg_reduce_max2_5.x = max_noftz(_reg_reduce_max2_5.x, max_noftz(_tmem_load_3[20], _tmem_load_3[21])); + _reg_reduce_max2_5.y = max_noftz(_reg_reduce_max2_5.y, max_noftz(_tmem_load_3[22], _tmem_load_3[23])); + _reg_reduce_max2_5.x = max_noftz(_reg_reduce_max2_5.x, max_noftz(_tmem_load_3[24], _tmem_load_3[25])); + _reg_reduce_max2_5.y = max_noftz(_reg_reduce_max2_5.y, max_noftz(_tmem_load_3[26], _tmem_load_3[27])); + _reg_reduce_max2_5.x = max_noftz(_reg_reduce_max2_5.x, max_noftz(_tmem_load_3[28], _tmem_load_3[29])); + _reg_reduce_max2_5.y = max_noftz(_reg_reduce_max2_5.y, max_noftz(_tmem_load_3[30], _tmem_load_3[31])); + float _tmem_load_3_max = row_max_reduce(_reg_reduce_max2_5); + float _max_5 = max_noftz(new_max, _tmem_load_3_max); + new_max = _max_5; + int frag_valid_3 = valid_sparse_cols - 96; + uint32_t _slice_lo_mask_4; + { + int _lim_6 = frag_valid_3; + if (_lim_6 <= 0) { _slice_lo_mask_4 = 0u; } + else if (_lim_6 >= 32) { _slice_lo_mask_4 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_4) : "r"(_lim_6)); + } + } + if (!(_slice_lo_mask_4 & (1u << 0))) _tmem_load_4[0] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 1))) _tmem_load_4[1] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 2))) _tmem_load_4[2] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 3))) _tmem_load_4[3] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 4))) _tmem_load_4[4] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 5))) _tmem_load_4[5] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 6))) _tmem_load_4[6] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 7))) _tmem_load_4[7] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 8))) _tmem_load_4[8] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 9))) _tmem_load_4[9] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 10))) _tmem_load_4[10] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 11))) _tmem_load_4[11] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 12))) _tmem_load_4[12] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 13))) _tmem_load_4[13] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 14))) _tmem_load_4[14] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 15))) _tmem_load_4[15] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 16))) _tmem_load_4[16] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 17))) _tmem_load_4[17] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 18))) _tmem_load_4[18] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 19))) _tmem_load_4[19] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 20))) _tmem_load_4[20] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 21))) _tmem_load_4[21] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 22))) _tmem_load_4[22] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 23))) _tmem_load_4[23] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 24))) _tmem_load_4[24] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 25))) _tmem_load_4[25] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 26))) _tmem_load_4[26] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 27))) _tmem_load_4[27] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 28))) _tmem_load_4[28] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 29))) _tmem_load_4[29] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 30))) _tmem_load_4[30] = -CAKE_INF; + if (!(_slice_lo_mask_4 & (1u << 31))) _tmem_load_4[31] = -CAKE_INF; + if (any_invalid != 0) { + unsigned int frag_mask_word_3 = invalid_cols[3]; + _tmem_load_4[0] = (((frag_mask_word_3 & 1) != 0) ? -CAKE_INF : _tmem_load_4[0]); + _tmem_load_4[1] = (((frag_mask_word_3 >> 1 & 1) != 0) ? -CAKE_INF : _tmem_load_4[1]); + _tmem_load_4[2] = (((frag_mask_word_3 >> 2 & 1) != 0) ? -CAKE_INF : _tmem_load_4[2]); + _tmem_load_4[3] = (((frag_mask_word_3 >> 3 & 1) != 0) ? -CAKE_INF : _tmem_load_4[3]); + _tmem_load_4[4] = (((frag_mask_word_3 >> 4 & 1) != 0) ? -CAKE_INF : _tmem_load_4[4]); + _tmem_load_4[5] = (((frag_mask_word_3 >> 5 & 1) != 0) ? -CAKE_INF : _tmem_load_4[5]); + _tmem_load_4[6] = (((frag_mask_word_3 >> 6 & 1) != 0) ? -CAKE_INF : _tmem_load_4[6]); + _tmem_load_4[7] = (((frag_mask_word_3 >> 7 & 1) != 0) ? -CAKE_INF : _tmem_load_4[7]); + _tmem_load_4[8] = (((frag_mask_word_3 >> 8 & 1) != 0) ? -CAKE_INF : _tmem_load_4[8]); + _tmem_load_4[9] = (((frag_mask_word_3 >> 9 & 1) != 0) ? -CAKE_INF : _tmem_load_4[9]); + _tmem_load_4[10] = (((frag_mask_word_3 >> 10 & 1) != 0) ? -CAKE_INF : _tmem_load_4[10]); + _tmem_load_4[11] = (((frag_mask_word_3 >> 11 & 1) != 0) ? -CAKE_INF : _tmem_load_4[11]); + _tmem_load_4[12] = (((frag_mask_word_3 >> 12 & 1) != 0) ? -CAKE_INF : _tmem_load_4[12]); + _tmem_load_4[13] = (((frag_mask_word_3 >> 13 & 1) != 0) ? -CAKE_INF : _tmem_load_4[13]); + _tmem_load_4[14] = (((frag_mask_word_3 >> 14 & 1) != 0) ? -CAKE_INF : _tmem_load_4[14]); + _tmem_load_4[15] = (((frag_mask_word_3 >> 15 & 1) != 0) ? -CAKE_INF : _tmem_load_4[15]); + _tmem_load_4[16] = (((frag_mask_word_3 >> 16 & 1) != 0) ? -CAKE_INF : _tmem_load_4[16]); + _tmem_load_4[17] = (((frag_mask_word_3 >> 17 & 1) != 0) ? -CAKE_INF : _tmem_load_4[17]); + _tmem_load_4[18] = (((frag_mask_word_3 >> 18 & 1) != 0) ? -CAKE_INF : _tmem_load_4[18]); + _tmem_load_4[19] = (((frag_mask_word_3 >> 19 & 1) != 0) ? -CAKE_INF : _tmem_load_4[19]); + _tmem_load_4[20] = (((frag_mask_word_3 >> 20 & 1) != 0) ? -CAKE_INF : _tmem_load_4[20]); + _tmem_load_4[21] = (((frag_mask_word_3 >> 21 & 1) != 0) ? -CAKE_INF : _tmem_load_4[21]); + _tmem_load_4[22] = (((frag_mask_word_3 >> 22 & 1) != 0) ? -CAKE_INF : _tmem_load_4[22]); + _tmem_load_4[23] = (((frag_mask_word_3 >> 23 & 1) != 0) ? -CAKE_INF : _tmem_load_4[23]); + _tmem_load_4[24] = (((frag_mask_word_3 >> 24 & 1) != 0) ? -CAKE_INF : _tmem_load_4[24]); + _tmem_load_4[25] = (((frag_mask_word_3 >> 25 & 1) != 0) ? -CAKE_INF : _tmem_load_4[25]); + _tmem_load_4[26] = (((frag_mask_word_3 >> 26 & 1) != 0) ? -CAKE_INF : _tmem_load_4[26]); + _tmem_load_4[27] = (((frag_mask_word_3 >> 27 & 1) != 0) ? -CAKE_INF : _tmem_load_4[27]); + _tmem_load_4[28] = (((frag_mask_word_3 >> 28 & 1) != 0) ? -CAKE_INF : _tmem_load_4[28]); + _tmem_load_4[29] = (((frag_mask_word_3 >> 29 & 1) != 0) ? -CAKE_INF : _tmem_load_4[29]); + _tmem_load_4[30] = (((frag_mask_word_3 >> 30 & 1) != 0) ? -CAKE_INF : _tmem_load_4[30]); + _tmem_load_4[31] = (((frag_mask_word_3 >> 31 & 1) != 0) ? -CAKE_INF : _tmem_load_4[31]); + } + float2 _reg_reduce_max2_7 = {-CAKE_INF, -CAKE_INF}; + _reg_reduce_max2_7.x = max_noftz(_reg_reduce_max2_7.x, max_noftz(_tmem_load_4[0], _tmem_load_4[1])); + _reg_reduce_max2_7.y = max_noftz(_reg_reduce_max2_7.y, max_noftz(_tmem_load_4[2], _tmem_load_4[3])); + _reg_reduce_max2_7.x = max_noftz(_reg_reduce_max2_7.x, max_noftz(_tmem_load_4[4], _tmem_load_4[5])); + _reg_reduce_max2_7.y = max_noftz(_reg_reduce_max2_7.y, max_noftz(_tmem_load_4[6], _tmem_load_4[7])); + _reg_reduce_max2_7.x = max_noftz(_reg_reduce_max2_7.x, max_noftz(_tmem_load_4[8], _tmem_load_4[9])); + _reg_reduce_max2_7.y = max_noftz(_reg_reduce_max2_7.y, max_noftz(_tmem_load_4[10], _tmem_load_4[11])); + _reg_reduce_max2_7.x = max_noftz(_reg_reduce_max2_7.x, max_noftz(_tmem_load_4[12], _tmem_load_4[13])); + _reg_reduce_max2_7.y = max_noftz(_reg_reduce_max2_7.y, max_noftz(_tmem_load_4[14], _tmem_load_4[15])); + _reg_reduce_max2_7.x = max_noftz(_reg_reduce_max2_7.x, max_noftz(_tmem_load_4[16], _tmem_load_4[17])); + _reg_reduce_max2_7.y = max_noftz(_reg_reduce_max2_7.y, max_noftz(_tmem_load_4[18], _tmem_load_4[19])); + _reg_reduce_max2_7.x = max_noftz(_reg_reduce_max2_7.x, max_noftz(_tmem_load_4[20], _tmem_load_4[21])); + _reg_reduce_max2_7.y = max_noftz(_reg_reduce_max2_7.y, max_noftz(_tmem_load_4[22], _tmem_load_4[23])); + _reg_reduce_max2_7.x = max_noftz(_reg_reduce_max2_7.x, max_noftz(_tmem_load_4[24], _tmem_load_4[25])); + _reg_reduce_max2_7.y = max_noftz(_reg_reduce_max2_7.y, max_noftz(_tmem_load_4[26], _tmem_load_4[27])); + _reg_reduce_max2_7.x = max_noftz(_reg_reduce_max2_7.x, max_noftz(_tmem_load_4[28], _tmem_load_4[29])); + _reg_reduce_max2_7.y = max_noftz(_reg_reduce_max2_7.y, max_noftz(_tmem_load_4[30], _tmem_load_4[31])); + float _tmem_load_4_max = row_max_reduce(_reg_reduce_max2_7); + float _max_6 = max_noftz(new_max, _tmem_load_4_max); + new_max = _max_6; + } + float _fma_0 = __fmaf_rn(row_max_val, softmax_scale_log2, (-new_max) * softmax_scale_log2); + float delta = _fma_0; + float _exp2_0 = approx_exp2(delta); + float exp_delta = _exp2_0; + float acc_scale = ((row_max_val > -CAKE_INF) ? exp_delta : 1.0f); + smem_stats_max[phase * 128 + stats_row] = acc_scale; + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(stats_addr + (phase) * 8); + row_max_val = new_max; + float safe_max = ((new_max == -CAKE_INF) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + float block_sum = 0.0f; + { + int p_off_all = ((phase != 0) ? 224 : 96); + int p_base_all = taddr + (unsigned int)p_off_all + (unsigned int)(tmem_row_base << 16); + const float2 _fma_b2_8 = {softmax_scale_log2, softmax_scale_log2}; + const float2 _fma_c2_9 = {-max_scaled, -max_scaled}; + float2 _fma_pair_10 = fma_f32x2(make_float2(_tmem_load_0[0], _tmem_load_0[1]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[0] = _fma_pair_10.x; + _tmem_load_0[1] = _fma_pair_10.y; + float2 _fma_pair_11 = fma_f32x2(make_float2(_tmem_load_0[2], _tmem_load_0[3]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[2] = _fma_pair_11.x; + _tmem_load_0[3] = _fma_pair_11.y; + float2 _fma_pair_12 = fma_f32x2(make_float2(_tmem_load_0[4], _tmem_load_0[5]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[4] = _fma_pair_12.x; + _tmem_load_0[5] = _fma_pair_12.y; + float2 _fma_pair_13 = fma_f32x2(make_float2(_tmem_load_0[6], _tmem_load_0[7]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[6] = _fma_pair_13.x; + _tmem_load_0[7] = _fma_pair_13.y; + float2 _fma_pair_14 = fma_f32x2(make_float2(_tmem_load_0[8], _tmem_load_0[9]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[8] = _fma_pair_14.x; + _tmem_load_0[9] = _fma_pair_14.y; + float2 _fma_pair_15 = fma_f32x2(make_float2(_tmem_load_0[10], _tmem_load_0[11]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[10] = _fma_pair_15.x; + _tmem_load_0[11] = _fma_pair_15.y; + float2 _fma_pair_16 = fma_f32x2(make_float2(_tmem_load_0[12], _tmem_load_0[13]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[12] = _fma_pair_16.x; + _tmem_load_0[13] = _fma_pair_16.y; + float2 _fma_pair_17 = fma_f32x2(make_float2(_tmem_load_0[14], _tmem_load_0[15]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[14] = _fma_pair_17.x; + _tmem_load_0[15] = _fma_pair_17.y; + float2 _fma_pair_18 = fma_f32x2(make_float2(_tmem_load_0[16], _tmem_load_0[17]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[16] = _fma_pair_18.x; + _tmem_load_0[17] = _fma_pair_18.y; + float2 _fma_pair_19 = fma_f32x2(make_float2(_tmem_load_0[18], _tmem_load_0[19]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[18] = _fma_pair_19.x; + _tmem_load_0[19] = _fma_pair_19.y; + float2 _fma_pair_20 = fma_f32x2(make_float2(_tmem_load_0[20], _tmem_load_0[21]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[20] = _fma_pair_20.x; + _tmem_load_0[21] = _fma_pair_20.y; + float2 _fma_pair_21 = fma_f32x2(make_float2(_tmem_load_0[22], _tmem_load_0[23]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[22] = _fma_pair_21.x; + _tmem_load_0[23] = _fma_pair_21.y; + float2 _fma_pair_22 = fma_f32x2(make_float2(_tmem_load_0[24], _tmem_load_0[25]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[24] = _fma_pair_22.x; + _tmem_load_0[25] = _fma_pair_22.y; + float2 _fma_pair_23 = fma_f32x2(make_float2(_tmem_load_0[26], _tmem_load_0[27]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[26] = _fma_pair_23.x; + _tmem_load_0[27] = _fma_pair_23.y; + float2 _fma_pair_24 = fma_f32x2(make_float2(_tmem_load_0[28], _tmem_load_0[29]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[28] = _fma_pair_24.x; + _tmem_load_0[29] = _fma_pair_24.y; + float2 _fma_pair_25 = fma_f32x2(make_float2(_tmem_load_0[30], _tmem_load_0[31]), _fma_b2_8, _fma_c2_9); + _tmem_load_0[30] = _fma_pair_25.x; + _tmem_load_0[31] = _fma_pair_25.y; + #pragma unroll + for (int _le = 0; _le < 32; _le++) { + _tmem_load_0[_le] = approx_exp2(_tmem_load_0[_le]); + } + float2 _reg_reduce_sum2_26 = make_float2(0.0f, 0.0f); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[0], _tmem_load_0[1])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[2], _tmem_load_0[3])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[4], _tmem_load_0[5])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[6], _tmem_load_0[7])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[8], _tmem_load_0[9])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[10], _tmem_load_0[11])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[12], _tmem_load_0[13])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[14], _tmem_load_0[15])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[16], _tmem_load_0[17])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[18], _tmem_load_0[19])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[20], _tmem_load_0[21])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[22], _tmem_load_0[23])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[24], _tmem_load_0[25])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[26], _tmem_load_0[27])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[28], _tmem_load_0[29])); + _reg_reduce_sum2_26 = add_f32x2(_reg_reduce_sum2_26, make_float2(_tmem_load_0[30], _tmem_load_0[31])); + float _tmem_load_0_sum = _reg_reduce_sum2_26.x + _reg_reduce_sum2_26.y; + block_sum = block_sum + _tmem_load_0_sum; + { + uint32_t _pv_packed[8]; + #pragma unroll + for (int _j = 0; _j < 8; _j++) { + uint32_t _pk; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_pk) : "f"(_tmem_load_0[0 + _j * 4]), "f"(_tmem_load_0[0 + _j * 4 + 1]), + "f"(_tmem_load_0[0 + _j * 4 + 2]), "f"(_tmem_load_0[0 + _j * 4 + 3])); + _pv_packed[_j] = _pk; + } + tmem_st_x8_u32(p_base_all, _pv_packed); + } + const float2 _fma_b2_27 = {softmax_scale_log2, softmax_scale_log2}; + const float2 _fma_c2_28 = {-max_scaled, -max_scaled}; + float2 _fma_pair_29 = fma_f32x2(make_float2(_tmem_load_2[0], _tmem_load_2[1]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[0] = _fma_pair_29.x; + _tmem_load_2[1] = _fma_pair_29.y; + float2 _fma_pair_30 = fma_f32x2(make_float2(_tmem_load_2[2], _tmem_load_2[3]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[2] = _fma_pair_30.x; + _tmem_load_2[3] = _fma_pair_30.y; + float2 _fma_pair_31 = fma_f32x2(make_float2(_tmem_load_2[4], _tmem_load_2[5]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[4] = _fma_pair_31.x; + _tmem_load_2[5] = _fma_pair_31.y; + float2 _fma_pair_32 = fma_f32x2(make_float2(_tmem_load_2[6], _tmem_load_2[7]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[6] = _fma_pair_32.x; + _tmem_load_2[7] = _fma_pair_32.y; + float2 _fma_pair_33 = fma_f32x2(make_float2(_tmem_load_2[8], _tmem_load_2[9]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[8] = _fma_pair_33.x; + _tmem_load_2[9] = _fma_pair_33.y; + float2 _fma_pair_34 = fma_f32x2(make_float2(_tmem_load_2[10], _tmem_load_2[11]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[10] = _fma_pair_34.x; + _tmem_load_2[11] = _fma_pair_34.y; + float2 _fma_pair_35 = fma_f32x2(make_float2(_tmem_load_2[12], _tmem_load_2[13]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[12] = _fma_pair_35.x; + _tmem_load_2[13] = _fma_pair_35.y; + float2 _fma_pair_36 = fma_f32x2(make_float2(_tmem_load_2[14], _tmem_load_2[15]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[14] = _fma_pair_36.x; + _tmem_load_2[15] = _fma_pair_36.y; + float2 _fma_pair_37 = fma_f32x2(make_float2(_tmem_load_2[16], _tmem_load_2[17]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[16] = _fma_pair_37.x; + _tmem_load_2[17] = _fma_pair_37.y; + float2 _fma_pair_38 = fma_f32x2(make_float2(_tmem_load_2[18], _tmem_load_2[19]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[18] = _fma_pair_38.x; + _tmem_load_2[19] = _fma_pair_38.y; + float2 _fma_pair_39 = fma_f32x2(make_float2(_tmem_load_2[20], _tmem_load_2[21]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[20] = _fma_pair_39.x; + _tmem_load_2[21] = _fma_pair_39.y; + float2 _fma_pair_40 = fma_f32x2(make_float2(_tmem_load_2[22], _tmem_load_2[23]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[22] = _fma_pair_40.x; + _tmem_load_2[23] = _fma_pair_40.y; + float2 _fma_pair_41 = fma_f32x2(make_float2(_tmem_load_2[24], _tmem_load_2[25]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[24] = _fma_pair_41.x; + _tmem_load_2[25] = _fma_pair_41.y; + float2 _fma_pair_42 = fma_f32x2(make_float2(_tmem_load_2[26], _tmem_load_2[27]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[26] = _fma_pair_42.x; + _tmem_load_2[27] = _fma_pair_42.y; + float2 _fma_pair_43 = fma_f32x2(make_float2(_tmem_load_2[28], _tmem_load_2[29]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[28] = _fma_pair_43.x; + _tmem_load_2[29] = _fma_pair_43.y; + float2 _fma_pair_44 = fma_f32x2(make_float2(_tmem_load_2[30], _tmem_load_2[31]), _fma_b2_27, _fma_c2_28); + _tmem_load_2[30] = _fma_pair_44.x; + _tmem_load_2[31] = _fma_pair_44.y; + #pragma unroll + for (int _le = 0; _le < 32; _le++) { + _tmem_load_2[_le] = approx_exp2(_tmem_load_2[_le]); + } + float2 _reg_reduce_sum2_45 = make_float2(0.0f, 0.0f); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[0], _tmem_load_2[1])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[2], _tmem_load_2[3])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[4], _tmem_load_2[5])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[6], _tmem_load_2[7])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[8], _tmem_load_2[9])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[10], _tmem_load_2[11])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[12], _tmem_load_2[13])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[14], _tmem_load_2[15])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[16], _tmem_load_2[17])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[18], _tmem_load_2[19])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[20], _tmem_load_2[21])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[22], _tmem_load_2[23])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[24], _tmem_load_2[25])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[26], _tmem_load_2[27])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[28], _tmem_load_2[29])); + _reg_reduce_sum2_45 = add_f32x2(_reg_reduce_sum2_45, make_float2(_tmem_load_2[30], _tmem_load_2[31])); + float _tmem_load_2_sum = _reg_reduce_sum2_45.x + _reg_reduce_sum2_45.y; + block_sum = block_sum + _tmem_load_2_sum; + { + uint32_t _pv_packed[8]; + #pragma unroll + for (int _j = 0; _j < 8; _j++) { + uint32_t _pk; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_pk) : "f"(_tmem_load_2[0 + _j * 4]), "f"(_tmem_load_2[0 + _j * 4 + 1]), + "f"(_tmem_load_2[0 + _j * 4 + 2]), "f"(_tmem_load_2[0 + _j * 4 + 3])); + _pv_packed[_j] = _pk; + } + tmem_st_x8_u32(p_base_all + 8, _pv_packed); + } + const float2 _fma_b2_46 = {softmax_scale_log2, softmax_scale_log2}; + const float2 _fma_c2_47 = {-max_scaled, -max_scaled}; + float2 _fma_pair_48 = fma_f32x2(make_float2(_tmem_load_3[0], _tmem_load_3[1]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[0] = _fma_pair_48.x; + _tmem_load_3[1] = _fma_pair_48.y; + float2 _fma_pair_49 = fma_f32x2(make_float2(_tmem_load_3[2], _tmem_load_3[3]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[2] = _fma_pair_49.x; + _tmem_load_3[3] = _fma_pair_49.y; + float2 _fma_pair_50 = fma_f32x2(make_float2(_tmem_load_3[4], _tmem_load_3[5]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[4] = _fma_pair_50.x; + _tmem_load_3[5] = _fma_pair_50.y; + float2 _fma_pair_51 = fma_f32x2(make_float2(_tmem_load_3[6], _tmem_load_3[7]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[6] = _fma_pair_51.x; + _tmem_load_3[7] = _fma_pair_51.y; + float2 _fma_pair_52 = fma_f32x2(make_float2(_tmem_load_3[8], _tmem_load_3[9]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[8] = _fma_pair_52.x; + _tmem_load_3[9] = _fma_pair_52.y; + float2 _fma_pair_53 = fma_f32x2(make_float2(_tmem_load_3[10], _tmem_load_3[11]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[10] = _fma_pair_53.x; + _tmem_load_3[11] = _fma_pair_53.y; + float2 _fma_pair_54 = fma_f32x2(make_float2(_tmem_load_3[12], _tmem_load_3[13]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[12] = _fma_pair_54.x; + _tmem_load_3[13] = _fma_pair_54.y; + float2 _fma_pair_55 = fma_f32x2(make_float2(_tmem_load_3[14], _tmem_load_3[15]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[14] = _fma_pair_55.x; + _tmem_load_3[15] = _fma_pair_55.y; + float2 _fma_pair_56 = fma_f32x2(make_float2(_tmem_load_3[16], _tmem_load_3[17]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[16] = _fma_pair_56.x; + _tmem_load_3[17] = _fma_pair_56.y; + float2 _fma_pair_57 = fma_f32x2(make_float2(_tmem_load_3[18], _tmem_load_3[19]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[18] = _fma_pair_57.x; + _tmem_load_3[19] = _fma_pair_57.y; + float2 _fma_pair_58 = fma_f32x2(make_float2(_tmem_load_3[20], _tmem_load_3[21]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[20] = _fma_pair_58.x; + _tmem_load_3[21] = _fma_pair_58.y; + float2 _fma_pair_59 = fma_f32x2(make_float2(_tmem_load_3[22], _tmem_load_3[23]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[22] = _fma_pair_59.x; + _tmem_load_3[23] = _fma_pair_59.y; + float2 _fma_pair_60 = fma_f32x2(make_float2(_tmem_load_3[24], _tmem_load_3[25]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[24] = _fma_pair_60.x; + _tmem_load_3[25] = _fma_pair_60.y; + float2 _fma_pair_61 = fma_f32x2(make_float2(_tmem_load_3[26], _tmem_load_3[27]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[26] = _fma_pair_61.x; + _tmem_load_3[27] = _fma_pair_61.y; + float2 _fma_pair_62 = fma_f32x2(make_float2(_tmem_load_3[28], _tmem_load_3[29]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[28] = _fma_pair_62.x; + _tmem_load_3[29] = _fma_pair_62.y; + float2 _fma_pair_63 = fma_f32x2(make_float2(_tmem_load_3[30], _tmem_load_3[31]), _fma_b2_46, _fma_c2_47); + _tmem_load_3[30] = _fma_pair_63.x; + _tmem_load_3[31] = _fma_pair_63.y; + #pragma unroll + for (int _le = 0; _le < 32; _le++) { + _tmem_load_3[_le] = approx_exp2(_tmem_load_3[_le]); + } + float2 _reg_reduce_sum2_64 = make_float2(0.0f, 0.0f); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[0], _tmem_load_3[1])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[2], _tmem_load_3[3])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[4], _tmem_load_3[5])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[6], _tmem_load_3[7])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[8], _tmem_load_3[9])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[10], _tmem_load_3[11])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[12], _tmem_load_3[13])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[14], _tmem_load_3[15])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[16], _tmem_load_3[17])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[18], _tmem_load_3[19])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[20], _tmem_load_3[21])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[22], _tmem_load_3[23])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[24], _tmem_load_3[25])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[26], _tmem_load_3[27])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[28], _tmem_load_3[29])); + _reg_reduce_sum2_64 = add_f32x2(_reg_reduce_sum2_64, make_float2(_tmem_load_3[30], _tmem_load_3[31])); + float _tmem_load_3_sum = _reg_reduce_sum2_64.x + _reg_reduce_sum2_64.y; + block_sum = block_sum + _tmem_load_3_sum; + { + uint32_t _pv_packed[8]; + #pragma unroll + for (int _j = 0; _j < 8; _j++) { + uint32_t _pk; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_pk) : "f"(_tmem_load_3[0 + _j * 4]), "f"(_tmem_load_3[0 + _j * 4 + 1]), + "f"(_tmem_load_3[0 + _j * 4 + 2]), "f"(_tmem_load_3[0 + _j * 4 + 3])); + _pv_packed[_j] = _pk; + } + tmem_st_x8_u32(p_base_all + 16, _pv_packed); + } + const float2 _fma_b2_65 = {softmax_scale_log2, softmax_scale_log2}; + const float2 _fma_c2_66 = {-max_scaled, -max_scaled}; + float2 _fma_pair_67 = fma_f32x2(make_float2(_tmem_load_4[0], _tmem_load_4[1]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[0] = _fma_pair_67.x; + _tmem_load_4[1] = _fma_pair_67.y; + float2 _fma_pair_68 = fma_f32x2(make_float2(_tmem_load_4[2], _tmem_load_4[3]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[2] = _fma_pair_68.x; + _tmem_load_4[3] = _fma_pair_68.y; + float2 _fma_pair_69 = fma_f32x2(make_float2(_tmem_load_4[4], _tmem_load_4[5]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[4] = _fma_pair_69.x; + _tmem_load_4[5] = _fma_pair_69.y; + float2 _fma_pair_70 = fma_f32x2(make_float2(_tmem_load_4[6], _tmem_load_4[7]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[6] = _fma_pair_70.x; + _tmem_load_4[7] = _fma_pair_70.y; + float2 _fma_pair_71 = fma_f32x2(make_float2(_tmem_load_4[8], _tmem_load_4[9]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[8] = _fma_pair_71.x; + _tmem_load_4[9] = _fma_pair_71.y; + float2 _fma_pair_72 = fma_f32x2(make_float2(_tmem_load_4[10], _tmem_load_4[11]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[10] = _fma_pair_72.x; + _tmem_load_4[11] = _fma_pair_72.y; + float2 _fma_pair_73 = fma_f32x2(make_float2(_tmem_load_4[12], _tmem_load_4[13]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[12] = _fma_pair_73.x; + _tmem_load_4[13] = _fma_pair_73.y; + float2 _fma_pair_74 = fma_f32x2(make_float2(_tmem_load_4[14], _tmem_load_4[15]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[14] = _fma_pair_74.x; + _tmem_load_4[15] = _fma_pair_74.y; + float2 _fma_pair_75 = fma_f32x2(make_float2(_tmem_load_4[16], _tmem_load_4[17]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[16] = _fma_pair_75.x; + _tmem_load_4[17] = _fma_pair_75.y; + float2 _fma_pair_76 = fma_f32x2(make_float2(_tmem_load_4[18], _tmem_load_4[19]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[18] = _fma_pair_76.x; + _tmem_load_4[19] = _fma_pair_76.y; + float2 _fma_pair_77 = fma_f32x2(make_float2(_tmem_load_4[20], _tmem_load_4[21]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[20] = _fma_pair_77.x; + _tmem_load_4[21] = _fma_pair_77.y; + float2 _fma_pair_78 = fma_f32x2(make_float2(_tmem_load_4[22], _tmem_load_4[23]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[22] = _fma_pair_78.x; + _tmem_load_4[23] = _fma_pair_78.y; + float2 _fma_pair_79 = fma_f32x2(make_float2(_tmem_load_4[24], _tmem_load_4[25]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[24] = _fma_pair_79.x; + _tmem_load_4[25] = _fma_pair_79.y; + float2 _fma_pair_80 = fma_f32x2(make_float2(_tmem_load_4[26], _tmem_load_4[27]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[26] = _fma_pair_80.x; + _tmem_load_4[27] = _fma_pair_80.y; + float2 _fma_pair_81 = fma_f32x2(make_float2(_tmem_load_4[28], _tmem_load_4[29]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[28] = _fma_pair_81.x; + _tmem_load_4[29] = _fma_pair_81.y; + float2 _fma_pair_82 = fma_f32x2(make_float2(_tmem_load_4[30], _tmem_load_4[31]), _fma_b2_65, _fma_c2_66); + _tmem_load_4[30] = _fma_pair_82.x; + _tmem_load_4[31] = _fma_pair_82.y; + #pragma unroll + for (int _le = 0; _le < 32; _le++) { + _tmem_load_4[_le] = approx_exp2(_tmem_load_4[_le]); + } + float2 _reg_reduce_sum2_83 = make_float2(0.0f, 0.0f); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[0], _tmem_load_4[1])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[2], _tmem_load_4[3])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[4], _tmem_load_4[5])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[6], _tmem_load_4[7])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[8], _tmem_load_4[9])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[10], _tmem_load_4[11])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[12], _tmem_load_4[13])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[14], _tmem_load_4[15])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[16], _tmem_load_4[17])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[18], _tmem_load_4[19])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[20], _tmem_load_4[21])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[22], _tmem_load_4[23])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[24], _tmem_load_4[25])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[26], _tmem_load_4[27])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[28], _tmem_load_4[29])); + _reg_reduce_sum2_83 = add_f32x2(_reg_reduce_sum2_83, make_float2(_tmem_load_4[30], _tmem_load_4[31])); + float _tmem_load_4_sum = _reg_reduce_sum2_83.x + _reg_reduce_sum2_83.y; + block_sum = block_sum + _tmem_load_4_sum; + { + uint32_t _pv_packed[8]; + #pragma unroll + for (int _j = 0; _j < 8; _j++) { + uint32_t _pk; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_pk) : "f"(_tmem_load_4[0 + _j * 4]), "f"(_tmem_load_4[0 + _j * 4 + 1]), + "f"(_tmem_load_4[0 + _j * 4 + 2]), "f"(_tmem_load_4[0 + _j * 4 + 3])); + _pv_packed[_j] = _pk; + } + tmem_st_x8_u32(p_base_all + 24, _pv_packed); + } + } + { + float _fma_1 = __fmaf_rn(row_sum_val, acc_scale, block_sum); + row_sum_val = _fma_1; + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(p_full_addr + (phase) * 8); + } + smem_stats_sum[stats_row] = row_sum_val; + smem_stats_final_max[stats_row] = row_max_val; + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(sum_ready_addr); + softmax_tile_cursor = softmax_tile_cursor + num_kv_tiles_1; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_dealloc_addr); + } + } + // ---- Role: correction_wg ---- + if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 184;"); + { // correction_wg_main + float softmax_scale_log2_1 = bmm1_scale[0] * 1.4426950408889634f; + float output_scale = bmm2_scale[0]; + const int wg_dummy_inc_1 = 0; + int all_num_kv_tiles_2 = (sparse_topk + 128 - 1) / 128; + const int tmem_row_base_1 = ((0) ? warp % 2 * 32 : warp % 4 * 32); + const int n_half_1 = ((0) ? (int)(warp % 4 / 2) : 0); + const int my_row_1 = tmem_row_base_1 + lane; + const int stats_row_1 = n_half_1 * 64 + my_row_1; + const int corr_row = tmem_row_base_1 << 16; + int correction_tile_cursor = 0; + unsigned int _phase_q_full_0_2 = 0; + unsigned int _phase_pv_done_0 = 0; + unsigned int _phase_o_done_0 = 0; + unsigned int _phase_sum_ready_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx_3 = bid; work_idx_3 < total_work_items; work_idx_3 += num_bids) { + int split_idx_2 = 0; + int query_idx_2 = work_idx_3 >> 1; + int v_chunk_1 = work_idx_3 & 1; + int tiles_per_split_2 = all_num_kv_tiles_2; + int first_tile_2 = split_idx_2 * tiles_per_split_2; + int num_kv_tiles_2 = tiles_per_split_2; + { + mbarrier_wait(q_full_addr, _phase_q_full_0_2); + _phase_q_full_0_2 ^= 1; + } + #pragma unroll 1 + for (int tile_2 = 0; tile_2 < num_kv_tiles_2; tile_2++) { + int pipeline_tile_1 = correction_tile_cursor + tile_2; + int phase_1 = pipeline_tile_1 & 1; + int stats_wait_phase = pipeline_tile_1 >> 1 & 1; + mbarrier_wait(stats_addr + (phase_1) * 8, stats_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + float acc_scale_1 = smem_stats_max[phase_1 * 128 + stats_row_1]; + if (tile_2 > 0) { + mbarrier_wait(pv_done_addr, _phase_pv_done_0); + _phase_pv_done_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + } + if (tile_2 > 0) { + int _vote_1 = __any_sync(0xFFFFFFFF, acc_scale_1 < 1.0f); + int any_rescale = _vote_1; + if (any_rescale != 0) { + #pragma unroll + for (int vs = 0; vs < 2; vs++) { + int o_base = taddr + 256 + (unsigned int)(vs * 128) + (unsigned int)corr_row; + #pragma unroll + for (int c = 0; c < 128; c += 16) { + float _tmem_load_5[16]; + tmem_ld_x16(&_tmem_load_5[0], o_base + c); + const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + mul_f32x2_inplace(&reinterpret_cast(_tmem_load_5)[_ls], _scale2_0); + tmem_st_x16_f32(o_base + c, _tmem_load_5); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + } + } + } + } + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(corr_done_addr + (phase_1) * 8); + } + mbarrier_wait(o_done_addr, _phase_o_done_0); + _phase_o_done_0 ^= 1; + mbarrier_wait(sum_ready_addr, _phase_sum_ready_0); + _phase_sum_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + float total_sum = smem_stats_sum[stats_row_1]; + float final_max = smem_stats_final_max[stats_row_1]; + float _rcp_0 = approx_rcp(total_sum); + float inv_sum = ((total_sum > 0.0f) ? _rcp_0 : 0.0f); + int head_idx = my_row_1; + int direct_o_offset = (query_idx_2 * num_heads + head_idx) * 512 + v_chunk_1 * 256; + int partial_o_offset = (query_idx_2 * num_heads + head_idx + split_idx_2) * 512 + ((0) ? v_chunk_1 : n_half_1) * 256; + int o_offset = direct_o_offset; + if ((0 & (int)(n_half_1 == 0) & (int)(head_idx < num_heads)) != 0) { + int stat_offset = query_idx_2 * num_heads + head_idx + split_idx_2; + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(total_sum)); + partial_lse[stat_offset] = ((total_sum > 0.0f) ? final_max * softmax_scale_log2_1 + _log2_0 : -CAKE_INF); + } + #pragma unroll + for (int vs_1 = 0; vs_1 < 2; vs_1++) { + int o_base_epi = taddr + 256 + (unsigned int)(vs_1 * 128) + (unsigned int)corr_row; + #pragma unroll + for (int c_1 = 0; c_1 < 128; c_1 += 32) { + float _tmem_load_6[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_6[0]), "=f"(_tmem_load_6[1]), "=f"(_tmem_load_6[2]), "=f"(_tmem_load_6[3]), "=f"(_tmem_load_6[4]), "=f"(_tmem_load_6[5]), "=f"(_tmem_load_6[6]), "=f"(_tmem_load_6[7]), "=f"(_tmem_load_6[8]), "=f"(_tmem_load_6[9]), "=f"(_tmem_load_6[10]), "=f"(_tmem_load_6[11]), "=f"(_tmem_load_6[12]), "=f"(_tmem_load_6[13]), "=f"(_tmem_load_6[14]), "=f"(_tmem_load_6[15]), "=f"(_tmem_load_6[16]), "=f"(_tmem_load_6[17]), "=f"(_tmem_load_6[18]), "=f"(_tmem_load_6[19]), "=f"(_tmem_load_6[20]), "=f"(_tmem_load_6[21]), "=f"(_tmem_load_6[22]), "=f"(_tmem_load_6[23]), "=f"(_tmem_load_6[24]), "=f"(_tmem_load_6[25]), "=f"(_tmem_load_6[26]), "=f"(_tmem_load_6[27]), "=f"(_tmem_load_6[28]), "=f"(_tmem_load_6[29]), "=f"(_tmem_load_6[30]), "=f"(_tmem_load_6[31]) + : "r"(o_base_epi + c_1)); + int gmem_base = o_offset + vs_1 * 128 + c_1; + if (head_idx < num_heads) { + #pragma unroll + for (int j = 0; j < 32; j += 16) { + { + const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; + #if __CUDA_ARCH__ >= 1000 + #pragma unroll + for (int _ps = 0; _ps < 8; _ps++) + mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_6[j])[_ps], _prescale2_1); + #else + #pragma unroll + for (int _ps = 0; _ps < 16; _ps++) + _tmem_load_6[j + _ps] *= inv_sum * output_scale; + #endif + { + __nv_bfloat162 _pk0 = __floats2bfloat162_rn(_tmem_load_6[j + 0], _tmem_load_6[j + 1]); + unsigned _pk_u0 = *reinterpret_cast(&_pk0); + __nv_bfloat162 _pk1 = __floats2bfloat162_rn(_tmem_load_6[j + 2], _tmem_load_6[j + 3]); + unsigned _pk_u1 = *reinterpret_cast(&_pk1); + __nv_bfloat162 _pk2 = __floats2bfloat162_rn(_tmem_load_6[j + 4], _tmem_load_6[j + 5]); + unsigned _pk_u2 = *reinterpret_cast(&_pk2); + __nv_bfloat162 _pk3 = __floats2bfloat162_rn(_tmem_load_6[j + 6], _tmem_load_6[j + 7]); + unsigned _pk_u3 = *reinterpret_cast(&_pk3); + __nv_bfloat162 _pk4 = __floats2bfloat162_rn(_tmem_load_6[j + 8], _tmem_load_6[j + 9]); + unsigned _pk_u4 = *reinterpret_cast(&_pk4); + __nv_bfloat162 _pk5 = __floats2bfloat162_rn(_tmem_load_6[j + 10], _tmem_load_6[j + 11]); + unsigned _pk_u5 = *reinterpret_cast(&_pk5); + __nv_bfloat162 _pk6 = __floats2bfloat162_rn(_tmem_load_6[j + 12], _tmem_load_6[j + 13]); + unsigned _pk_u6 = *reinterpret_cast(&_pk6); + __nv_bfloat162 _pk7 = __floats2bfloat162_rn(_tmem_load_6[j + 14], _tmem_load_6[j + 15]); + unsigned _pk_u7 = *reinterpret_cast(&_pk7); + asm volatile( + "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"((void*)(&((__nv_bfloat16*)(O + (gmem_base + j)))[0])), "r"(_pk_u0), "r"(_pk_u1), "r"(_pk_u2), "r"(_pk_u3), "r"(_pk_u4), "r"(_pk_u5), "r"(_pk_u6), "r"(_pk_u7) : "memory"); + } + } + } + } + } + } + correction_tile_cursor = correction_tile_cursor + num_kv_tiles_2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_dealloc_addr); + } + } + // ---- Role: mma_warp ---- + if (warp == 8) { + { // mma_warp_main + const int wg2_dummy_1 = 0; + int all_num_kv_tiles_3 = (sparse_topk + 128 - 1) / 128; + unsigned int mma_k_stage = 0; + unsigned int mma_v_stage = 0; + unsigned int mma_kv_stage = 0; + int mma_tile_cursor = 0; + unsigned int _phase_s_seeded_0 = 0; + unsigned int _phase_q_full_0_3 = 0; + unsigned int _phase_q_pair_ready_0 = 0; + unsigned int _phase_k_full = 0; + unsigned int _phase_kv_full = 0; + unsigned int _phase_kv_pair_ready = 0; + unsigned int _phase_v_full = 0; + #pragma unroll 1 + for (unsigned int work_idx_4 = bid; work_idx_4 < total_work_items; work_idx_4 += num_bids) { + int split_idx_3 = 0; + int tiles_per_split_3 = all_num_kv_tiles_3; + int first_tile_3 = split_idx_3 * tiles_per_split_3; + int num_kv_tiles_3 = tiles_per_split_3; + { + { + mbarrier_wait(s_seeded_addr, _phase_s_seeded_0); + _phase_s_seeded_0 ^= 1; + } + asm volatile("tcgen05.fence::after_thread_sync;"); + } + { + mbarrier_wait(q_full_addr, _phase_q_full_0_3); + _phase_q_full_0_3 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + } + int first_pv = 1; + #pragma unroll 1 + for (int tile_3 = 0; tile_3 < num_kv_tiles_3; tile_3++) { + int pipeline_tile_2 = mma_tile_cursor + tile_3; + int phase_2 = pipeline_tile_2 & 1; + int score_col = ((phase_2 != 0) ? 128 : 0); + { + #pragma unroll 1 + for (int n = 0; n < 4; n++) { + mbarrier_wait(kv_full_addr + (mma_kv_stage) * 8, _phase_kv_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + { + int _mma_a_lo_4 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (n) * 1024); + int _mma_b_lo_4 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 1024); + { + uint64_t _mma_ss_a_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_4); + uint64_t _mma_ss_b_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_4); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem_scratch + (score_col)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 136314896, ((n == 0) ? 0 : 1)); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem_scratch + (score_col)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 136314896, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem_scratch + (score_col)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 136314896, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem_scratch + (score_col)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 136314896, 1); + } + } + if (n == 3) { + elect_commit(s_full_addr + (phase_2) * 8); + } + elect_commit(kv_empty_addr + (mma_kv_stage) * 8); + } + mma_kv_stage += 1; + if (mma_kv_stage == 8) { mma_kv_stage = 0; _phase_kv_full ^= 1; _phase_kv_pair_ready ^= 1; } + } + } + if (tile_3 > 0) { + int prev_pipeline_tile = pipeline_tile_2 - 1; + int prev_phase = prev_pipeline_tile & 1; + int pv_wait_phase = prev_pipeline_tile >> 1 & 1; + mbarrier_wait(p_full_addr + (prev_phase) * 8, pv_wait_phase); + mbarrier_wait(corr_done_addr + (prev_phase) * 8, pv_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int p_col = ((prev_phase != 0) ? 224 : 96); + { + #pragma unroll + for (int vs_2 = 0; vs_2 < 2; vs_2++) { + mbarrier_wait(kv_full_addr + (mma_kv_stage) * 8, _phase_kv_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + int output_col = 256 + vs_2 * 128; + { + int _mma_b_lo_8 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (mma_kv_stage) * 1024); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 dhi, blo, ta, id;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 136380432;\n\t" + "mov.b32 ta, %2;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p0;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" + "}\n" + :: "r"((tmem_tmem_scratch + (output_col))), "r"(_mma_b_lo_8), "r"(tmem_tmem_scratch + p_col), "r"(((first_pv) ? 0 : 1))); + elect_commit(kv_empty_addr + (mma_kv_stage) * 8); + } + mma_kv_stage += 1; + if (mma_kv_stage == 8) { mma_kv_stage = 0; _phase_kv_full ^= 1; _phase_kv_pair_ready ^= 1; } + } + } + first_pv = 0; + elect_commit(pv_done_addr); + } + } + int last_pipeline_tile = mma_tile_cursor + num_kv_tiles_3 - 1; + int last_phase = last_pipeline_tile & 1; + int drain_wait_phase = last_pipeline_tile >> 1 & 1; + mbarrier_wait(p_full_addr + (last_phase) * 8, drain_wait_phase); + mbarrier_wait(corr_done_addr + (last_phase) * 8, drain_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int p_col_last = ((last_phase != 0) ? 224 : 96); + { + #pragma unroll + for (int vs_3 = 0; vs_3 < 2; vs_3++) { + mbarrier_wait(kv_full_addr + (mma_kv_stage) * 8, _phase_kv_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + int output_col_d = 256 + vs_3 * 128; + { + int _mma_b_lo_12 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (mma_kv_stage) * 1024); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 dhi, blo, ta, id;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 136380432;\n\t" + "mov.b32 ta, %2;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p0;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" + "}\n" + :: "r"((tmem_tmem_scratch + (output_col_d))), "r"(_mma_b_lo_12), "r"(tmem_tmem_scratch + p_col_last), "r"(((first_pv) ? 0 : 1))); + elect_commit(kv_empty_addr + (mma_kv_stage) * 8); + } + mma_kv_stage += 1; + if (mma_kv_stage == 8) { mma_kv_stage = 0; _phase_kv_full ^= 1; _phase_kv_pair_ready ^= 1; } + } + } + elect_commit(q_empty_addr); + elect_commit(o_done_addr); + mma_tile_cursor = mma_tile_cursor + num_kv_tiles_3; + } + mbarrier_arrive(tmem_dealloc_addr); + unsigned int _phase_tmem_dealloc_0 = 0; + mbarrier_wait(tmem_dealloc_addr, _phase_tmem_dealloc_0); + _phase_tmem_dealloc_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + { + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + } + } + // ---- Role: empty1 ---- + if (warp >= 14 && warp <= 15) { + { // empty1_main + const int wg2_dummy_2 = 0; + unsigned int _phase_q_full_0_4 = 0; + { + #pragma unroll 1 + for (unsigned int work_idx_5 = bid; work_idx_5 < total_work_items; work_idx_5 += num_bids) { + mbarrier_wait(q_full_addr, _phase_q_full_0_4); + _phase_q_full_0_4 ^= 1; + } + } + } + } + + // Cleanup +} + +} // extern "C" diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_5703afd5b398f2b6c980_binding.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_5703afd5b398f2b6c980_binding.cu new file mode 100644 index 00000000000..e9e716552b4 --- /dev/null +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_5703afd5b398f2b6c980_binding.cu @@ -0,0 +1,461 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +// Generated by the Cake source exporter — do not edit. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_5703afd5b398f2b6c980'. +#include +#include +#include +#include + +#include "tvm_ffi_utils.h" + +#include +#include +#include +#include +#include +#include +#include + +extern "C" __global__ void kernel_cake_dsv4_5703afd5b398f2b6c980(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items); + + +namespace cake_host_shim_5678cf46270e8859 { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t CakeDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +// 4D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + int64_t d3 = t.size(t.ndim() - 3); + TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + TVM_FFI_CHECK(d1 % 128 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 128; + uint64_t global_dim[4] = {(uint64_t)(128), (uint64_t)(d2), (uint64_t)((d1 / 128)), (uint64_t)(d3)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (128, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 128; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + (uint64_t)((carrier_stride_1 * 8) / 8), + (uint64_t)((carrier_stride_2 * 8) / 8), + }; + uint32_t box_dim[4] = {128u, 64u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_swa_kv' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_swa_kv' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; + TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (128, 1) exceeds resolved global dims for 'tmap_swa_kv'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + }; + uint32_t box_dim[2] = {128u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_compressed_kv' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_compressed_kv(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_compressed_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_compressed_kv' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_compressed_kv' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_compressed_kv"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_compressed_kv' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved a non-positive global dim"; + TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (128, 1) exceeds resolved global dims for 'tmap_compressed_kv'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + }; + uint32_t box_dim[2] = {128u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_kv') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_compressed_kv, TensorView arg_O, TensorView arg_partial_lse, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_num_query_tokens, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 1, 8, 1); + CheckContiguous(arg_tmap_q, "tmap_q"); + CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); + CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 1, 8, 1); + CheckCudaTensor(arg_tmap_compressed_kv, "tmap_compressed_kv"); + CheckDtype(arg_tmap_compressed_kv, "tmap_compressed_kv", 1, 8, 1); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_partial_lse, "partial_lse"); + CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); + CheckContiguous(arg_partial_lse, "partial_lse"); + CheckCudaTensor(arg_swa_indices, "swa_indices"); + CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); + CheckContiguous(arg_swa_indices, "swa_indices"); + CheckCudaTensor(arg_compressed_indices, "compressed_indices"); + CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); + CheckContiguous(arg_compressed_indices, "compressed_indices"); + CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); + CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); + CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); + CheckContiguous(arg_bmm1_scale, "bmm1_scale"); + CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); + CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); + CheckContiguous(arg_bmm2_scale, "bmm2_scale"); + TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) + << "scalar 'num_heads' value " << arg_num_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) + << "scalar 'swa_index_stride' value " << arg_swa_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) + << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) + << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_query_tokens >= -2147483648LL && arg_num_query_tokens <= 2147483647LL, ValueError) + << "scalar 'num_query_tokens' value " << arg_num_query_tokens + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) + << "scalar 'sparse_topk' value " << arg_sparse_topk + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) + << "scalar 'has_sinks' value " << arg_has_sinks + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); + CheckSameCudaDevice(arg_tmap_compressed_kv, arg_tmap_q, "tmap_compressed_kv", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_partial_lse, arg_tmap_q, "partial_lse", "tmap_q"); + CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); + CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); + CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); + CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") must be divisible by cluster dims (2, 1, 1)"; + + cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); + CUtensorMap p_tmap_compressed_kv = EncodeTma_tmap_compressed_kv(arg_tmap_compressed_kv); + __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); + float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); + int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); + int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); + int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); + float* p_sinks = static_cast(arg_sinks.data_ptr()); + float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); + float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); + int32_t v_num_heads = (int32_t)arg_num_heads; + int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; + int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; + int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; + int32_t v_num_query_tokens = (int32_t)arg_num_query_tokens; + int32_t v_sparse_topk = (int32_t)arg_sparse_topk; + int32_t v_has_sinks = (int32_t)arg_has_sinks; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_tmap_compressed_kv, &p_O, &p_partial_lse, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_num_query_tokens, &v_sparse_topk, &v_has_sinks, &v_total_work_items}; + + static std::mutex smem_mu; + static auto* smem_status_by_device = + new std::unordered_map(); + cudaError_t smem_status = cudaSuccess; + { + std::lock_guard lock(smem_mu); + auto it = smem_status_by_device->find(dev.device_id); + if (it == smem_status_by_device->end()) { + int current_device = -1; + smem_status = cudaGetDevice(¤t_device); + if (smem_status == cudaSuccess && current_device != dev.device_id) { + smem_status = cudaErrorInvalidDevice; + } + if (smem_status == cudaSuccess) { + smem_status = cudaFuncSetAttribute( + reinterpret_cast(kernel_cake_dsv4_5703afd5b398f2b6c980), + cudaFuncAttributeMaxDynamicSharedMemorySize, + 223744); + } + smem_status_by_device->emplace(dev.device_id, smem_status); + } else { + smem_status = it->second; + } + } + TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) + << "cudaFuncSetAttribute for kernel_cake_dsv4_5703afd5b398f2b6c980 failed: " + << cudaGetErrorString(smem_status); + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(512u, 1u, 1u); + + // Extended direct-source launch for a linked CUDA kernel symbol. + cudaLaunchAttribute attrs[2]{}; + int n = 0; + attrs[n].id = cudaLaunchAttributeClusterDimension; + attrs[n].val.clusterDim.x = 2u; + attrs[n].val.clusterDim.y = 1u; + attrs[n].val.clusterDim.z = 1u; + ++n; + attrs[n].id = cudaLaunchAttributeClusterSchedulingPolicyPreference; + attrs[n].val.clusterSchedulingPolicyPreference = cudaClusterSchedulingPolicySpread; + ++n; + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes = 223744u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; + cudaError_t launch_status = cudaLaunchKernelExC( + &config, reinterpret_cast(kernel_cake_dsv4_5703afd5b398f2b6c980), kargs); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernelExC for kernel_cake_dsv4_5703afd5b398f2b6c980 failed: " + << cudaGetErrorString(launch_status); +} + +} // namespace cake_host_shim_5678cf46270e8859 + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_5678cf46270e8859::Run); diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_5703afd5b398f2b6c980_kernel.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_5703afd5b398f2b6c980_kernel.cu new file mode 100644 index 00000000000..e481fa741c8 --- /dev/null +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_5703afd5b398f2b6c980_kernel.cu @@ -0,0 +1,2803 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +#define CAKE_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_SCRATCH_OFFSET 0 +#define NUM_K_PIPE_STAGES 3 +#define NUM_V_PIPE_STAGES 2 +#define NUM_KV_PIPE_STAGES 1 +#define NUM_INDEX_PIPE_STAGES 6 +#define SMEM_SMEM_Q_OFF 1024 +#define SMEM_SMEM_Q_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_STRIDE 8192 +#define SMEM_SMEM_KV_OFF 33792 +#define SMEM_SMEM_KV_STAGE_BYTES 16384 +#define SMEM_SMEM_KV_STRIDE 16384 +#define SMEM_SMEM_V_OFF 33792 +#define SMEM_SMEM_V_STAGE_BYTES 16384 +#define SMEM_SMEM_V_STRIDE 16384 +#define SMEM_SMEM_Q_FULL_OFF 1024 +#define SMEM_SMEM_Q_FULL_STAGE_BYTES 32768 +#define SMEM_SMEM_Q_FULL_STRIDE 32768 +#define SMEM_SMEM_K_FULL_OFF 33792 +#define SMEM_SMEM_K_FULL_STAGE_BYTES 32768 +#define SMEM_SMEM_K_FULL_STRIDE 32768 +#define SMEM_SMEM_V_FULL_OFF 132096 +#define SMEM_SMEM_V_FULL_STAGE_BYTES 32768 +#define SMEM_SMEM_V_FULL_STRIDE 32768 +#define SMEM_SMEM_STATS_MAX_OFF 197632 +#define SMEM_SMEM_STATS_MAX_STAGE_BYTES 1024 +#define SMEM_SMEM_STATS_MAX_STRIDE 1024 +#define SMEM_SMEM_STATS_SUM_OFF 198656 +#define SMEM_SMEM_STATS_SUM_STAGE_BYTES 512 +#define SMEM_SMEM_STATS_SUM_STRIDE 512 +#define SMEM_SMEM_STATS_FINAL_MAX_OFF 199168 +#define SMEM_SMEM_STATS_FINAL_MAX_STAGE_BYTES 512 +#define SMEM_SMEM_STATS_FINAL_MAX_STRIDE 512 +#define SMEM_SMEM_SOFTMAX_WARP_PAIR_EXCHANGE_OFF 199680 +#define SMEM_SMEM_SOFTMAX_WARP_PAIR_EXCHANGE_STAGE_BYTES 1024 +#define SMEM_SMEM_SOFTMAX_WARP_PAIR_EXCHANGE_STRIDE 1024 +#define SMEM_SMEM_CORR_WARP_PAIR_EXCHANGE_OFF 200704 +#define SMEM_SMEM_CORR_WARP_PAIR_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_CORR_WARP_PAIR_EXCHANGE_STRIDE 512 +#define SMEM_SMEM_SPARSE_INDICES_OFF 201216 +#define SMEM_SMEM_SPARSE_INDICES_STAGE_BYTES 1024 +#define SMEM_SMEM_SPARSE_INDICES_STRIDE 1024 +#define SMEM_SMEM_P_FP8_OFF 207360 +#define SMEM_SMEM_P_FP8_STAGE_BYTES 8192 +#define SMEM_SMEM_P_FP8_STRIDE 8192 +#define SMEM_TOTAL 223744 +#define THREADS 512 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count) : "memory"); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f8f6f4_cta2( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" + "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "tcgen05.mma.cta_group::2.kind::f8f6f4 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" + "}\n" + :: "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x4(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x4.b32" + " {%0, %1, %2, %3}, [%4];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); +} + + +__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ float2 ex2_emulation_f32x2_value(float2 value) { + const float c0 = 1.0f, c1 = 0.695146143436431884765625f; + const float c2 = 0.227564394474029541015625f, c3 = 0.077119089663028717041015625f; + const float magic = 12582912.0f; + float x0 = max_noftz(value.x, -127.0f), x1 = max_noftz(value.y, -127.0f); + float2 xc2 = make_float2(x0, x1), magic2 = make_float2(magic, magic); + float2 xr2; + asm("add.rm.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xr2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&magic2)); + float2 c3_2 = make_float2(c3, c3), c2_2 = make_float2(c2, c2); + float2 c1_2 = make_float2(c1, c1), c0_2 = make_float2(c0, c0); + float2 xrb2, xfrac2; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xrb2) + : "l"(*(unsigned long long*)&xr2), "l"(*(unsigned long long*)&magic2)); + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xfrac2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&xrb2)); + float2 poly2; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&c3_2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c2_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c1_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c0_2)); + int x0r_i, x1r_i, p0_i, p1_i; + asm("mov.b64 {%0, %1}, %2;" : "=r"(x0r_i), "=r"(x1r_i) : "l"(*(unsigned long long*)&xr2)); + asm("mov.b64 {%0, %1}, %2;" : "=r"(p0_i), "=r"(p1_i) : "l"(*(unsigned long long*)&poly2)); + float r0, r1; + asm("mov.b32 %0, %1;" : "=f"(r0) : "r"((x0r_i << 23) + p0_i)); + asm("mov.b32 %0, %1;" : "=f"(r1) : "r"((x1r_i << 23) + p1_i)); + return make_float2(r0, r1); +} + +__device__ __forceinline__ void ex2_emulation_f32x2(float* x0_ptr, float* x1_ptr) { + float2 result = ex2_emulation_f32x2_value(make_float2(*x0_ptr, *x1_ptr)); + *x0_ptr = result.x; *x1_ptr = result.y; +} + +__device__ __forceinline__ void softmax_frag_exp2_cast( + float* sv, uint32_t* pv, int use_emu) +{ + #pragma unroll + for (int j = 0; j < 16; j++) { + if (use_emu && j >= 12) + ex2_emulation_f32x2(&sv[j*2], &sv[j*2+1]); + else { + sv[j*2] = approx_exp2(sv[j*2]); + sv[j*2+1] = approx_exp2(sv[j*2+1]); + } + } + #pragma unroll + for (int j = 0; j < 16; j++) { + __nv_bfloat162 bf = __float22bfloat162_rn({sv[j*2], sv[j*2+1]}); + pv[j] = reinterpret_cast(bf); + } +} + + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_2d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3}], [%4];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr) { + // Canonical .shared::cta form for non-multicast gather4, matching + // trtllm-gen / cuda_ptx and the PTX ISA qualifier order + // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, + // .shared::cluster is reserved for the multicast variant (ctaMask). + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem_cta2( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr) { + // Canonical .shared::cta form; see tma_gather4_gmem2smem above. + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem_mc( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr, unsigned short cta_mask) { + // Multicast variant: the PTX grammar ties the .shared::cluster + // destination to .multicast::cluster + ctaMask (cf. cuda_ptx / + // SM100_TMA_LOAD_MULTICAST_2D_GATHER4 in CUTLASS). + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" + ".mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem_mc_cta2( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr, unsigned short cta_mask) { + // Multicast + cta_group::2 variant; see tma_gather4_gmem2smem_mc. + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" + ".mbarrier::complete_tx::bytes.multicast::cluster.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %1;\n\t" + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" + "}\n" + :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); +} + + +__device__ __forceinline__ void tmem_st_x8_u32(int addr, const uint32_t* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x8.b32" + " [%0], {%1,%2,%3,%4,%5,%6,%7,%8};" + :: "r"(addr), + "r"(src[0]), "r"(src[1]), "r"(src[2]), "r"(src[3]), + "r"(src[4]), "r"(src[5]), "r"(src[6]), "r"(src[7])); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(512, 1) __cluster_dims__(2,1,1) void +kernel_cake_dsv4_5703afd5b398f2b6c980(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); + + const int mbar_base = smem; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define k_full_addr (mbar_base + 16) + #define k_empty_addr (mbar_base + 40) + #define v_full_addr (mbar_base + 64) + #define v_empty_addr (mbar_base + 80) + #define index_full_addr (mbar_base + 96) + #define index_empty_addr (mbar_base + 144) + #define s_full_addr (mbar_base + 192) + #define p_full_addr (mbar_base + 208) + #define corr_done_addr (mbar_base + 224) + #define stats_addr (mbar_base + 240) + #define sum_ready_addr (mbar_base + 256) + #define o_done_addr (mbar_base + 264) + #define pv_done_addr (mbar_base + 272) + #define s_seeded_addr (mbar_base + 280) + #define q_pair_ready_addr (mbar_base + 288) + #define kv_pair_ready_addr (mbar_base + 296) + #define pv_pair_ready_addr (mbar_base + 304) + #define tmem_dealloc_addr (mbar_base + 320) + #define tmem_dealloc_peer_addr (mbar_base + 328) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 2; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + + // Kernel setup ops + uint8_t* smem_q = reinterpret_cast(smem_raw + 1024); + const int smem_q_addr = smem + 1024; + uint8_t* smem_kv = reinterpret_cast(smem_raw + 33792); + const int smem_kv_addr = smem + 33792; + uint8_t* smem_v = reinterpret_cast(smem_raw + 33792); + const int smem_v_addr = smem + 33792; + uint8_t* smem_q_full = reinterpret_cast(smem_raw + 1024); + const int smem_q_full_addr = smem + 1024; + uint8_t* smem_k_full = reinterpret_cast(smem_raw + 33792); + const int smem_k_full_addr = smem + 33792; + uint8_t* smem_v_full = reinterpret_cast(smem_raw + 132096); + const int smem_v_full_addr = smem + 132096; + float* smem_stats_max = reinterpret_cast(smem_raw + 197632); + const int smem_stats_max_addr = smem + 197632; + float* smem_stats_sum = reinterpret_cast(smem_raw + 198656); + const int smem_stats_sum_addr = smem + 198656; + float* smem_stats_final_max = reinterpret_cast(smem_raw + 199168); + const int smem_stats_final_max_addr = smem + 199168; + float* smem_softmax_warp_pair_exchange = reinterpret_cast(smem_raw + 199680); + const int smem_softmax_warp_pair_exchange_addr = smem + 199680; + float* smem_corr_warp_pair_exchange = reinterpret_cast(smem_raw + 200704); + const int smem_corr_warp_pair_exchange_addr = smem + 200704; + int* smem_sparse_indices = reinterpret_cast(smem_raw + 201216); + const int smem_sparse_indices_addr = smem + 201216; + uint8_t* smem_p_fp8 = reinterpret_cast(smem_raw + 207360); + const int smem_p_fp8_addr = smem + 207360; + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // Mbarrier init (21 pipeline groups, 0 ordered-sequence groups, 42 barriers) + // Mbarriers at smem_raw[0..336) + + if (warp == 12) { + uint32_t leader = elect_sync(); + if (leader) { + // q_full: 1 barriers, init_count=2 + mbarrier_init(smem + 0, 2); + // q_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 8, 1); + // --- pipeline 'k_pipe' --- + // k_full: 3 barriers, init_count=2 + mbarrier_init(smem + 16, 2); + mbarrier_init(smem + 24, 2); + mbarrier_init(smem + 32, 2); + // k_empty: 3 barriers, init_count=1 + mbarrier_init(smem + 40, 1); + mbarrier_init(smem + 48, 1); + mbarrier_init(smem + 56, 1); + // --- pipeline 'v_pipe' --- + // v_full: 2 barriers, init_count=2 + mbarrier_init(smem + 64, 2); + mbarrier_init(smem + 72, 2); + // v_empty: 2 barriers, init_count=1 + mbarrier_init(smem + 80, 1); + mbarrier_init(smem + 88, 1); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + if (warp == 9) { + uint32_t leader = elect_sync(); + if (leader) { + // --- pipeline 'index_pipe' --- + // index_full: 6 barriers, init_count=32 + mbarrier_init(smem + 96, 32); + mbarrier_init(smem + 104, 32); + mbarrier_init(smem + 112, 32); + mbarrier_init(smem + 120, 32); + mbarrier_init(smem + 128, 32); + mbarrier_init(smem + 136, 32); + // index_empty: 6 barriers, init_count=256 + mbarrier_init(smem + 144, 256); + mbarrier_init(smem + 152, 256); + mbarrier_init(smem + 160, 256); + mbarrier_init(smem + 168, 256); + mbarrier_init(smem + 176, 256); + mbarrier_init(smem + 184, 256); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // s_full: stages (0,), init_count=1 + mbarrier_init(smem + 192, 1); + // o_done: 1 barriers, init_count=1 + mbarrier_init(smem + 264, 1); + // pv_done: 1 barriers, init_count=1 + mbarrier_init(smem + 272, 1); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + if (warp == 1) { + uint32_t leader = elect_sync(); + if (leader) { + // s_full: stages (1,), init_count=1 + mbarrier_init(smem + 200, 1); + // p_full: 2 barriers, init_count=128 + mbarrier_init(smem + 208, 128); + mbarrier_init(smem + 216, 128); + // corr_done: stages (0,), init_count=128 + mbarrier_init(smem + 224, 128); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + if (warp == 5) { + uint32_t leader = elect_sync(); + if (leader) { + // corr_done: stages (1,), init_count=128 + mbarrier_init(smem + 232, 128); + // stats: 2 barriers, init_count=128 + mbarrier_init(smem + 240, 128); + mbarrier_init(smem + 248, 128); + // sum_ready: 1 barriers, init_count=128 + mbarrier_init(smem + 256, 128); + // tmem_dealloc_peer: 1 barriers, init_count=32 + mbarrier_init(smem + 328, 32); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + if (warp == 2) { + uint32_t leader = elect_sync(); + if (leader) { + // s_seeded: 1 barriers, init_count=256 + mbarrier_init(smem + 280, 256); + // q_pair_ready: 1 barriers, init_count=64 + mbarrier_init(smem + 288, 64); + // kv_pair_ready: 1 barriers, init_count=64 + mbarrier_init(smem + 296, 64); + // pv_pair_ready: stages (0,), init_count=64 + mbarrier_init(smem + 304, 64); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + if (warp == 4) { + uint32_t leader = elect_sync(); + if (leader) { + // pv_pair_ready: stages (1,), init_count=64 + mbarrier_init(smem + 312, 64); + // tmem_dealloc: 1 barriers, init_count=416 + mbarrier_init(smem + 320, 416); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncwarp(); + + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 336); + if (warp == 0) { + int _tmem_hold = smem + 336; + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); + } + + __syncthreads(); + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_tmem_scratch = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Inc phase consumes the registers released above. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 136;"); + } + + // ---- Role: index_warp ---- + if (warp == 9) { + { // index_warp_main + const int index_dummy = 0; + unsigned int _phase_index_empty = 1; + unsigned int _phase_q_full_0 = 0; + { + int all_num_kv_tiles = (sparse_topk + 128 - 1) / 128; + unsigned int index_prod_stage = 0; + #pragma unroll + for (unsigned int static_work = 0; static_work < 1; static_work++) { + int split_idx = 0; + int query_idx = bid / 2; + int tiles_per_split = all_num_kv_tiles + 1 - 1; + int first_tile = split_idx * tiles_per_split; + int num_index_passes = (tiles_per_split + 1) / 2; + #pragma unroll 1 + for (int index_pass = 0; index_pass < num_index_passes; index_pass++) { + mbarrier_wait(index_empty_addr + (index_prod_stage) * 8, _phase_index_empty); + int index_stage_base = smem_sparse_indices_addr + index_prod_stage * 1024; + int stage_row_base = index_pass * 256; + #pragma unroll + for (int index_half = 0; index_half < 2; index_half++) { + int index_offset = index_half * 128 + lane * 4; + int tile_col_base = first_tile * 128 + stage_row_base + index_half * 128; + int* row_tile_ptr = ((tile_col_base < 128) ? swa_indices + (query_idx * swa_index_stride) : compressed_indices + (query_idx * compressed_index_stride + (tile_col_base - 128))); + int lane_col = tile_col_base + lane * 4; + if (lane_col + 4 <= sparse_topk) { + asm volatile("cp.async.cg.shared::cta.global.L2::128B [%0], [%1], 16;" + :: "r"(index_stage_base + index_offset * 4), "l"(row_tile_ptr + (lane * 4))); + } else { + int tail_rows[4]; + #pragma unroll + for (int row_i = 0; row_i < 4; row_i++) { + tail_rows[row_i] = -1; + } + #pragma unroll + for (int row_i_1 = 0; row_i_1 < 4; row_i_1++) { + if (lane_col + row_i_1 < sparse_topk) { + tail_rows[row_i_1] = row_tile_ptr[lane * 4 + row_i_1]; + } + } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(smem_sparse_indices_addr + ((unsigned int)index_stage_base - smem_sparse_indices_addr + (unsigned int)(index_offset * 4))), "r"(tail_rows[0]), "r"(tail_rows[1]), "r"(tail_rows[2]), "r"(tail_rows[3]) : "memory"); + } + } + asm volatile("cp.async.commit_group;"); + asm volatile("cp.async.wait_group 0;"); + mbarrier_arrive(index_full_addr + (index_prod_stage) * 8); + index_prod_stage += 1; + if (index_prod_stage == 6) { index_prod_stage = 0; _phase_index_empty ^= 1; } + } + } + } + } + } + // ---- Role: load_warp ---- + if (warp >= 12 && warp <= 15) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 72;"); + { // load_warp_main + const int wg2_dummy = 0; + const int load_warp_rank = warp - ((1) ? 12 : 9); + int all_num_kv_tiles_1 = (sparse_topk + 128 - 1) / 128; + unsigned int load_k_stage = 0; + unsigned int load_v_stage = 0; + unsigned int load_kv_stage = 0; + unsigned int load_k_index_stage = 0; + unsigned int load_v_index_stage = 0; + int k_cta_offset = bid % 2 * 64; + unsigned int _phase_q_empty_0 = 1; + unsigned int _phase_index_full = 0; + unsigned int _phase_k_empty = 1; + unsigned int _phase_k_empty_1 = 1; + unsigned int _phase_v_empty = 1; + #pragma unroll + for (unsigned int work_idx = 0; work_idx < 1; work_idx++) { + int split_idx_1 = 0; + int query_idx_1 = bid / 2 >> 1; + int v_chunk = bid / 2 & 1; + { + split_idx_1 = 0; + query_idx_1 = bid / 2; + v_chunk = bid % 2; + } + int tiles_per_split_1 = all_num_kv_tiles_1 + 1 - 1; + int first_tile_1 = split_idx_1 * tiles_per_split_1; + int num_kv_tiles = tiles_per_split_1; + int sparse_extent = num_kv_tiles * 128; + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + if (load_warp_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + #pragma unroll + for (int q_stage = 0; q_stage < 4; q_stage++) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global.mbarrier::complete_tx::bytes.multicast::cluster.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5}], [%6], %7;" + :: "r"(smem_q_full_addr + (unsigned int)(q_stage * 8192)), "l"((&tmap_q)), "r"(0), "r"(bid % 2 * 64), "r"(q_stage), "r"(query_idx_1), + "r"(((q_full_addr) & 0xFEFFFFFF)), "h"((uint16_t)(1 << bid % 2)) : "memory"); + } + } + } + { + mbarrier_wait(index_full_addr + (load_k_index_stage) * 8, _phase_index_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + } + int k_index_stage_base = smem_sparse_indices_addr + load_k_index_stage * 1024; + { + mbarrier_wait(k_empty_addr + (load_k_stage) * 8, _phase_k_empty); + if (load_warp_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + } + } + int k_dst = smem_k_full_addr + load_k_stage * 32768 + (unsigned int)(load_warp_rank * 8192); + int group = lane; + if (group < 16) { + int group_offset = k_cta_offset + group * 4; + int raw_rows[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows[0])), "=r"(*reinterpret_cast(&raw_rows[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows[(0) + 3])) + : "r"(k_index_stage_base + group_offset * 4)); + int row0 = ((raw_rows[0] >= 0) ? raw_rows[0] : 0); + int row1 = ((raw_rows[1] >= 0) ? raw_rows[1] : 0); + int row2 = ((raw_rows[2] >= 0) ? raw_rows[2] : 0); + int row3 = ((raw_rows[3] >= 0) ? raw_rows[3] : 0); + if (first_tile_1 == 0) { + tma_gather4_gmem2smem_mc_cta2(k_dst + group * 512, (&tmap_swa_kv), load_warp_rank * 128, row0, row1, row2, row3, ((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); + } else { + tma_gather4_gmem2smem_mc_cta2(k_dst + group * 512, (&tmap_compressed_kv), load_warp_rank * 128, row0, row1, row2, row3, ((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); + } + } + load_k_stage += 1; + if (load_k_stage == 3) { load_k_stage = 0; _phase_k_empty ^= 1; } + } + #pragma unroll 1 + for (int tile = 1; tile < num_kv_tiles; tile++) { + if ((1 & (int)((tile & 1) == 0)) != 0) { + mbarrier_wait(index_full_addr + (load_k_index_stage) * 8, _phase_index_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + } + k_index_stage_base = smem_sparse_indices_addr + load_k_index_stage * 1024; + { + mbarrier_wait(k_empty_addr + (load_k_stage) * 8, _phase_k_empty); + if (load_warp_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + } + } + int k_dst_1 = smem_k_full_addr + load_k_stage * 32768 + (unsigned int)(load_warp_rank * 8192); + int group_1 = lane; + if (group_1 < 16) { + int group_offset_1 = (tile & 1) * 128 + k_cta_offset + group_1 * 4; + int raw_rows_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_1[0])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 3])) + : "r"(k_index_stage_base + group_offset_1 * 4)); + int row0_1 = ((raw_rows_1[0] >= 0) ? raw_rows_1[0] : 0); + int row1_1 = ((raw_rows_1[1] >= 0) ? raw_rows_1[1] : 0); + int row2_1 = ((raw_rows_1[2] >= 0) ? raw_rows_1[2] : 0); + int row3_1 = ((raw_rows_1[3] >= 0) ? raw_rows_1[3] : 0); + tma_gather4_gmem2smem_mc_cta2(k_dst_1 + group_1 * 512, (&tmap_compressed_kv), load_warp_rank * 128, row0_1, row1_1, row2_1, row3_1, ((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); + } + load_k_stage += 1; + if (load_k_stage == 3) { load_k_stage = 0; _phase_k_empty ^= 1; } + if ((tile & 1) != 0) { + load_k_index_stage += 1; + if (load_k_index_stage == 6) { load_k_index_stage = 0; _phase_index_full ^= 1; } + } + } + int prev_tile = tile - 1; + int global_prev_tile = first_tile_1 + prev_tile; + int v_index_stage_base = smem_sparse_indices_addr + load_v_index_stage * 1024; + { + mbarrier_wait(v_empty_addr + (load_v_stage) * 8, _phase_v_empty); + if (load_warp_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + } + } + if (lane < 16) { + int v_dst = smem_v_full_addr + load_v_stage * 32768 + (unsigned int)(load_warp_rank / 2 * 16384); + int v_col = v_chunk * 256 + load_warp_rank / 2 * 128; + int group_2 = (load_warp_rank & 1) * 16 + lane; + int group_offset_2 = (prev_tile & 1) * 128 + group_2 * 4; + int raw_rows_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_2[0])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 3])) + : "r"(v_index_stage_base + group_offset_2 * 4)); + int row0_2 = ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : 0); + int row1_2 = ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : 0); + int row2_2 = ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : 0); + int row3_2 = ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : 0); + if (global_prev_tile == 0) { + tma_gather4_gmem2smem_mc_cta2(v_dst + group_2 * 512, (&tmap_swa_kv), v_col, row0_2, row1_2, row2_2, row3_2, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); + } else { + tma_gather4_gmem2smem_mc_cta2(v_dst + group_2 * 512, (&tmap_compressed_kv), v_col, row0_2, row1_2, row2_2, row3_2, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); + } + } + load_v_stage += 1; + if (load_v_stage == 2) { load_v_stage = 0; _phase_v_empty ^= 1; } + if ((prev_tile & 1) != 0) { + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(index_empty_addr + (load_v_index_stage) * 8); + load_v_index_stage = load_v_index_stage + 1; + if (load_v_index_stage == 6) { + load_v_index_stage = 0; + } + } + } + } + int last_tile = num_kv_tiles - 1; + int global_last_tile = first_tile_1 + last_tile; + { + mbarrier_wait(v_empty_addr + (load_v_stage) * 8, _phase_v_empty); + if (load_warp_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + } + } + if (lane < 16) { + int v_dst_1 = smem_v_full_addr + load_v_stage * 32768 + (unsigned int)(load_warp_rank / 2 * 16384); + int v_col_1 = v_chunk * 256 + load_warp_rank / 2 * 128; + int group_3 = (load_warp_rank & 1) * 16 + lane; + int group_offset_3 = (last_tile & 1) * 128 + group_3 * 4; + int raw_rows_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_3[0])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 3])) + : "r"(smem_sparse_indices_addr + load_v_index_stage * 1024 + (unsigned int)(group_offset_3 * 4))); + int row0_3 = ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : 0); + int row1_3 = ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : 0); + int row2_3 = ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : 0); + int row3_3 = ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : 0); + if (global_last_tile == 0) { + tma_gather4_gmem2smem_mc_cta2(v_dst_1 + group_3 * 512, (&tmap_swa_kv), v_col_1, row0_3, row1_3, row2_3, row3_3, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); + } else { + tma_gather4_gmem2smem_mc_cta2(v_dst_1 + group_3 * 512, (&tmap_compressed_kv), v_col_1, row0_3, row1_3, row2_3, row3_3, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); + } + } + load_v_stage += 1; + if (load_v_stage == 2) { load_v_stage = 0; _phase_v_empty ^= 1; } + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(index_empty_addr + (load_v_index_stage) * 8); + load_v_index_stage = load_v_index_stage + 1; + if (load_v_index_stage == 6) { + load_v_index_stage = 0; + } + } + } + mbarrier_arrive(tmem_dealloc_addr); + } + } + // ---- Role: softmax_wg ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 152;"); + { // softmax_wg_main + float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; + const int wg_dummy_inc = 0; + int all_num_kv_tiles_2 = (sparse_topk + 128 - 1) / 128; + const int tmem_row_base = ((1) ? warp % 2 * 32 : warp % 4 * 32); + const int tmem_score_row_base = ((1) ? (int)(warp % 4 * 32) : tmem_row_base); + const int n_half = ((1) ? (int)(warp % 4 / 2) : 0); + const int my_row = tmem_row_base + lane; + const int stats_row = n_half * 64 + my_row; + int softmax_tile_cursor = 0; + unsigned int softmax_index_stage = 0; + unsigned int _phase_q_full_0_1 = 0; + #pragma unroll + for (unsigned int work_idx_1 = 0; work_idx_1 < 1; work_idx_1++) { + int split_idx_2 = 0; + int query_idx_2 = bid / 2 >> 1; + { + split_idx_2 = 0; + query_idx_2 = bid / 2; + } + int tiles_per_split_2 = all_num_kv_tiles_2 + 1 - 1; + int first_tile_2 = split_idx_2 * tiles_per_split_2; + int num_kv_tiles_1 = tiles_per_split_2; + int _max_0 = ((sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) : (0)); + int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); + int active_topk = _min_0; + float row_max_val = -CAKE_INF; + float row_sum_val = 0.0f; + int sink_head = ((1) ? bid % 2 * 64 + my_row : my_row); + if (has_sinks != 0 && sink_head < num_heads && split_idx_2 == 0) { + row_max_val = sinks[sink_head] * 1.4426950408889634f / softmax_scale_log2; + row_sum_val = 1.0f; + } + #pragma unroll 1 + for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { + int pipeline_tile = softmax_tile_cursor + tile_1; + int phase = pipeline_tile & 1; + int s_wait_phase = pipeline_tile >> 1 & 1; + mbarrier_wait(s_full_addr + (phase) * 8, s_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int s_off = ((phase != 0) ? 128 : 0); + int s_base = taddr + (unsigned int)s_off + (unsigned int)(tmem_score_row_base << 16); + float new_max = row_max_val; + int valid_sparse_cols = ((active_topk < sparse_topk) ? active_topk : sparse_topk); + valid_sparse_cols = valid_sparse_cols - (first_tile_2 + tile_1) * 128 - n_half * 64; + if (valid_sparse_cols < 0) { + valid_sparse_cols = 0; + } + if (valid_sparse_cols > 64) { + valid_sparse_cols = 64; + } + int index_tile_addr = smem_sparse_indices_addr + softmax_index_stage * 1024 + (unsigned int)((tile_1 & 1) * 512); + int index_lane_addr = index_tile_addr + (n_half * 64 + lane) * 4; + int staged_index[2]; + unsigned int invalid_cols[2]; + unsigned int any_invalid = 0; + #pragma unroll + for (int w = 0; w < 2; w++) { + asm volatile("ld.shared.b32 %0, [%1];" : "=r"(*reinterpret_cast(&staged_index[w])) : "r"(index_lane_addr + w * 128)); + unsigned int _vote_0 = __ballot_sync(0xFFFFFFFF, staged_index[w] < 0); + invalid_cols[w] = _vote_0; + any_invalid = any_invalid | invalid_cols[w]; + } + if ((tile_1 & 1) != 0 || tile_1 + 1 == num_kv_tiles_1) { + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(index_empty_addr + (softmax_index_stage) * 8); + softmax_index_stage = softmax_index_stage + 1; + if (softmax_index_stage == 6) { + softmax_index_stage = 0; + } + } + float _tmem_load_0[4]; + tmem_ld_x4(&_tmem_load_0[0], s_base); + float _tmem_load_1[64]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]) + : "r"(s_base)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]) + : "r"(s_base + 32)); + { + uint32_t _slice_lo_mask_0; + { + int _lim_0 = valid_sparse_cols; + if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_0 & (1u << 0))) _tmem_load_1[0] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 1))) _tmem_load_1[1] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 2))) _tmem_load_1[2] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 3))) _tmem_load_1[3] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 4))) _tmem_load_1[4] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 5))) _tmem_load_1[5] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 6))) _tmem_load_1[6] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 7))) _tmem_load_1[7] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 8))) _tmem_load_1[8] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 9))) _tmem_load_1[9] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 10))) _tmem_load_1[10] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 11))) _tmem_load_1[11] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 12))) _tmem_load_1[12] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 13))) _tmem_load_1[13] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 14))) _tmem_load_1[14] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 15))) _tmem_load_1[15] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 16))) _tmem_load_1[16] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 17))) _tmem_load_1[17] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 18))) _tmem_load_1[18] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 19))) _tmem_load_1[19] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 20))) _tmem_load_1[20] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 21))) _tmem_load_1[21] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 22))) _tmem_load_1[22] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 23))) _tmem_load_1[23] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 24))) _tmem_load_1[24] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 25))) _tmem_load_1[25] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 26))) _tmem_load_1[26] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 27))) _tmem_load_1[27] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 28))) _tmem_load_1[28] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 29))) _tmem_load_1[29] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 30))) _tmem_load_1[30] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 31))) _tmem_load_1[31] = -CAKE_INF; + uint32_t _slice_lo_mask_1; + { + int _lim_1 = valid_sparse_cols - 32; + if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) _tmem_load_1[32] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) _tmem_load_1[33] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) _tmem_load_1[34] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) _tmem_load_1[35] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) _tmem_load_1[36] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) _tmem_load_1[37] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) _tmem_load_1[38] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) _tmem_load_1[39] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) _tmem_load_1[40] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) _tmem_load_1[41] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) _tmem_load_1[42] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) _tmem_load_1[43] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) _tmem_load_1[44] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) _tmem_load_1[45] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) _tmem_load_1[46] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) _tmem_load_1[47] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) _tmem_load_1[48] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) _tmem_load_1[49] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) _tmem_load_1[50] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) _tmem_load_1[51] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) _tmem_load_1[52] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) _tmem_load_1[53] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) _tmem_load_1[54] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) _tmem_load_1[55] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) _tmem_load_1[56] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) _tmem_load_1[57] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) _tmem_load_1[58] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) _tmem_load_1[59] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) _tmem_load_1[60] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) _tmem_load_1[61] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) _tmem_load_1[62] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) _tmem_load_1[63] = -CAKE_INF; + if (any_invalid != 0) { + #pragma unroll + for (int w_1 = 0; w_1 < 2; w_1++) { + unsigned int mask_word = invalid_cols[w_1]; + _tmem_load_1[w_1 * 32] = (((mask_word & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32]); + _tmem_load_1[w_1 * 32 + 1] = (((mask_word >> 1 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 1]); + _tmem_load_1[w_1 * 32 + 2] = (((mask_word >> 2 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 2]); + _tmem_load_1[w_1 * 32 + 3] = (((mask_word >> 3 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 3]); + _tmem_load_1[w_1 * 32 + 4] = (((mask_word >> 4 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 4]); + _tmem_load_1[w_1 * 32 + 5] = (((mask_word >> 5 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 5]); + _tmem_load_1[w_1 * 32 + 6] = (((mask_word >> 6 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 6]); + _tmem_load_1[w_1 * 32 + 7] = (((mask_word >> 7 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 7]); + _tmem_load_1[w_1 * 32 + 8] = (((mask_word >> 8 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 8]); + _tmem_load_1[w_1 * 32 + 9] = (((mask_word >> 9 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 9]); + _tmem_load_1[w_1 * 32 + 10] = (((mask_word >> 10 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 10]); + _tmem_load_1[w_1 * 32 + 11] = (((mask_word >> 11 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 11]); + _tmem_load_1[w_1 * 32 + 12] = (((mask_word >> 12 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 12]); + _tmem_load_1[w_1 * 32 + 13] = (((mask_word >> 13 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 13]); + _tmem_load_1[w_1 * 32 + 14] = (((mask_word >> 14 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 14]); + _tmem_load_1[w_1 * 32 + 15] = (((mask_word >> 15 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 15]); + _tmem_load_1[w_1 * 32 + 16] = (((mask_word >> 16 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 16]); + _tmem_load_1[w_1 * 32 + 17] = (((mask_word >> 17 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 17]); + _tmem_load_1[w_1 * 32 + 18] = (((mask_word >> 18 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 18]); + _tmem_load_1[w_1 * 32 + 19] = (((mask_word >> 19 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 19]); + _tmem_load_1[w_1 * 32 + 20] = (((mask_word >> 20 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 20]); + _tmem_load_1[w_1 * 32 + 21] = (((mask_word >> 21 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 21]); + _tmem_load_1[w_1 * 32 + 22] = (((mask_word >> 22 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 22]); + _tmem_load_1[w_1 * 32 + 23] = (((mask_word >> 23 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 23]); + _tmem_load_1[w_1 * 32 + 24] = (((mask_word >> 24 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 24]); + _tmem_load_1[w_1 * 32 + 25] = (((mask_word >> 25 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 25]); + _tmem_load_1[w_1 * 32 + 26] = (((mask_word >> 26 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 26]); + _tmem_load_1[w_1 * 32 + 27] = (((mask_word >> 27 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 27]); + _tmem_load_1[w_1 * 32 + 28] = (((mask_word >> 28 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 28]); + _tmem_load_1[w_1 * 32 + 29] = (((mask_word >> 29 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 29]); + _tmem_load_1[w_1 * 32 + 30] = (((mask_word >> 30 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 30]); + _tmem_load_1[w_1 * 32 + 31] = (((mask_word >> 31 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 31]); + } + } + float2 _reg_reduce_max2_2 = {-CAKE_INF, -CAKE_INF}; + row_max_x32_accum(&_tmem_load_1[0], _reg_reduce_max2_2); + row_max_x32_accum(&_tmem_load_1[32], _reg_reduce_max2_2); + float _tmem_load_1_max = row_max_reduce(_reg_reduce_max2_2); + float _max_1 = max_noftz(new_max, _tmem_load_1_max); + new_max = _max_1; + smem_softmax_warp_pair_exchange[phase * 128 + stats_row] = new_max; + asm volatile("barrier.sync %0, 64;" :: "r"(2 + warp % 2) : "memory"); + float _max_2 = max_noftz(new_max, smem_softmax_warp_pair_exchange[phase * 128 + (stats_row ^ 64)]); + new_max = _max_2; + } + float _fma_0 = __fmaf_rn(row_max_val, softmax_scale_log2, (-new_max) * softmax_scale_log2); + float delta = _fma_0; + float _exp2_0 = approx_exp2(delta); + float exp_delta = _exp2_0; + float acc_scale = ((row_max_val > -CAKE_INF) ? exp_delta : 1.0f); + smem_stats_max[phase * 128 + stats_row] = acc_scale; + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(stats_addr + (phase) * 8); + row_max_val = new_max; + float safe_max = ((new_max == -CAKE_INF) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + float block_sum = 0.0f; + { + { + const float2 _fma_b2_3 = {softmax_scale_log2, softmax_scale_log2}; + const float2 _fma_c2_4 = {-max_scaled, -max_scaled}; + #pragma unroll + for (int _lf = 0; _lf < 32; _lf++) + fma_f32x2_inplace(&reinterpret_cast(_tmem_load_1)[_lf], _fma_b2_3, _fma_c2_4); + #pragma unroll + for (int _le = 0; _le < 64; _le++) { + _tmem_load_1[_le] = approx_exp2(_tmem_load_1[_le]); + } + uint32_t _fp8_0[16]; + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[0]), "f"(_tmem_load_1[1]), + "f"(_tmem_load_1[2]), "f"(_tmem_load_1[3])); + _fp8_0[0] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[4]), "f"(_tmem_load_1[5]), + "f"(_tmem_load_1[6]), "f"(_tmem_load_1[7])); + _fp8_0[1] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[8]), "f"(_tmem_load_1[9]), + "f"(_tmem_load_1[10]), "f"(_tmem_load_1[11])); + _fp8_0[2] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[12]), "f"(_tmem_load_1[13]), + "f"(_tmem_load_1[14]), "f"(_tmem_load_1[15])); + _fp8_0[3] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[16]), "f"(_tmem_load_1[17]), + "f"(_tmem_load_1[18]), "f"(_tmem_load_1[19])); + _fp8_0[4] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[20]), "f"(_tmem_load_1[21]), + "f"(_tmem_load_1[22]), "f"(_tmem_load_1[23])); + _fp8_0[5] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[24]), "f"(_tmem_load_1[25]), + "f"(_tmem_load_1[26]), "f"(_tmem_load_1[27])); + _fp8_0[6] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[28]), "f"(_tmem_load_1[29]), + "f"(_tmem_load_1[30]), "f"(_tmem_load_1[31])); + _fp8_0[7] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[32]), "f"(_tmem_load_1[33]), + "f"(_tmem_load_1[34]), "f"(_tmem_load_1[35])); + _fp8_0[8] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[36]), "f"(_tmem_load_1[37]), + "f"(_tmem_load_1[38]), "f"(_tmem_load_1[39])); + _fp8_0[9] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[40]), "f"(_tmem_load_1[41]), + "f"(_tmem_load_1[42]), "f"(_tmem_load_1[43])); + _fp8_0[10] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[44]), "f"(_tmem_load_1[45]), + "f"(_tmem_load_1[46]), "f"(_tmem_load_1[47])); + _fp8_0[11] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[48]), "f"(_tmem_load_1[49]), + "f"(_tmem_load_1[50]), "f"(_tmem_load_1[51])); + _fp8_0[12] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[52]), "f"(_tmem_load_1[53]), + "f"(_tmem_load_1[54]), "f"(_tmem_load_1[55])); + _fp8_0[13] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[56]), "f"(_tmem_load_1[57]), + "f"(_tmem_load_1[58]), "f"(_tmem_load_1[59])); + _fp8_0[14] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[60]), "f"(_tmem_load_1[61]), + "f"(_tmem_load_1[62]), "f"(_tmem_load_1[63])); + _fp8_0[15] = _packed; + } + int p_stage_base = smem_p_fp8_addr + (unsigned int)(phase * 8192); + #pragma unroll + for (int p_vec = 0; p_vec < 4; p_vec++) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_stage_base + (my_row * 128 + (n_half * 4 + p_vec) * 16 ^ (my_row * 128 + (n_half * 4 + p_vec) * 16 >> 7 & 7) << 4))), "r"(_fp8_0[p_vec * 4]), "r"(_fp8_0[p_vec * 4 + 1]), "r"(_fp8_0[p_vec * 4 + 2]), "r"(_fp8_0[p_vec * 4 + 3]) : "memory"); + } + } + } + { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + { + asm volatile("" ::: "memory"); + } + mbarrier_arrive(p_full_addr + (phase) * 8); + { + { + asm volatile("" ::: "memory"); + const float2* _reg_reduce_src2_5 = reinterpret_cast(&_tmem_load_1[0]); + float2 _reg_reduce_sum2_5_0 = make_float2(0.0f, 0.0f); + float2 _reg_reduce_sum2_5_1 = make_float2(0.0f, 0.0f); + float2 _reg_reduce_sum2_5_2 = make_float2(0.0f, 0.0f); + float2 _reg_reduce_sum2_5_3 = make_float2(0.0f, 0.0f); + #pragma unroll + for (int _rr = 0; _rr < 8; _rr++) { + add_f32x2_inplace(&_reg_reduce_sum2_5_0, _reg_reduce_src2_5[_rr * 4 + 0]); + add_f32x2_inplace(&_reg_reduce_sum2_5_1, _reg_reduce_src2_5[_rr * 4 + 1]); + add_f32x2_inplace(&_reg_reduce_sum2_5_2, _reg_reduce_src2_5[_rr * 4 + 2]); + add_f32x2_inplace(&_reg_reduce_sum2_5_3, _reg_reduce_src2_5[_rr * 4 + 3]); + } + add_f32x2_inplace(&_reg_reduce_sum2_5_0, _reg_reduce_sum2_5_1); + add_f32x2_inplace(&_reg_reduce_sum2_5_2, _reg_reduce_sum2_5_3); + add_f32x2_inplace(&_reg_reduce_sum2_5_0, _reg_reduce_sum2_5_2); + float _tmem_load_1_sum = _reg_reduce_sum2_5_0.x + _reg_reduce_sum2_5_0.y; + block_sum = _tmem_load_1_sum; + float _fma_2 = __fmaf_rn(row_sum_val, acc_scale, block_sum); + row_sum_val = _fma_2; + } + } + } + smem_stats_sum[stats_row] = row_sum_val; + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(sum_ready_addr); + softmax_tile_cursor = softmax_tile_cursor + num_kv_tiles_1; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_dealloc_addr); + } + } + // ---- Role: correction_wg ---- + if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 144;"); + { // correction_wg_main + float softmax_scale_log2_1 = bmm1_scale[0] * 1.4426950408889634f; + float output_scale = bmm2_scale[0]; + const int wg_dummy_inc_1 = 0; + int all_num_kv_tiles_3 = (sparse_topk + 128 - 1) / 128; + const int tmem_row_base_1 = ((1) ? warp % 2 * 32 : warp % 4 * 32); + const int n_half_1 = ((1) ? (int)(warp % 4 / 2) : 0); + const int my_row_1 = tmem_row_base_1 + lane; + const int stats_row_1 = n_half_1 * 64 + my_row_1; + const int corr_row = tmem_row_base_1 << 16; + int correction_tile_cursor = 0; + unsigned int _phase_q_full_0_2 = 0; + unsigned int _phase_pv_done_0 = 0; + unsigned int _phase_o_done_0 = 0; + unsigned int _phase_sum_ready_0 = 0; + #pragma unroll + for (unsigned int work_idx_2 = 0; work_idx_2 < 1; work_idx_2++) { + int split_idx_3 = 0; + int query_idx_3 = bid / 2 >> 1; + int v_chunk_1 = bid / 2 & 1; + { + split_idx_3 = 0; + query_idx_3 = bid / 2; + v_chunk_1 = bid % 2; + } + int tiles_per_split_3 = all_num_kv_tiles_3 + 1 - 1; + int first_tile_3 = split_idx_3 * tiles_per_split_3; + int num_kv_tiles_2 = tiles_per_split_3; + #pragma unroll 1 + for (int tile_2 = 0; tile_2 < num_kv_tiles_2; tile_2++) { + int pipeline_tile_1 = correction_tile_cursor + tile_2; + int phase_1 = pipeline_tile_1 & 1; + int stats_wait_phase = pipeline_tile_1 >> 1 & 1; + mbarrier_wait(stats_addr + (phase_1) * 8, stats_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + float acc_scale_1 = smem_stats_max[phase_1 * 128 + stats_row_1]; + if (tile_2 > 0) { + mbarrier_wait(pv_done_addr, _phase_pv_done_0); + _phase_pv_done_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + } + if (tile_2 > 0) { + int _vote_1 = __any_sync(0xFFFFFFFF, acc_scale_1 < 1.0f); + int any_rescale = _vote_1; + if (any_rescale != 0) { + #pragma unroll + for (int vs = 0; vs < 2; vs++) { + int o_base = taddr + 256 + (unsigned int)(vs * 128) + (unsigned int)corr_row; + #pragma unroll + for (int c = 0; c < 128; c += 16) { + float _tmem_load_4[16]; + tmem_ld_x16(&_tmem_load_4[0], o_base + c); + const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + mul_f32x2_inplace(&reinterpret_cast(_tmem_load_4)[_ls], _scale2_0); + tmem_st_x16_f32(o_base + c, _tmem_load_4); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + } + } + } + } + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(corr_done_addr + (phase_1) * 8); + } + mbarrier_wait(o_done_addr, _phase_o_done_0); + _phase_o_done_0 ^= 1; + mbarrier_wait(sum_ready_addr, _phase_sum_ready_0); + _phase_sum_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + float total_sum = smem_stats_sum[stats_row_1]; + float final_max = smem_stats_final_max[stats_row_1]; + { + smem_corr_warp_pair_exchange[stats_row_1] = total_sum; + asm volatile("barrier.sync %0, 64;" :: "r"(4 + warp % 2) : "memory"); + total_sum = total_sum + smem_corr_warp_pair_exchange[stats_row_1 ^ 64]; + } + float _rcp_0 = approx_rcp(total_sum); + float inv_sum = ((total_sum > 0.0f) ? _rcp_0 : 0.0f); + int head_idx = my_row_1; + { + head_idx = bid % 2 * 64 + my_row_1; + } + int direct_o_offset = (query_idx_3 * num_heads + head_idx) * 512 + v_chunk_1 * 256; + int partial_o_offset = (query_idx_3 * num_heads + head_idx + split_idx_3) * 512 + ((0) ? v_chunk_1 : n_half_1) * 256; + int o_offset = partial_o_offset; + #pragma unroll + for (int vs_1 = 0; vs_1 < 2; vs_1++) { + int o_base_epi = taddr + 256 + (unsigned int)(vs_1 * 128) + (unsigned int)corr_row; + #pragma unroll + for (int c_1 = 0; c_1 < 128; c_1 += 64) { + float _tmem_load_5[64]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x64.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, [%64];" + : "=f"(_tmem_load_5[0]), "=f"(_tmem_load_5[1]), "=f"(_tmem_load_5[2]), "=f"(_tmem_load_5[3]), "=f"(_tmem_load_5[4]), "=f"(_tmem_load_5[5]), "=f"(_tmem_load_5[6]), "=f"(_tmem_load_5[7]), "=f"(_tmem_load_5[8]), "=f"(_tmem_load_5[9]), "=f"(_tmem_load_5[10]), "=f"(_tmem_load_5[11]), "=f"(_tmem_load_5[12]), "=f"(_tmem_load_5[13]), "=f"(_tmem_load_5[14]), "=f"(_tmem_load_5[15]), "=f"(_tmem_load_5[16]), "=f"(_tmem_load_5[17]), "=f"(_tmem_load_5[18]), "=f"(_tmem_load_5[19]), "=f"(_tmem_load_5[20]), "=f"(_tmem_load_5[21]), "=f"(_tmem_load_5[22]), "=f"(_tmem_load_5[23]), "=f"(_tmem_load_5[24]), "=f"(_tmem_load_5[25]), "=f"(_tmem_load_5[26]), "=f"(_tmem_load_5[27]), "=f"(_tmem_load_5[28]), "=f"(_tmem_load_5[29]), "=f"(_tmem_load_5[30]), "=f"(_tmem_load_5[31]), "=f"(_tmem_load_5[32]), "=f"(_tmem_load_5[33]), "=f"(_tmem_load_5[34]), "=f"(_tmem_load_5[35]), "=f"(_tmem_load_5[36]), "=f"(_tmem_load_5[37]), "=f"(_tmem_load_5[38]), "=f"(_tmem_load_5[39]), "=f"(_tmem_load_5[40]), "=f"(_tmem_load_5[41]), "=f"(_tmem_load_5[42]), "=f"(_tmem_load_5[43]), "=f"(_tmem_load_5[44]), "=f"(_tmem_load_5[45]), "=f"(_tmem_load_5[46]), "=f"(_tmem_load_5[47]), "=f"(_tmem_load_5[48]), "=f"(_tmem_load_5[49]), "=f"(_tmem_load_5[50]), "=f"(_tmem_load_5[51]), "=f"(_tmem_load_5[52]), "=f"(_tmem_load_5[53]), "=f"(_tmem_load_5[54]), "=f"(_tmem_load_5[55]), "=f"(_tmem_load_5[56]), "=f"(_tmem_load_5[57]), "=f"(_tmem_load_5[58]), "=f"(_tmem_load_5[59]), "=f"(_tmem_load_5[60]), "=f"(_tmem_load_5[61]), "=f"(_tmem_load_5[62]), "=f"(_tmem_load_5[63]) + : "r"(o_base_epi + c_1)); + int gmem_base = o_offset + vs_1 * 128 + c_1; + if (bid % 2 == 0) { + #pragma unroll + for (int j = 0; j < 64; j += 16) { + { + const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; + #if __CUDA_ARCH__ >= 1000 + #pragma unroll + for (int _ps = 0; _ps < 8; _ps++) + mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_5[j])[_ps], _prescale2_1); + #else + #pragma unroll + for (int _ps = 0; _ps < 16; _ps++) + _tmem_load_5[j + _ps] *= inv_sum * output_scale; + #endif + { + __nv_bfloat162 _pk0 = __floats2bfloat162_rn(_tmem_load_5[j + 0], _tmem_load_5[j + 1]); + unsigned _pk_u0 = *reinterpret_cast(&_pk0); + __nv_bfloat162 _pk1 = __floats2bfloat162_rn(_tmem_load_5[j + 2], _tmem_load_5[j + 3]); + unsigned _pk_u1 = *reinterpret_cast(&_pk1); + __nv_bfloat162 _pk2 = __floats2bfloat162_rn(_tmem_load_5[j + 4], _tmem_load_5[j + 5]); + unsigned _pk_u2 = *reinterpret_cast(&_pk2); + __nv_bfloat162 _pk3 = __floats2bfloat162_rn(_tmem_load_5[j + 6], _tmem_load_5[j + 7]); + unsigned _pk_u3 = *reinterpret_cast(&_pk3); + __nv_bfloat162 _pk4 = __floats2bfloat162_rn(_tmem_load_5[j + 8], _tmem_load_5[j + 9]); + unsigned _pk_u4 = *reinterpret_cast(&_pk4); + __nv_bfloat162 _pk5 = __floats2bfloat162_rn(_tmem_load_5[j + 10], _tmem_load_5[j + 11]); + unsigned _pk_u5 = *reinterpret_cast(&_pk5); + __nv_bfloat162 _pk6 = __floats2bfloat162_rn(_tmem_load_5[j + 12], _tmem_load_5[j + 13]); + unsigned _pk_u6 = *reinterpret_cast(&_pk6); + __nv_bfloat162 _pk7 = __floats2bfloat162_rn(_tmem_load_5[j + 14], _tmem_load_5[j + 15]); + unsigned _pk_u7 = *reinterpret_cast(&_pk7); + asm volatile( + "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"((void*)(&((__nv_bfloat16*)(O + (gmem_base + j)))[0])), "r"(_pk_u0), "r"(_pk_u1), "r"(_pk_u2), "r"(_pk_u3), "r"(_pk_u4), "r"(_pk_u5), "r"(_pk_u6), "r"(_pk_u7) : "memory"); + } + } + } + } + } + } + correction_tile_cursor = correction_tile_cursor + num_kv_tiles_2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_dealloc_addr); + } + } + // ---- Role: mma_warp ---- + if (warp == 8) { + { // mma_warp_main + const int wg2_dummy_1 = 0; + int all_num_kv_tiles_4 = (sparse_topk + 128 - 1) / 128; + unsigned int mma_k_stage = 0; + unsigned int mma_v_stage = 0; + unsigned int mma_kv_stage = 0; + int mma_tile_cursor = 0; + unsigned int _phase_s_seeded_0 = 0; + unsigned int _phase_q_full_0_3 = 0; + unsigned int _phase_q_pair_ready_0 = 0; + unsigned int _phase_k_full = 0; + unsigned int _phase_k_full_1 = 0; + unsigned int _phase_kv_pair_ready = 0; + unsigned int _phase_v_full = 0; + #pragma unroll + for (unsigned int work_idx_3 = 0; work_idx_3 < 1; work_idx_3++) { + int split_idx_4 = 0; + { + split_idx_4 = 0; + } + int tiles_per_split_4 = all_num_kv_tiles_4 + 1 - 1; + int first_tile_4 = split_idx_4 * tiles_per_split_4; + int num_kv_tiles_3 = tiles_per_split_4; + { + if (cta_rank == 0) { + mbarrier_wait(q_full_addr, _phase_q_full_0_3); + _phase_q_full_0_3 ^= 1; + } + } + int first_pv = 1; + #pragma unroll 1 + for (int tile_3 = 0; tile_3 < num_kv_tiles_3; tile_3++) { + int pipeline_tile_2 = mma_tile_cursor + tile_3; + int phase_2 = pipeline_tile_2 & 1; + int score_col = ((phase_2 != 0) ? 128 : 0); + { + if (cta_rank == 0) { + mbarrier_wait(k_full_addr + (mma_k_stage) * 8, _phase_k_full); + int _mma_a_lo_0 = (((smem_q_full_addr) >> 4) & 0x3FFF) + (0) * 2048; + int _mma_b_lo_0 = (((smem_k_full_addr) >> 4) & 0x3FFF) + (mma_k_stage) * 2048; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136314896;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (score_col))), "r"(0)); + int _mma_a_lo_1 = (((smem_q_full_addr + 8192) >> 4) & 0x3FFF) + (0) * 2048; + int _mma_b_lo_1 = (((smem_k_full_addr + 8192) >> 4) & 0x3FFF) + (mma_k_stage) * 2048; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136314896;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + int _mma_a_lo_2 = (((smem_q_full_addr + 16384) >> 4) & 0x3FFF) + (0) * 2048; + int _mma_b_lo_2 = (((smem_k_full_addr + 16384) >> 4) & 0x3FFF) + (mma_k_stage) * 2048; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136314896;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + int _mma_a_lo_3 = (((smem_q_full_addr + 24576) >> 4) & 0x3FFF) + (0) * 2048; + int _mma_b_lo_3 = (((smem_k_full_addr + 24576) >> 4) & 0x3FFF) + (mma_k_stage) * 2048; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136314896;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + if (cta_rank == 0) { + elect_commit_cg2_multicast(s_full_addr + (phase_2) * 8, (uint16_t)(3)); + } + if (cta_rank == 0) { + elect_commit_cg2_multicast(k_empty_addr + (mma_k_stage) * 8, (uint16_t)(3)); + } + } + mma_k_stage += 1; + if (mma_k_stage == 3) { mma_k_stage = 0; _phase_k_full ^= 1; } + } + if (tile_3 > 0) { + int prev_pipeline_tile = pipeline_tile_2 - 1; + int prev_phase = prev_pipeline_tile & 1; + int pv_wait_phase = prev_pipeline_tile >> 1 & 1; + mbarrier_wait(p_full_addr + (prev_phase) * 8, pv_wait_phase); + mbarrier_wait(corr_done_addr + (prev_phase) * 8, pv_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + { + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cluster.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(pv_pair_ready_addr + (unsigned int)(prev_phase * 8)), "r"(0) : "memory"); + if (cta_rank == 0) { + mbarrier_wait_cluster_hint(pv_pair_ready_addr + (prev_phase) * 8, pv_wait_phase, 10000000); + } + } + int p_col = ((prev_phase != 0) ? 224 : 96); + { + if (cta_rank == 0) { + mbarrier_wait(v_full_addr + (mma_v_stage) * 8, _phase_v_full); + int _mma_a_lo_6 = (((smem_p_fp8_addr) >> 4) & 0x3FFF) + (prev_phase) * 512; + int _mma_b_lo_6 = ((((smem_v_full_addr) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138477584;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_6), "r"(_mma_b_lo_6), "r"((tmem_tmem_scratch + (256))), "r"(((first_pv) ? 0 : 1))); + int _mma_a_lo_7 = (((smem_p_fp8_addr) >> 4) & 0x3FFF) + (prev_phase) * 512; + int _mma_b_lo_7 = ((((smem_v_full_addr + 16384) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138477584;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_7), "r"(_mma_b_lo_7), "r"((tmem_tmem_scratch + (384))), "r"(((first_pv) ? 0 : 1))); + if (cta_rank == 0) { + elect_commit_cg2_multicast(v_empty_addr + (mma_v_stage) * 8, (uint16_t)(3)); + } + } + mma_v_stage += 1; + if (mma_v_stage == 2) { mma_v_stage = 0; _phase_v_full ^= 1; } + } + first_pv = 0; + if (cta_rank == 0) { + elect_commit_cg2_multicast(pv_done_addr, (uint16_t)(3)); + } + } + } + int last_pipeline_tile = mma_tile_cursor + num_kv_tiles_3 - 1; + int last_phase = last_pipeline_tile & 1; + int drain_wait_phase = last_pipeline_tile >> 1 & 1; + mbarrier_wait(p_full_addr + (last_phase) * 8, drain_wait_phase); + mbarrier_wait(corr_done_addr + (last_phase) * 8, drain_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + { + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cluster.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(pv_pair_ready_addr + (unsigned int)(last_phase * 8)), "r"(0) : "memory"); + if (cta_rank == 0) { + mbarrier_wait_cluster_hint(pv_pair_ready_addr + (last_phase) * 8, drain_wait_phase, 10000000); + } + } + int p_col_last = ((last_phase != 0) ? 224 : 96); + { + if (cta_rank == 0) { + mbarrier_wait(v_full_addr + (mma_v_stage) * 8, _phase_v_full); + int _mma_a_lo_10 = (((smem_p_fp8_addr) >> 4) & 0x3FFF) + (last_phase) * 512; + int _mma_b_lo_10 = ((((smem_v_full_addr) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138477584;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_10), "r"(_mma_b_lo_10), "r"((tmem_tmem_scratch + (256))), "r"(((first_pv) ? 0 : 1))); + int _mma_a_lo_11 = (((smem_p_fp8_addr) >> 4) & 0x3FFF) + (last_phase) * 512; + int _mma_b_lo_11 = ((((smem_v_full_addr + 16384) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138477584;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_11), "r"(_mma_b_lo_11), "r"((tmem_tmem_scratch + (384))), "r"(((first_pv) ? 0 : 1))); + if (cta_rank == 0) { + elect_commit_cg2_multicast(v_empty_addr + (mma_v_stage) * 8, (uint16_t)(3)); + } + } + mma_v_stage += 1; + if (mma_v_stage == 2) { mma_v_stage = 0; _phase_v_full ^= 1; } + } + if (cta_rank == 0) { + elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + } + if (cta_rank == 0) { + elect_commit_cg2_multicast(o_done_addr, (uint16_t)(3)); + } + mma_tile_cursor = mma_tile_cursor + num_kv_tiles_3; + } + mbarrier_arrive(tmem_dealloc_addr); + unsigned int _phase_tmem_dealloc_0 = 0; + mbarrier_wait(tmem_dealloc_addr, _phase_tmem_dealloc_0); + _phase_tmem_dealloc_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + { + int tmem_dealloc_peer_rank = bid % 2 ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cluster.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_dealloc_peer_addr), "r"(tmem_dealloc_peer_rank) : "memory"); + mbarrier_wait(tmem_dealloc_peer_addr, 0); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + } + } + // ---- Role: empty1 ---- + if (warp >= 10 && warp <= 11) { + { // empty1_main + const int wg2_dummy_2 = 0; + unsigned int _phase_q_full_0_4 = 0; + } + } + + // Cleanup +} + +} // extern "C" diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_714df747e199d0adbe08_binding.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_714df747e199d0adbe08_binding.cu new file mode 100644 index 00000000000..d7c15323daa --- /dev/null +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_714df747e199d0adbe08_binding.cu @@ -0,0 +1,606 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +// Generated by the Cake source exporter — do not edit. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_714df747e199d0adbe08'. +#include +#include +#include +#include + +#include "tvm_ffi_utils.h" + +#include +#include +#include +#include +#include +#include +#include + +extern "C" __global__ void kernel_cake_dsv4_714df747e199d0adbe08(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, const __grid_constant__ CUtensorMap O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); + + +namespace cake_host_shim_dab7193014e792b3 { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t CakeDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +// 4D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + int64_t d3 = t.size(t.ndim() - 3); + TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(d3)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 64u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_swa_k' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_swa_k(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_swa_k' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_swa_k' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_swa_k' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_swa_k"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_swa_k' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_swa_k' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_k'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_swa_k') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_compressed_k' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_compressed_k(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_compressed_k' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_compressed_k' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_compressed_k' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_compressed_k"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_compressed_k' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_compressed_k' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_k'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_k') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_swa_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_swa_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_swa_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_swa_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_swa_v' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_swa_v"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_swa_v' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_swa_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_v'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_swa_v') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_compressed_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_compressed_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_compressed_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_compressed_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_compressed_v' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_compressed_v"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_compressed_v' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_compressed_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_v'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_v') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'O' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_O(const TensorView& t, int64_t tp_num_heads, int64_t tp_num_query_tokens) { + TVM_FFI_CHECK(t.ndim() >= 1, ValueError) + << "TMA source 'O' must have at least 1 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'O' must have unit innermost stride, got " << t.stride(-1); + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(8), (uint64_t)(5), (uint64_t)((tp_num_query_tokens * tp_num_heads))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'O' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1] && 1u <= global_dim[2] && 32u <= global_dim[3], ValueError) + << "TMA box (64, 1, 1, 32) exceeds resolved global dims for 'O'"; + int64_t carrier_stride_0 = 64; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 512; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = 2560; + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 1u, 1u, 32u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'O') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_compressed_k, TensorView arg_tmap_swa_v, TensorView arg_tmap_compressed_v, TensorView arg_O, TensorView arg_partial_lse, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_num_query_tokens, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckContiguous(arg_tmap_q, "tmap_q"); + CheckCudaTensor(arg_tmap_swa_k, "tmap_swa_k"); + CheckDtype(arg_tmap_swa_k, "tmap_swa_k", 4, 16, 1); + CheckCudaTensor(arg_tmap_compressed_k, "tmap_compressed_k"); + CheckDtype(arg_tmap_compressed_k, "tmap_compressed_k", 4, 16, 1); + CheckCudaTensor(arg_tmap_swa_v, "tmap_swa_v"); + CheckDtype(arg_tmap_swa_v, "tmap_swa_v", 4, 16, 1); + CheckCudaTensor(arg_tmap_compressed_v, "tmap_compressed_v"); + CheckDtype(arg_tmap_compressed_v, "tmap_compressed_v", 4, 16, 1); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_partial_lse, "partial_lse"); + CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); + CheckContiguous(arg_partial_lse, "partial_lse"); + CheckCudaTensor(arg_swa_indices, "swa_indices"); + CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); + CheckContiguous(arg_swa_indices, "swa_indices"); + CheckCudaTensor(arg_compressed_indices, "compressed_indices"); + CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); + CheckContiguous(arg_compressed_indices, "compressed_indices"); + CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); + CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); + CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); + CheckContiguous(arg_bmm1_scale, "bmm1_scale"); + CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); + CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); + CheckContiguous(arg_bmm2_scale, "bmm2_scale"); + TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) + << "scalar 'num_heads' value " << arg_num_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_query_tokens >= -2147483648LL && arg_num_query_tokens <= 2147483647LL, ValueError) + << "scalar 'num_query_tokens' value " << arg_num_query_tokens + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) + << "scalar 'swa_index_stride' value " << arg_swa_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) + << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) + << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) + << "scalar 'sparse_topk' value " << arg_sparse_topk + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) + << "scalar 'has_sinks' value " << arg_has_sinks + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_swa_k, arg_tmap_q, "tmap_swa_k", "tmap_q"); + CheckSameCudaDevice(arg_tmap_compressed_k, arg_tmap_q, "tmap_compressed_k", "tmap_q"); + CheckSameCudaDevice(arg_tmap_swa_v, arg_tmap_q, "tmap_swa_v", "tmap_q"); + CheckSameCudaDevice(arg_tmap_compressed_v, arg_tmap_q, "tmap_compressed_v", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_partial_lse, arg_tmap_q, "partial_lse", "tmap_q"); + CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); + CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); + CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); + CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") must be divisible by cluster dims (2, 1, 1)"; + + cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_swa_k = EncodeTma_tmap_swa_k(arg_tmap_swa_k); + CUtensorMap p_tmap_compressed_k = EncodeTma_tmap_compressed_k(arg_tmap_compressed_k); + CUtensorMap p_tmap_swa_v = EncodeTma_tmap_swa_v(arg_tmap_swa_v); + CUtensorMap p_tmap_compressed_v = EncodeTma_tmap_compressed_v(arg_tmap_compressed_v); + CUtensorMap p_O = EncodeTma_O(arg_O, (int64_t)arg_num_heads, (int64_t)arg_num_query_tokens); + float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); + int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); + int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); + int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); + float* p_sinks = static_cast(arg_sinks.data_ptr()); + float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); + float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); + int32_t v_num_heads = (int32_t)arg_num_heads; + int32_t v_num_query_tokens = (int32_t)arg_num_query_tokens; + int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; + int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; + int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; + int32_t v_sparse_topk = (int32_t)arg_sparse_topk; + int32_t v_has_sinks = (int32_t)arg_has_sinks; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + void* kargs[] = {&p_tmap_q, &p_tmap_swa_k, &p_tmap_compressed_k, &p_tmap_swa_v, &p_tmap_compressed_v, &p_O, &p_partial_lse, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_num_query_tokens, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_has_sinks, &v_total_work_items}; + + static std::mutex smem_mu; + static auto* smem_status_by_device = + new std::unordered_map(); + cudaError_t smem_status = cudaSuccess; + { + std::lock_guard lock(smem_mu); + auto it = smem_status_by_device->find(dev.device_id); + if (it == smem_status_by_device->end()) { + int current_device = -1; + smem_status = cudaGetDevice(¤t_device); + if (smem_status == cudaSuccess && current_device != dev.device_id) { + smem_status = cudaErrorInvalidDevice; + } + if (smem_status == cudaSuccess) { + smem_status = cudaFuncSetAttribute( + reinterpret_cast(kernel_cake_dsv4_714df747e199d0adbe08), + cudaFuncAttributeMaxDynamicSharedMemorySize, + 231808); + } + smem_status_by_device->emplace(dev.device_id, smem_status); + } else { + smem_status = it->second; + } + } + TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) + << "cudaFuncSetAttribute for kernel_cake_dsv4_714df747e199d0adbe08 failed: " + << cudaGetErrorString(smem_status); + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(512u, 1u, 1u); + + // Extended direct-source launch for a linked CUDA kernel symbol. + cudaLaunchAttribute attrs[2]{}; + int n = 0; + attrs[n].id = cudaLaunchAttributeClusterDimension; + attrs[n].val.clusterDim.x = 2u; + attrs[n].val.clusterDim.y = 1u; + attrs[n].val.clusterDim.z = 1u; + ++n; + attrs[n].id = cudaLaunchAttributeClusterSchedulingPolicyPreference; + attrs[n].val.clusterSchedulingPolicyPreference = cudaClusterSchedulingPolicySpread; + ++n; + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes = 231808u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; + cudaError_t launch_status = cudaLaunchKernelExC( + &config, reinterpret_cast(kernel_cake_dsv4_714df747e199d0adbe08), kargs); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernelExC for kernel_cake_dsv4_714df747e199d0adbe08 failed: " + << cudaGetErrorString(launch_status); +} + +} // namespace cake_host_shim_dab7193014e792b3 + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_dab7193014e792b3::Run); diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_714df747e199d0adbe08_kernel.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_714df747e199d0adbe08_kernel.cu new file mode 100644 index 00000000000..cec2500ed5f --- /dev/null +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_714df747e199d0adbe08_kernel.cu @@ -0,0 +1,2378 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define CAKE_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_SCRATCH_OFFSET 0 +#define NUM_K_PIPE_STAGES 4 +#define NUM_V_PIPE_STAGES 2 +#define SMEM_SMEM_Q_OFF 0 +#define SMEM_SMEM_Q_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_STRIDE 8192 +#define SMEM_SMEM_K_OFF 65536 +#define SMEM_SMEM_K_STAGE_BYTES 16384 +#define SMEM_SMEM_K_STRIDE 16384 +#define SMEM_SMEM_V_LO_OFF 131072 +#define SMEM_SMEM_V_LO_STAGE_BYTES 16384 +#define SMEM_SMEM_V_LO_STRIDE 32768 +#define SMEM_SMEM_V_HI_OFF 147456 +#define SMEM_SMEM_V_HI_STAGE_BYTES 16384 +#define SMEM_SMEM_V_HI_STRIDE 32768 +#define SMEM_SMEM_P_OFF 196608 +#define SMEM_SMEM_P_STAGE_BYTES 8192 +#define SMEM_SMEM_P_STRIDE 8192 +#define SMEM_SMEM_INDICES_OFF 229376 +#define SMEM_SMEM_INDICES_STAGE_BYTES 1024 +#define SMEM_SMEM_INDICES_STRIDE 1024 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_OFF 230400 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_SOFTMAX_EXCHANGE_STRIDE 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_OFF 230912 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_EPILOGUE_EXCHANGE_STRIDE 512 +#define SMEM_SMEM_INDEX_VALID_OFF 231424 +#define SMEM_SMEM_INDEX_VALID_STAGE_BYTES 64 +#define SMEM_SMEM_INDEX_VALID_STRIDE 64 +#define SMEM_TOTAL 231808 +#define THREADS 512 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count) : "memory"); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16_cta2( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" + "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" + "}\n" + :: "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x4(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x4.b32" + " {%0, %1, %2, %3}, [%4];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); +} + + +__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ float2 ex2_emulation_f32x2_value(float2 value) { + const float c0 = 1.0f, c1 = 0.695146143436431884765625f; + const float c2 = 0.227564394474029541015625f, c3 = 0.077119089663028717041015625f; + const float magic = 12582912.0f; + float x0 = max_noftz(value.x, -127.0f), x1 = max_noftz(value.y, -127.0f); + float2 xc2 = make_float2(x0, x1), magic2 = make_float2(magic, magic); + float2 xr2; + asm("add.rm.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xr2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&magic2)); + float2 c3_2 = make_float2(c3, c3), c2_2 = make_float2(c2, c2); + float2 c1_2 = make_float2(c1, c1), c0_2 = make_float2(c0, c0); + float2 xrb2, xfrac2; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xrb2) + : "l"(*(unsigned long long*)&xr2), "l"(*(unsigned long long*)&magic2)); + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xfrac2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&xrb2)); + float2 poly2; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&c3_2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c2_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c1_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c0_2)); + int x0r_i, x1r_i, p0_i, p1_i; + asm("mov.b64 {%0, %1}, %2;" : "=r"(x0r_i), "=r"(x1r_i) : "l"(*(unsigned long long*)&xr2)); + asm("mov.b64 {%0, %1}, %2;" : "=r"(p0_i), "=r"(p1_i) : "l"(*(unsigned long long*)&poly2)); + float r0, r1; + asm("mov.b32 %0, %1;" : "=f"(r0) : "r"((x0r_i << 23) + p0_i)); + asm("mov.b32 %0, %1;" : "=f"(r1) : "r"((x1r_i << 23) + p1_i)); + return make_float2(r0, r1); +} + +__device__ __forceinline__ void ex2_emulation_f32x2(float* x0_ptr, float* x1_ptr) { + float2 result = ex2_emulation_f32x2_value(make_float2(*x0_ptr, *x1_ptr)); + *x0_ptr = result.x; *x1_ptr = result.y; +} + +__device__ __forceinline__ void softmax_frag_exp2_cast( + float* sv, uint32_t* pv, int use_emu) +{ + #pragma unroll + for (int j = 0; j < 16; j++) { + if (use_emu && j >= 12) + ex2_emulation_f32x2(&sv[j*2], &sv[j*2+1]); + else { + sv[j*2] = approx_exp2(sv[j*2]); + sv[j*2+1] = approx_exp2(sv[j*2+1]); + } + } + #pragma unroll + for (int j = 0; j < 16; j++) { + __nv_bfloat162 bf = __float22bfloat162_rn({sv[j*2], sv[j*2+1]}); + pv[j] = reinterpret_cast(bf); + } +} + + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_2d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3}], [%4];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr) { + // Canonical .shared::cta form for non-multicast gather4, matching + // trtllm-gen / cuda_ptx and the PTX ISA qualifier order + // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, + // .shared::cluster is reserved for the multicast variant (ctaMask). + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem_cta2( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr) { + // Canonical .shared::cta form; see tma_gather4_gmem2smem above. + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem_mc( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr, unsigned short cta_mask) { + // Multicast variant: the PTX grammar ties the .shared::cluster + // destination to .multicast::cluster + ctaMask (cf. cuda_ptx / + // SM100_TMA_LOAD_MULTICAST_2D_GATHER4 in CUTLASS). + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" + ".mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem_mc_cta2( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr, unsigned short cta_mask) { + // Multicast + cta_group::2 variant; see tma_gather4_gmem2smem_mc. + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" + ".mbarrier::complete_tx::bytes.multicast::cluster.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void tma_store_4d( + const void *tmap, int x, int y, int z, int w, unsigned smem_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.global.shared::cta.tile.bulk_group" + " [%0, {%1, %2, %3, %4}], [%5];" + :: "l"(tmap), "r"(x), "r"(y), "r"(z), "r"(w), "r"(smem_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %1;\n\t" + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" + "}\n" + :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); +} + +extern "C" { + +__global__ __launch_bounds__(512, 1) __cluster_dims__(2,1,1) void +kernel_cake_dsv4_714df747e199d0adbe08(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, const __grid_constant__ CUtensorMap O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); + + const int mbar_base = smem + 231488; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define k_full_addr (mbar_base + 16) + #define k_empty_addr (mbar_base + 48) + #define v_full_addr (mbar_base + 80) + #define v_empty_addr (mbar_base + 96) + #define s_full_addr (mbar_base + 112) + #define s_empty_addr (mbar_base + 128) + #define p_full_addr (mbar_base + 144) + #define p_empty_addr (mbar_base + 160) + #define o_empty_addr (mbar_base + 176) + #define stats_addr (mbar_base + 184) + #define stats_empty_addr (mbar_base + 200) + #define o_full_addr (mbar_base + 216) + #define epi_q_free_addr (mbar_base + 224) + #define o_first_slice_seeded_addr (mbar_base + 232) + #define tmem_dealloc_peer_addr (mbar_base + 240) + #define index_valid_full_addr (mbar_base + 248) + #define index_valid_empty_addr (mbar_base + 264) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 2; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + + // Kernel setup ops + __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 0); + const int smem_q_addr = smem + 0; + __nv_bfloat16* smem_k = reinterpret_cast<__nv_bfloat16*>(smem_raw + 65536); + const int smem_k_addr = smem + 65536; + __nv_bfloat16* smem_v_lo = reinterpret_cast<__nv_bfloat16*>(smem_raw + 131072); + const int smem_v_lo_addr = smem + 131072; + __nv_bfloat16* smem_v_hi = reinterpret_cast<__nv_bfloat16*>(smem_raw + 147456); + const int smem_v_hi_addr = smem + 147456; + __nv_bfloat16* smem_p = reinterpret_cast<__nv_bfloat16*>(smem_raw + 196608); + const int smem_p_addr = smem + 196608; + int* smem_indices = reinterpret_cast(smem_raw + 229376); + const int smem_indices_addr = smem + 229376; + float* smem_softmax_exchange = reinterpret_cast(smem_raw + 230400); + const int smem_softmax_exchange_addr = smem + 230400; + float* smem_epilogue_exchange = reinterpret_cast(smem_raw + 230912); + const int smem_epilogue_exchange_addr = smem + 230912; + unsigned int* smem_index_valid = reinterpret_cast(smem_raw + 231424); + const int smem_index_valid_addr = smem + 231424; + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // Mbarrier init (19 pipeline groups, 0 ordered-sequence groups, 35 barriers) + // Mbarriers at smem_raw[231488..231768) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // q_full: 1 barriers, init_count=1 + mbarrier_init(smem + 231488, 1); + // q_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 231496, 1); + // --- pipeline 'k_pipe' --- + // k_full: 4 barriers, init_count=1 + mbarrier_init(smem + 231504, 1); + mbarrier_init(smem + 231512, 1); + mbarrier_init(smem + 231520, 1); + mbarrier_init(smem + 231528, 1); + // k_empty: 4 barriers, init_count=1 + mbarrier_init(smem + 231536, 1); + mbarrier_init(smem + 231544, 1); + mbarrier_init(smem + 231552, 1); + mbarrier_init(smem + 231560, 1); + // --- pipeline 'v_pipe' --- + // v_full: 2 barriers, init_count=1 + mbarrier_init(smem + 231568, 1); + mbarrier_init(smem + 231576, 1); + // v_empty: 2 barriers, init_count=1 + mbarrier_init(smem + 231584, 1); + mbarrier_init(smem + 231592, 1); + // s_full: 2 barriers, init_count=1 + mbarrier_init(smem + 231600, 1); + mbarrier_init(smem + 231608, 1); + // s_empty: 2 barriers, init_count=256 + mbarrier_init(smem + 231616, 256); + mbarrier_init(smem + 231624, 256); + // p_full: 2 barriers, init_count=256 + mbarrier_init(smem + 231632, 256); + mbarrier_init(smem + 231640, 256); + // p_empty: 2 barriers, init_count=1 + mbarrier_init(smem + 231648, 1); + mbarrier_init(smem + 231656, 1); + // o_empty: 1 barriers, init_count=256 + mbarrier_init(smem + 231664, 256); + // stats: 2 barriers, init_count=128 + mbarrier_init(smem + 231672, 128); + mbarrier_init(smem + 231680, 128); + // stats_empty: 2 barriers, init_count=128 + mbarrier_init(smem + 231688, 128); + mbarrier_init(smem + 231696, 128); + // o_full: 1 barriers, init_count=1 + mbarrier_init(smem + 231704, 1); + // epi_q_free: 1 barriers, init_count=4 + mbarrier_init(smem + 231712, 4); + // o_first_slice_seeded: 1 barriers, init_count=256 + mbarrier_init(smem + 231720, 256); + // tmem_dealloc_peer: 1 barriers, init_count=32 + mbarrier_init(smem + 231728, 32); + // index_valid_full: 2 barriers, init_count=32 + mbarrier_init(smem + 231736, 32); + mbarrier_init(smem + 231744, 32); + // index_valid_empty: 2 barriers, init_count=128 + mbarrier_init(smem + 231752, 128); + mbarrier_init(smem + 231760, 128); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 231768); + if (warp == 0) { + int _tmem_hold = smem + 231768; + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); + } + + __syncthreads(); + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_tmem_scratch = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 15) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 64;"); + } + + // ---- Role: softmax_wg ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); + { // softmax_wg_main + const int softmax_dummy = 0; + float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; + const int head_row_base = warp % 2 * 32; + const int score_row_base = warp % 4 * 32; + const int n_half = warp % 4 / 2; + const int my_row = head_row_base + lane; + const int exchange_idx = n_half * 64 + my_row; + int tile_cursor = 0; + int valid_iter = 0; + float seed_zero[4]; + #pragma unroll + for (int seed_i = 0; seed_i < 4; seed_i++) { + seed_zero[seed_i] = 0.0f; + } + #pragma unroll + for (int seed_stage = 0; seed_stage < 2; seed_stage++) { + #pragma unroll + for (int seed_col = 128 + seed_stage * 128; seed_col < 128 + (seed_stage + 1) * 128; seed_col += 4) { + int seed_addr = taddr + (unsigned int)seed_col + (unsigned int)(score_row_base << 16); + tmem_st_x4_f32(seed_addr, seed_zero); + } + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(o_first_slice_seeded_addr); + int peer_rank = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_first_slice_seeded_addr), "r"(peer_rank) : "memory"); + #pragma unroll 1 + for (unsigned int work_idx = cluster_id; work_idx < total_work_items; work_idx += num_clusters) { + int split_idx = work_idx % (unsigned int)(((1) ? 5 : 1)); + int query_idx = work_idx / (unsigned int)(((1) ? 5 : 1)); + int _max_0 = ((sparse_topk_lens[query_idx] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx] + sparse_topk_lens_offset) : (0)); + int _min_2 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); + int active_topk = _min_2; + int owner_capacity = ((sparse_topk + 128 - 1) / 128 + ((1) ? 5 : 1) - 1) / ((1) ? 5 : 1); + int first_owner_tile = work_idx % (unsigned int)(((1) ? 5 : 1)) * (unsigned int)owner_capacity; + int _min_3 = ((owner_capacity) < ((sparse_topk + 128 - 1) / 128 - first_owner_tile) ? (owner_capacity) : ((sparse_topk + 128 - 1) / 128 - first_owner_tile)); + int num_kv_tiles = _min_3; + int first_kv_tile = first_owner_tile; + int head_idx = cta_rank * 64 + my_row; + float row_max_val = -CAKE_INF; + float row_sum_val = 0.0f; + if (has_sinks != 0 && head_idx < num_heads && split_idx == 0) { + row_max_val = sinks[head_idx] * 1.4426950408889634f / softmax_scale_log2; + if (n_half == 0) { + row_sum_val = 1.0f; + } + } + int valid_slot = valid_iter & 1; + int valid_full_phase = valid_iter >> 1 & 1; + mbarrier_wait(index_valid_full_addr + (valid_slot) * 8, valid_full_phase); + int valid_word_base = valid_slot * 8 + n_half * 2; + #pragma unroll 1 + for (int tile = 0; tile < num_kv_tiles; tile++) { + int pipeline_tile = tile_cursor + tile; + int score_phase = pipeline_tile & 1; + int s_full_phase = pipeline_tile >> 1 & 1; + mbarrier_wait(s_full_addr + (score_phase) * 8, s_full_phase); + int score_offset = ((score_phase != 0) ? 64 : 0); + int score_base = taddr + (unsigned int)score_offset + (unsigned int)(score_row_base << 16); + float _tmem_load_0[64]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]) + : "r"(score_base)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]) + : "r"(score_base + 32)); + int valid_cols = active_topk - (first_kv_tile + tile) * 128 - n_half * 64; + if (valid_cols < 0) { + valid_cols = 0; + } + if (valid_cols > 64) { + valid_cols = 64; + } + unsigned int keep_lo = smem_index_valid[valid_word_base + tile * 4]; + unsigned int keep_hi = smem_index_valid[valid_word_base + tile * 4 + 1]; + uint32_t _thresh_mask_0; + { + int _lim_0 = valid_cols; + if (_lim_0 <= 0) { _thresh_mask_0 = 0u; } + else if (_lim_0 >= 32) { _thresh_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_thresh_mask_0) : "r"(_lim_0)); + } + } + uint32_t _thresh_mask_1; + { + int _lim_1 = valid_cols - 32; + if (_lim_1 <= 0) { _thresh_mask_1 = 0u; } + else if (_lim_1 >= 32) { _thresh_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_thresh_mask_1) : "r"(_lim_1)); + } + } + if (!(_thresh_mask_0 & keep_lo & (1u << 0))) _tmem_load_0[0] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 1))) _tmem_load_0[1] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 2))) _tmem_load_0[2] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 3))) _tmem_load_0[3] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 4))) _tmem_load_0[4] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 5))) _tmem_load_0[5] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 6))) _tmem_load_0[6] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 7))) _tmem_load_0[7] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 8))) _tmem_load_0[8] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 9))) _tmem_load_0[9] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 10))) _tmem_load_0[10] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 11))) _tmem_load_0[11] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 12))) _tmem_load_0[12] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 13))) _tmem_load_0[13] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 14))) _tmem_load_0[14] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 15))) _tmem_load_0[15] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 16))) _tmem_load_0[16] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 17))) _tmem_load_0[17] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 18))) _tmem_load_0[18] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 19))) _tmem_load_0[19] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 20))) _tmem_load_0[20] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 21))) _tmem_load_0[21] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 22))) _tmem_load_0[22] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 23))) _tmem_load_0[23] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 24))) _tmem_load_0[24] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 25))) _tmem_load_0[25] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 26))) _tmem_load_0[26] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 27))) _tmem_load_0[27] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 28))) _tmem_load_0[28] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 29))) _tmem_load_0[29] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 30))) _tmem_load_0[30] = -CAKE_INF; + if (!(_thresh_mask_0 & keep_lo & (1u << 31))) _tmem_load_0[31] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 0))) _tmem_load_0[32] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 1))) _tmem_load_0[33] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 2))) _tmem_load_0[34] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 3))) _tmem_load_0[35] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 4))) _tmem_load_0[36] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 5))) _tmem_load_0[37] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 6))) _tmem_load_0[38] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 7))) _tmem_load_0[39] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 8))) _tmem_load_0[40] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 9))) _tmem_load_0[41] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 10))) _tmem_load_0[42] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 11))) _tmem_load_0[43] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 12))) _tmem_load_0[44] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 13))) _tmem_load_0[45] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 14))) _tmem_load_0[46] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 15))) _tmem_load_0[47] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 16))) _tmem_load_0[48] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 17))) _tmem_load_0[49] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 18))) _tmem_load_0[50] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 19))) _tmem_load_0[51] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 20))) _tmem_load_0[52] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 21))) _tmem_load_0[53] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 22))) _tmem_load_0[54] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 23))) _tmem_load_0[55] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 24))) _tmem_load_0[56] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 25))) _tmem_load_0[57] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 26))) _tmem_load_0[58] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 27))) _tmem_load_0[59] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 28))) _tmem_load_0[60] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 29))) _tmem_load_0[61] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 30))) _tmem_load_0[62] = -CAKE_INF; + if (!(_thresh_mask_1 & keep_hi & (1u << 31))) _tmem_load_0[63] = -CAKE_INF; + float2 _reg_reduce_max2_2 = {-CAKE_INF, -CAKE_INF}; + row_max_x32_accum(&_tmem_load_0[0], _reg_reduce_max2_2); + row_max_x32_accum(&_tmem_load_0[32], _reg_reduce_max2_2); + float _tmem_load_0_max = row_max_reduce(_reg_reduce_max2_2); + float _max_1 = max_noftz(_tmem_load_0_max, row_max_val); + float new_max = _max_1; + smem_softmax_exchange[exchange_idx] = new_max; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + float _max_2 = max_noftz(new_max, smem_softmax_exchange[exchange_idx ^ 64]); + new_max = _max_2; + asm volatile("barrier.sync 3, 128;" ::: "memory"); + float no_correction = (((new_max - row_max_val) * softmax_scale_log2 <= 0.0f) ? 1.0f : 0.0f); + float delta = softmax_scale_log2 * (row_max_val - new_max); + float _exp2_0 = approx_exp2(delta); + float exp_delta = _exp2_0; + float acc_scale = ((row_max_val > -CAKE_INF) ? exp_delta : 1.0f); + row_max_val = new_max; + float safe_max = ((new_max == -CAKE_INF) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + const float2 _fma_b2_3 = {softmax_scale_log2, softmax_scale_log2}; + const float2 _fma_c2_4 = {-max_scaled, -max_scaled}; + #pragma unroll + for (int _lf = 0; _lf < 32; _lf++) + fma_f32x2_inplace(&reinterpret_cast(_tmem_load_0)[_lf], _fma_b2_3, _fma_c2_4); + #pragma unroll + for (int _le = 0; _le < 64; _le++) { + _tmem_load_0[_le] = approx_exp2(_tmem_load_0[_le]); + } + float2 _reg_reduce_sum2_5 = make_float2(0.0f, 0.0f); + softmax_block_sum(&_tmem_load_0[0], &_reg_reduce_sum2_5); + softmax_block_sum(&_tmem_load_0[32], &_reg_reduce_sum2_5); + float _tmem_load_0_sum = _reg_reduce_sum2_5.x + _reg_reduce_sum2_5.y; + row_sum_val = row_sum_val * acc_scale + _tmem_load_0_sum; + int stats_empty_phase = pipeline_tile >> 1 & 1 ^ 1; + mbarrier_wait(stats_empty_addr + (score_phase) * 8, stats_empty_phase); + float meta[4]; + meta[0] = row_sum_val; + meta[1] = row_max_val; + meta[2] = acc_scale; + meta[3] = no_correction; + int meta_addr = taddr + 384 + (unsigned int)(score_phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(head_row_base << 16); + tmem_st_x4_f32(meta_addr, meta); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + mbarrier_arrive(stats_addr + (score_phase) * 8); + int p_stage = pipeline_tile % 2; + int p_empty_phase = pipeline_tile / 2 & 1 ^ 1; + mbarrier_wait(p_empty_addr + (p_stage) * 8, p_empty_phase); + uint32_t _tmem_load_0_bf16[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_0[_lp*2 + 0], _tmem_load_0[_lp*2+1 + 0])); + _tmem_load_0_bf16[_lp] = *(uint32_t*)&_bf2; + } + int smem_p_stage = p_stage * 2 + n_half; + int p_base_smem = smem_p_addr + (unsigned int)(smem_p_stage * 8192); + #pragma unroll + for (int vec = 0; vec < 8; vec++) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_base_smem + (my_row * 128 + vec * 16 ^ (my_row * 128 + vec * 16 >> 7 & 7) << 4))), "r"(_tmem_load_0_bf16[vec * 4]), "r"(_tmem_load_0_bf16[vec * 4 + 1]), "r"(_tmem_load_0_bf16[vec * 4 + 2]), "r"(_tmem_load_0_bf16[vec * 4 + 3]) : "memory"); + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(p_full_addr + (unsigned int)(p_stage * 8)), "r"(0) : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(s_empty_addr + (unsigned int)(score_phase * 8)), "r"(0) : "memory"); + } + mbarrier_arrive(index_valid_empty_addr + (valid_slot) * 8); + valid_iter = valid_iter + 1; + tile_cursor = tile_cursor + num_kv_tiles; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + } + } + // ---- Role: correction_wg ---- + if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); + { // correction_wg_main + const int correction_dummy = 0; + float softmax_scale_log2_1 = bmm1_scale[0] * 1.4426950408889634f; + float output_scale = bmm2_scale[0]; + const int head_row_base_1 = warp % 2 * 32; + const int output_row_base = warp % 4 * 32; + const int n_half_1 = warp % 4 / 2; + const int my_row_1 = head_row_base_1 + lane; + const int corr_row = head_row_base_1 << 16; + int tile_cursor_1 = 0; + unsigned int _phase_o_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = cluster_id; work_idx_1 < total_work_items; work_idx_1 += num_clusters) { + int split_idx_1 = work_idx_1 % (unsigned int)(((1) ? 5 : 1)); + int query_idx_1 = work_idx_1 / (unsigned int)(((1) ? 5 : 1)); + int owner_capacity_1 = ((sparse_topk + 128 - 1) / 128 + ((1) ? 5 : 1) - 1) / ((1) ? 5 : 1); + int first_owner_tile_1 = work_idx_1 % (unsigned int)(((1) ? 5 : 1)) * (unsigned int)owner_capacity_1; + int _min_4 = ((owner_capacity_1) < ((sparse_topk + 128 - 1) / 128 - first_owner_tile_1) ? (owner_capacity_1) : ((sparse_topk + 128 - 1) / 128 - first_owner_tile_1)); + int num_kv_tiles_1 = _min_4; + float final_sum = 0.0f; + float final_max = -CAKE_INF; + #pragma unroll 1 + for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { + int pipeline_tile_1 = tile_cursor_1 + tile_1; + int score_phase_1 = pipeline_tile_1 & 1; + int stats_phase = pipeline_tile_1 >> 1 & 1; + mbarrier_wait(stats_addr + (score_phase_1) * 8, stats_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int meta_addr_1 = taddr + 384 + (unsigned int)(score_phase_1 * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row; + float _tmem_load_1[4]; + tmem_ld_x4(&_tmem_load_1[0], meta_addr_1); + final_sum = _tmem_load_1[0]; + final_max = _tmem_load_1[1]; + float acc_scale_1 = _tmem_load_1[2]; + float no_correction_1 = _tmem_load_1[3]; + if (tile_1 > 0) { + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _vote_0 = __all_sync(0xFFFFFFFF, no_correction_1 == 1.0f); + int skip_correction = _vote_0; + if (skip_correction == 0) { + #pragma unroll + for (int local_slice = 0; local_slice < 2; local_slice++) { + #pragma unroll + for (int acc_stage = 0; acc_stage < 2; acc_stage++) { + int o_base = taddr + 128 + (unsigned int)(acc_stage * 128) + (unsigned int)(local_slice * 64) + (unsigned int)(output_row_base << 16); + #pragma unroll + for (int c = 0; c < 64; c += 16) { + float _tmem_load_2[16]; + tmem_ld_x16(&_tmem_load_2[0], o_base + c); + const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + mul_f32x2_inplace(&reinterpret_cast(_tmem_load_2)[_ls], _scale2_0); + tmem_st_x16_f32(o_base + c, _tmem_load_2); + } + } + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + } + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + } + mbarrier_arrive(stats_empty_addr + (score_phase_1) * 8); + } + mbarrier_wait(o_full_addr, _phase_o_full_0); + _phase_o_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + const int exchange_idx_1 = n_half_1 * 64 + my_row_1; + smem_epilogue_exchange[exchange_idx_1] = final_sum; + asm volatile("barrier.sync 4, 128;" ::: "memory"); + final_sum = final_sum + smem_epilogue_exchange[exchange_idx_1 ^ 64]; + asm volatile("barrier.sync 4, 128;" ::: "memory"); + float _rcp_0 = approx_rcp(final_sum); + float inv_sum = ((final_sum > 0.0f) ? _rcp_0 : 0.0f); + int head_idx_1 = cta_rank * 64 + my_row_1; + int output_offset = ((1) ? ((query_idx_1 * num_heads + head_idx_1) * ((1) ? 5 : 1) + split_idx_1) * 512 : (query_idx_1 * num_heads + head_idx_1) * 512); + if (n_half_1 == 0 && head_idx_1 < num_heads) { + int stat_offset = (query_idx_1 * num_heads + head_idx_1) * ((1) ? 5 : 1) + split_idx_1; + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(final_sum)); + partial_lse[stat_offset] = ((final_sum > 0.0f) ? final_max * softmax_scale_log2_1 + _log2_0 : -CAKE_INF); + } + float epi_scale = inv_sum * output_scale; + int epi_box_base = smem_q_addr + (unsigned int)(warp % 4 * 16384); + int epi_row0 = query_idx_1 * num_heads + cta_rank * 64 + head_row_base_1; + #pragma unroll + for (int acc_stage_1 = 0; acc_stage_1 < 2; acc_stage_1++) { + #pragma unroll + for (int local_slice_1 = 0; local_slice_1 < 2; local_slice_1++) { + int logical_slice = acc_stage_1 * 2 * 2 + n_half_1 * 2 + local_slice_1; + int epi_box = epi_box_base + (acc_stage_1 * 2 + local_slice_1) * 4096; + int o_base_1 = taddr + 128 + (unsigned int)(acc_stage_1 * 128) + (unsigned int)(local_slice_1 * 64) + (unsigned int)(output_row_base << 16); + #pragma unroll + for (int c_1 = 0; c_1 < 64; c_1 += 32) { + float _tmem_load_3[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]) + : "r"(o_base_1 + c_1)); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + float epi_scaled[32]; + #pragma unroll + for (int e = 0; e < 32; e++) { + epi_scaled[e] = _tmem_load_3[e] * epi_scale; + } + unsigned int epi_packed[16]; + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(epi_scaled[_lp*2 + 0], epi_scaled[_lp*2+1 + 0])); + epi_packed[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((epi_box + (lane * 128 + c_1 * 2 ^ (lane * 128 + c_1 * 2 >> 7 & 7) << 4))), "r"(epi_packed[0]), "r"(epi_packed[1]), "r"(epi_packed[2]), "r"(epi_packed[3]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((epi_box + (lane * 128 + (c_1 * 2 + 16) ^ (lane * 128 + (c_1 * 2 + 16) >> 7 & 7) << 4))), "r"(epi_packed[4]), "r"(epi_packed[5]), "r"(epi_packed[6]), "r"(epi_packed[7]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((epi_box + (lane * 128 + (c_1 * 2 + 32) ^ (lane * 128 + (c_1 * 2 + 32) >> 7 & 7) << 4))), "r"(epi_packed[8]), "r"(epi_packed[9]), "r"(epi_packed[10]), "r"(epi_packed[11]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((epi_box + (lane * 128 + (c_1 * 2 + 48) ^ (lane * 128 + (c_1 * 2 + 48) >> 7 & 7) << 4))), "r"(epi_packed[12]), "r"(epi_packed[13]), "r"(epi_packed[14]), "r"(epi_packed[15]) : "memory"); + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + __syncwarp(); + if (elect_sync()) { + tma_store_4d((&O), 0, logical_slice, split_idx_1, epi_row0, epi_box); + asm volatile("cp.async.bulk.commit_group;"); + } + } + } + if (elect_sync()) { + asm volatile("cp.async.bulk.wait_group.read 0;"); + mbarrier_arrive(epi_q_free_addr); + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(o_empty_addr), "r"(0) : "memory"); + tile_cursor_1 = tile_cursor_1 + num_kv_tiles_1; + } + if (elect_sync()) { + asm volatile("cp.async.bulk.wait_group 0;"); + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + } + } + // ---- Role: mma_warp ---- + if (warp == 8) { + { // mma_warp_main + const int mma_dummy = 0; + unsigned int k_stage = 0; + unsigned int v_stage = 0; + int tile_cursor_2 = 0; + unsigned int _phase_o_first_slice_seeded_0 = 0; + mbarrier_wait(o_first_slice_seeded_addr, _phase_o_first_slice_seeded_0); + _phase_o_first_slice_seeded_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + unsigned int _phase_q_full_0 = 0; + unsigned int _phase_k_full = 0; + unsigned int _phase_o_empty_0 = 1; + unsigned int _phase_v_full = 0; + #pragma unroll 1 + for (unsigned int work_idx_2 = cluster_id; work_idx_2 < total_work_items; work_idx_2 += num_clusters) { + int owner_capacity_2 = ((sparse_topk + 128 - 1) / 128 + ((1) ? 5 : 1) - 1) / ((1) ? 5 : 1); + int first_owner_tile_2 = work_idx_2 % (unsigned int)(((1) ? 5 : 1)) * (unsigned int)owner_capacity_2; + int _min_1 = ((owner_capacity_2) < ((sparse_topk + 128 - 1) / 128 - first_owner_tile_2) ? (owner_capacity_2) : ((sparse_topk + 128 - 1) / 128 - first_owner_tile_2)); + int num_kv_tiles_2 = _min_1; + if (cta_rank == 0) { + mbarrier_wait(q_full_addr, _phase_q_full_0); + _phase_q_full_0 ^= 1; + int first_pv = 1; + #pragma unroll 1 + for (int tile_2 = 0; tile_2 < num_kv_tiles_2; tile_2++) { + int pipeline_tile_2 = tile_cursor_2 + tile_2; + int score_phase_2 = pipeline_tile_2 & 1; + int s_empty_phase = pipeline_tile_2 >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (score_phase_2) * 8, s_empty_phase); + int score_col = ((score_phase_2 != 0) ? 64 : 0); + #pragma unroll 1 + for (int qk_stage = 0; qk_stage < 4; qk_stage++) { + mbarrier_wait(k_full_addr + (k_stage) * 8, _phase_k_full); + int _mma_a_lo_0 = (((smem_q_addr) >> 4) & 0x3FFF) + (qk_stage * 2) * 512; + int _mma_b_lo_0 = (((smem_k_addr) >> 4) & 0x3FFF) + (k_stage) * 1024; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (score_col))), "r"(((qk_stage == 0) ? 0 : 1))); + int _mma_a_lo_1 = (((smem_q_addr) >> 4) & 0x3FFF) + (qk_stage * 2 + 1) * 512; + int _mma_b_lo_1 = (((smem_k_addr + 8192) >> 4) & 0x3FFF) + (k_stage) * 1024; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); + elect_commit_cg2_multicast(k_empty_addr + (k_stage) * 8, (uint16_t)(3)); + k_stage += 1; + if (k_stage == 4) { k_stage = 0; _phase_k_full ^= 1; } + } + elect_commit_cg2_multicast(s_full_addr + (score_phase_2) * 8, (uint16_t)(3)); + if (tile_2 > 0) { + int prev_pipeline_tile = pipeline_tile_2 - 1; + int p_stage_1 = prev_pipeline_tile % 2; + int p_full_phase = prev_pipeline_tile / 2 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (p_stage_1) * 8, p_full_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps = 0; ps < 2; ps++) { + int smem_p_stage_1 = p_stage_1 * 2 + ps; + mbarrier_wait(v_full_addr + (v_stage) * 8, _phase_v_full); + #pragma unroll + for (int acc_stage_2 = 0; acc_stage_2 < 2; acc_stage_2++) { + int output_col = 128 + acc_stage_2 * 128; + if (acc_stage_2 == 0) { + int _mma_a_lo_2 = (((smem_p_addr) >> 4) & 0x3FFF) + (smem_p_stage_1) * 512; + int _mma_b_lo_2 = ((((smem_v_lo_addr) >> 4) & 0x3FFF) | 0x2000000) + (v_stage) * 2048; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (output_col))), "r"(((((ps == 0) ? first_pv : 0)) ? 0 : 1))); + } else { + int _mma_a_lo_3 = (((smem_p_addr) >> 4) & 0x3FFF) + (smem_p_stage_1) * 512; + int _mma_b_lo_3 = ((((smem_v_hi_addr) >> 4) & 0x3FFF) | 0x2000000) + (v_stage) * 2048; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"((tmem_tmem_scratch + (output_col))), "r"(((((ps == 0) ? first_pv : 0)) ? 0 : 1))); + } + } + elect_commit_cg2_multicast(v_empty_addr + (v_stage) * 8, (uint16_t)(3)); + v_stage += 1; + if (v_stage == 2) { v_stage = 0; _phase_v_full ^= 1; } + } + first_pv = 0; + elect_commit_cg2_multicast(p_empty_addr + (p_stage_1) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + } + } + int last_pipeline_tile = tile_cursor_2 + num_kv_tiles_2 - 1; + int last_p_stage = last_pipeline_tile % 2; + int last_p_phase = last_pipeline_tile / 2 & 1; + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + mbarrier_wait(p_full_addr + (last_p_stage) * 8, last_p_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int ps_1 = 0; ps_1 < 2; ps_1++) { + int smem_p_stage_2 = last_p_stage * 2 + ps_1; + mbarrier_wait(v_full_addr + (v_stage) * 8, _phase_v_full); + #pragma unroll + for (int acc_stage_3 = 0; acc_stage_3 < 2; acc_stage_3++) { + int output_col_1 = 128 + acc_stage_3 * 128; + if (acc_stage_3 == 0) { + int _mma_a_lo_4 = (((smem_p_addr) >> 4) & 0x3FFF) + (smem_p_stage_2) * 512; + int _mma_b_lo_4 = ((((smem_v_lo_addr) >> 4) & 0x3FFF) | 0x2000000) + (v_stage) * 2048; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"((tmem_tmem_scratch + (output_col_1))), "r"(((((ps_1 == 0) ? first_pv : 0)) ? 0 : 1))); + } else { + int _mma_a_lo_5 = (((smem_p_addr) >> 4) & 0x3FFF) + (smem_p_stage_2) * 512; + int _mma_b_lo_5 = ((((smem_v_hi_addr) >> 4) & 0x3FFF) | 0x2000000) + (v_stage) * 2048; + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_5), "r"(_mma_b_lo_5), "r"((tmem_tmem_scratch + (output_col_1))), "r"(((((ps_1 == 0) ? first_pv : 0)) ? 0 : 1))); + } + } + elect_commit_cg2_multicast(v_empty_addr + (v_stage) * 8, (uint16_t)(3)); + v_stage += 1; + if (v_stage == 2) { v_stage = 0; _phase_v_full ^= 1; } + } + elect_commit_cg2_multicast(p_empty_addr + (last_p_stage) * 8, (uint16_t)(3)); + elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); + elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + } + tile_cursor_2 = tile_cursor_2 + num_kv_tiles_2; + } + if (cta_rank == 0) { + #pragma unroll + for (int tail_offset = 0; tail_offset < 2; tail_offset++) { + int tail_tile = tile_cursor_2 + tail_offset; + int tail_stage = tail_tile & 1; + int tail_phase = tail_tile >> 1 & 1 ^ 1; + mbarrier_wait(s_empty_addr + (tail_stage) * 8, tail_phase); + } + mbarrier_wait(o_empty_addr, _phase_o_empty_0); + _phase_o_empty_0 ^= 1; + } + asm volatile("tcgen05.fence::before_thread_sync;"); + int peer_rank_1 = cta_rank ^ 1; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(tmem_dealloc_peer_addr), "r"(peer_rank_1) : "memory"); + mbarrier_wait(tmem_dealloc_peer_addr, 0); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + } + // ---- Role: load_wg ---- + if (warp >= 9 && warp <= 12) { + { // load_wg_main + const int load_dummy = 0; + const int load_warp_rank = warp - 9; + unsigned int k_stage_1 = 0; + unsigned int v_stage_1 = 0; + unsigned int _phase_q_empty_0 = 1; + unsigned int _phase_epi_q_free_0 = 1; + unsigned int _phase_k_empty = 1; + unsigned int _phase_v_empty = 1; + #pragma unroll 1 + for (unsigned int work_idx_3 = cluster_id; work_idx_3 < total_work_items; work_idx_3 += num_clusters) { + int split_idx_2 = work_idx_3 % (unsigned int)(((1) ? 5 : 1)); + int query_idx_2 = work_idx_3 / (unsigned int)(((1) ? 5 : 1)); + int owner_capacity_3 = ((sparse_topk + 128 - 1) / 128 + ((1) ? 5 : 1) - 1) / ((1) ? 5 : 1); + int first_owner_tile_3 = work_idx_3 % (unsigned int)(((1) ? 5 : 1)) * (unsigned int)owner_capacity_3; + int _min_0 = ((owner_capacity_3) < ((sparse_topk + 128 - 1) / 128 - first_owner_tile_3) ? (owner_capacity_3) : ((sparse_topk + 128 - 1) / 128 - first_owner_tile_3)); + int num_kv_tiles_3 = _min_0; + int first_kv_tile_1 = first_owner_tile_3; + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + mbarrier_wait(epi_q_free_addr, _phase_epi_q_free_0); + _phase_epi_q_free_0 ^= 1; + if (load_warp_rank == 0) { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(131072)) : "memory"); + } + } + if (elect_sync()) { + #pragma unroll + for (int q_stage = 0; q_stage < 8; q_stage++) { + tma_4d_gmem2smem_cta2(smem_q_addr + (unsigned int)(q_stage * 8192), (&tmap_q), 0, cta_rank * 64, q_stage, query_idx_2, ((q_full_addr) & 0xFEFFFFFF)); + } + } + } + asm volatile("barrier.sync 2, 160;" ::: "memory"); + #pragma unroll + for (int qk_stage_1 = 0; qk_stage_1 < 4; qk_stage_1++) { + mbarrier_wait(k_empty_addr + (k_stage_1) * 8, _phase_k_empty); + if (load_warp_rank == 0) { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + } + } + } + int group = load_warp_rank * 4 + lane; + if (lane < 4) { + int index_offset = cta_rank * 64 + group * 4; + int raw_rows[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows[0])), "=r"(*reinterpret_cast(&raw_rows[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows[(0) + 3])) + : "r"(smem_indices_addr + (unsigned int)(index_offset * 4))); + if (first_kv_tile_1 == 0) { + tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group * 512), (&tmap_swa_k), qk_stage_1 * 128, ((raw_rows[0] >= 0) ? raw_rows[0] : 0), ((raw_rows[1] >= 0) ? raw_rows[1] : 0), ((raw_rows[2] >= 0) ? raw_rows[2] : 0), ((raw_rows[3] >= 0) ? raw_rows[3] : 0), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group * 512) + 8192, (&tmap_swa_k), qk_stage_1 * 128 + 64, ((raw_rows[0] >= 0) ? raw_rows[0] : 0), ((raw_rows[1] >= 0) ? raw_rows[1] : 0), ((raw_rows[2] >= 0) ? raw_rows[2] : 0), ((raw_rows[3] >= 0) ? raw_rows[3] : 0), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + } else { + tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group * 512), (&tmap_compressed_k), qk_stage_1 * 128, ((raw_rows[0] >= 0) ? raw_rows[0] : 0), ((raw_rows[1] >= 0) ? raw_rows[1] : 0), ((raw_rows[2] >= 0) ? raw_rows[2] : 0), ((raw_rows[3] >= 0) ? raw_rows[3] : 0), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group * 512) + 8192, (&tmap_compressed_k), qk_stage_1 * 128 + 64, ((raw_rows[0] >= 0) ? raw_rows[0] : 0), ((raw_rows[1] >= 0) ? raw_rows[1] : 0), ((raw_rows[2] >= 0) ? raw_rows[2] : 0), ((raw_rows[3] >= 0) ? raw_rows[3] : 0), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + } + } + k_stage_1 += 1; + if (k_stage_1 == 4) { k_stage_1 = 0; _phase_k_empty ^= 1; } + } + #pragma unroll 1 + for (int tile_3 = 1; tile_3 < num_kv_tiles_3; tile_3++) { + #pragma unroll + for (int qk_stage_2 = 0; qk_stage_2 < 4; qk_stage_2++) { + mbarrier_wait(k_empty_addr + (k_stage_1) * 8, _phase_k_empty); + if (load_warp_rank == 0) { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + } + } + } + int group_1 = load_warp_rank * 4 + lane; + if (lane < 4) { + int index_offset_1 = tile_3 * 128 + cta_rank * 64 + group_1 * 4; + int raw_rows_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_1[0])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 3])) + : "r"(smem_indices_addr + (unsigned int)(index_offset_1 * 4))); + tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group_1 * 512), (&tmap_compressed_k), qk_stage_2 * 128, ((raw_rows_1[0] >= 0) ? raw_rows_1[0] : 0), ((raw_rows_1[1] >= 0) ? raw_rows_1[1] : 0), ((raw_rows_1[2] >= 0) ? raw_rows_1[2] : 0), ((raw_rows_1[3] >= 0) ? raw_rows_1[3] : 0), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group_1 * 512) + 8192, (&tmap_compressed_k), qk_stage_2 * 128 + 64, ((raw_rows_1[0] >= 0) ? raw_rows_1[0] : 0), ((raw_rows_1[1] >= 0) ? raw_rows_1[1] : 0), ((raw_rows_1[2] >= 0) ? raw_rows_1[2] : 0), ((raw_rows_1[3] >= 0) ? raw_rows_1[3] : 0), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + } + k_stage_1 += 1; + if (k_stage_1 == 4) { k_stage_1 = 0; _phase_k_empty ^= 1; } + } + int prev_tile = tile_3 - 1; + #pragma unroll + for (int ps_2 = 0; ps_2 < 2; ps_2++) { + mbarrier_wait(v_empty_addr + (v_stage_1) * 8, _phase_v_empty); + if (load_warp_rank == 0) { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(65536)) : "memory"); + } + } + } + int group_2 = load_warp_rank * 4 + lane; + if (lane < 4) { + int index_offset_2 = prev_tile * 128 + ps_2 * 64 + group_2 * 4; + int raw_rows_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_2[0])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 3])) + : "r"(smem_indices_addr + (unsigned int)(index_offset_2 * 4))); + #pragma unroll + for (int acc_stage_4 = 0; acc_stage_4 < 2; acc_stage_4++) { + int dst_v = ((acc_stage_4 == 0) ? smem_v_lo_addr + v_stage_1 * 32768 : smem_v_hi_addr + v_stage_1 * 32768); + if (first_kv_tile_1 + prev_tile == 0) { + tma_gather4_gmem2smem_mc_cta2(dst_v + group_2 * 512, (&tmap_swa_v), acc_stage_4 * 128 * 2 + cta_rank * 128, ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : 0), ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : 0), ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : 0), ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : 0), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + tma_gather4_gmem2smem_mc_cta2(dst_v + group_2 * 512 + 8192, (&tmap_swa_v), acc_stage_4 * 128 * 2 + cta_rank * 128 + 64, ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : 0), ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : 0), ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : 0), ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : 0), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + } else { + tma_gather4_gmem2smem_mc_cta2(dst_v + group_2 * 512, (&tmap_compressed_v), acc_stage_4 * 128 * 2 + cta_rank * 128, ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : 0), ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : 0), ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : 0), ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : 0), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + tma_gather4_gmem2smem_mc_cta2(dst_v + group_2 * 512 + 8192, (&tmap_compressed_v), acc_stage_4 * 128 * 2 + cta_rank * 128 + 64, ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : 0), ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : 0), ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : 0), ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : 0), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + } + } + } + v_stage_1 += 1; + if (v_stage_1 == 2) { v_stage_1 = 0; _phase_v_empty ^= 1; } + } + } + int last_tile = num_kv_tiles_3 - 1; + #pragma unroll + for (int ps_3 = 0; ps_3 < 2; ps_3++) { + mbarrier_wait(v_empty_addr + (v_stage_1) * 8, _phase_v_empty); + if (load_warp_rank == 0) { + if (cta_rank == 0) { + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(65536)) : "memory"); + } + } + } + int group_3 = load_warp_rank * 4 + lane; + if (lane < 4) { + int index_offset_3 = last_tile * 128 + ps_3 * 64 + group_3 * 4; + int raw_rows_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&raw_rows_3[0])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 3])) + : "r"(smem_indices_addr + (unsigned int)(index_offset_3 * 4))); + #pragma unroll + for (int acc_stage_5 = 0; acc_stage_5 < 2; acc_stage_5++) { + int dst_v_1 = ((acc_stage_5 == 0) ? smem_v_lo_addr + v_stage_1 * 32768 : smem_v_hi_addr + v_stage_1 * 32768); + if (first_kv_tile_1 + last_tile == 0) { + tma_gather4_gmem2smem_mc_cta2(dst_v_1 + group_3 * 512, (&tmap_swa_v), acc_stage_5 * 128 * 2 + cta_rank * 128, ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : 0), ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : 0), ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : 0), ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : 0), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + tma_gather4_gmem2smem_mc_cta2(dst_v_1 + group_3 * 512 + 8192, (&tmap_swa_v), acc_stage_5 * 128 * 2 + cta_rank * 128 + 64, ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : 0), ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : 0), ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : 0), ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : 0), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + } else { + tma_gather4_gmem2smem_mc_cta2(dst_v_1 + group_3 * 512, (&tmap_compressed_v), acc_stage_5 * 128 * 2 + cta_rank * 128, ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : 0), ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : 0), ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : 0), ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : 0), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + tma_gather4_gmem2smem_mc_cta2(dst_v_1 + group_3 * 512 + 8192, (&tmap_compressed_v), acc_stage_5 * 128 * 2 + cta_rank * 128 + 64, ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : 0), ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : 0), ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : 0), ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : 0), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); + } + } + } + v_stage_1 += 1; + if (v_stage_1 == 2) { v_stage_1 = 0; _phase_v_empty ^= 1; } + } + asm volatile("barrier.sync 2, 160;" ::: "memory"); + } + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + } + } + // ---- Role: index_warp ---- + if (warp == 13) { + { // index_warp_main + const int index_dummy = 0; + int valid_iter_1 = 0; + #pragma unroll 1 + for (unsigned int work_idx_4 = cluster_id; work_idx_4 < total_work_items; work_idx_4 += num_clusters) { + int split_idx_3 = work_idx_4 % (unsigned int)(((1) ? 5 : 1)); + int query_idx_3 = work_idx_4 / (unsigned int)(((1) ? 5 : 1)); + int owner_col_base = split_idx_3 * ((9 + ((1) ? 5 : 1) - 1) / ((1) ? 5 : 1)) * 128; + int* swa_row = swa_indices + (query_idx_3 * swa_index_stride); + int* compressed_row = compressed_indices + (query_idx_3 * compressed_index_stride); + int valid_slot_1 = valid_iter_1 & 1; + int valid_empty_phase = valid_iter_1 >> 1 & 1 ^ 1; + mbarrier_wait(index_valid_empty_addr + (valid_slot_1) * 8, valid_empty_phase); + #pragma unroll + for (int index_pass = 0; index_pass < (9 + ((1) ? 5 : 1) - 1) / ((1) ? 5 : 1); index_pass++) { + int index_offset_4 = index_pass * 128 + lane * 4; + int combined_col = owner_col_base + index_offset_4; + int* index_src = ((combined_col < 128) ? swa_row + combined_col : compressed_row + (combined_col - 128)); + int values[4]; + #pragma unroll + for (int i = 0; i < 4; i++) { + values[i] = -1; + } + if (combined_col + 4 <= sparse_topk) { + int _vec_load_0[4]; + { + const int4* _ivptr_0 = reinterpret_cast(index_src + 0); + int4 _ivld_0; + _ivld_0 = *_ivptr_0; + _vec_load_0[0 + 0] = _ivld_0.x; + _vec_load_0[0 + 1] = _ivld_0.y; + _vec_load_0[0 + 2] = _ivld_0.z; + _vec_load_0[0 + 3] = _ivld_0.w; + } + #pragma unroll + for (int i_1 = 0; i_1 < 4; i_1++) { + values[i_1] = _vec_load_0[i_1]; + } + } + if (combined_col < sparse_topk && combined_col + 4 > sparse_topk) { + #pragma unroll + for (int i_2 = 0; i_2 < 4; i_2++) { + if (combined_col + i_2 < sparse_topk) { + values[i_2] = index_src[i_2]; + } + } + } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(smem_indices_addr + (unsigned int)(index_offset_4 * 4)), "r"(values[0]), "r"(values[1]), "r"(values[2]), "r"(values[3]) : "memory"); + unsigned int valid_word = (((values[0] >= 0) ? 1 : 0) | ((values[1] >= 0) ? 1 : 0) << 1 | ((values[2] >= 0) ? 1 : 0) << 2 | ((values[3] >= 0) ? 1 : 0) << 3) << (lane & 7) * 4; + unsigned int _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, valid_word, 1); + valid_word = valid_word | _shfl_xor_0; + unsigned int _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, valid_word, 2); + valid_word = valid_word | _shfl_xor_1; + unsigned int _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, valid_word, 4); + valid_word = valid_word | _shfl_xor_2; + if ((lane & 7) == 0) { + smem_index_valid[valid_slot_1 * 8 + index_pass * 4 + (lane >> 3)] = valid_word; + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(index_valid_full_addr + (valid_slot_1) * 8); + asm volatile("barrier.sync 2, 160;" ::: "memory"); + asm volatile("barrier.sync 2, 160;" ::: "memory"); + valid_iter_1 = valid_iter_1 + 1; + } + } + } + // ---- Role: empty_warps ---- + if (warp >= 14 && warp <= 15) { + { // empty_warps_main + const int empty_dummy = 0; + #pragma unroll 1 + for (unsigned int work_idx_5 = cluster_id; work_idx_5 < total_work_items; work_idx_5 += num_clusters) { + asm volatile("tcgen05.fence::after_thread_sync;"); + } + } + } + + // Cleanup +} + +} // extern "C" diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_9c3b9b4fa7cba6e44ffc_binding.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_9c3b9b4fa7cba6e44ffc_binding.cu new file mode 100644 index 00000000000..e0f1513b40a --- /dev/null +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_9c3b9b4fa7cba6e44ffc_binding.cu @@ -0,0 +1,461 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +// Generated by the Cake source exporter — do not edit. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_9c3b9b4fa7cba6e44ffc'. +#include +#include +#include +#include + +#include "tvm_ffi_utils.h" + +#include +#include +#include +#include +#include +#include +#include + +extern "C" __global__ void kernel_cake_dsv4_9c3b9b4fa7cba6e44ffc(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items); + + +namespace cake_host_shim_42a0596262e73abf { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t CakeDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +// 4D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + int64_t d3 = t.size(t.ndim() - 3); + TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + TVM_FFI_CHECK(d1 % 128 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 128; + uint64_t global_dim[4] = {(uint64_t)(128), (uint64_t)(d2), (uint64_t)((d1 / 128)), (uint64_t)(d3)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (128, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 128; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + (uint64_t)((carrier_stride_1 * 8) / 8), + (uint64_t)((carrier_stride_2 * 8) / 8), + }; + uint32_t box_dim[4] = {128u, 64u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_swa_kv' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_swa_kv' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; + TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (128, 1) exceeds resolved global dims for 'tmap_swa_kv'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + }; + uint32_t box_dim[2] = {128u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_compressed_kv' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_compressed_kv(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_compressed_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_compressed_kv' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_compressed_kv' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_compressed_kv"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_compressed_kv' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved a non-positive global dim"; + TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (128, 1) exceeds resolved global dims for 'tmap_compressed_kv'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + }; + uint32_t box_dim[2] = {128u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_kv') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_compressed_kv, TensorView arg_O, TensorView arg_partial_lse, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_num_query_tokens, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 1, 8, 1); + CheckContiguous(arg_tmap_q, "tmap_q"); + CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); + CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 1, 8, 1); + CheckCudaTensor(arg_tmap_compressed_kv, "tmap_compressed_kv"); + CheckDtype(arg_tmap_compressed_kv, "tmap_compressed_kv", 1, 8, 1); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_partial_lse, "partial_lse"); + CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); + CheckContiguous(arg_partial_lse, "partial_lse"); + CheckCudaTensor(arg_swa_indices, "swa_indices"); + CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); + CheckContiguous(arg_swa_indices, "swa_indices"); + CheckCudaTensor(arg_compressed_indices, "compressed_indices"); + CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); + CheckContiguous(arg_compressed_indices, "compressed_indices"); + CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); + CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); + CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); + CheckContiguous(arg_bmm1_scale, "bmm1_scale"); + CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); + CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); + CheckContiguous(arg_bmm2_scale, "bmm2_scale"); + TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) + << "scalar 'num_heads' value " << arg_num_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) + << "scalar 'swa_index_stride' value " << arg_swa_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) + << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) + << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_query_tokens >= -2147483648LL && arg_num_query_tokens <= 2147483647LL, ValueError) + << "scalar 'num_query_tokens' value " << arg_num_query_tokens + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) + << "scalar 'sparse_topk' value " << arg_sparse_topk + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) + << "scalar 'has_sinks' value " << arg_has_sinks + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); + CheckSameCudaDevice(arg_tmap_compressed_kv, arg_tmap_q, "tmap_compressed_kv", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_partial_lse, arg_tmap_q, "partial_lse", "tmap_q"); + CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); + CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); + CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); + CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") must be divisible by cluster dims (2, 1, 1)"; + + cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); + CUtensorMap p_tmap_compressed_kv = EncodeTma_tmap_compressed_kv(arg_tmap_compressed_kv); + __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); + float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); + int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); + int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); + int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); + float* p_sinks = static_cast(arg_sinks.data_ptr()); + float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); + float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); + int32_t v_num_heads = (int32_t)arg_num_heads; + int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; + int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; + int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; + int32_t v_num_query_tokens = (int32_t)arg_num_query_tokens; + int32_t v_sparse_topk = (int32_t)arg_sparse_topk; + int32_t v_has_sinks = (int32_t)arg_has_sinks; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_tmap_compressed_kv, &p_O, &p_partial_lse, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_num_query_tokens, &v_sparse_topk, &v_has_sinks, &v_total_work_items}; + + static std::mutex smem_mu; + static auto* smem_status_by_device = + new std::unordered_map(); + cudaError_t smem_status = cudaSuccess; + { + std::lock_guard lock(smem_mu); + auto it = smem_status_by_device->find(dev.device_id); + if (it == smem_status_by_device->end()) { + int current_device = -1; + smem_status = cudaGetDevice(¤t_device); + if (smem_status == cudaSuccess && current_device != dev.device_id) { + smem_status = cudaErrorInvalidDevice; + } + if (smem_status == cudaSuccess) { + smem_status = cudaFuncSetAttribute( + reinterpret_cast(kernel_cake_dsv4_9c3b9b4fa7cba6e44ffc), + cudaFuncAttributeMaxDynamicSharedMemorySize, + 190976); + } + smem_status_by_device->emplace(dev.device_id, smem_status); + } else { + smem_status = it->second; + } + } + TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) + << "cudaFuncSetAttribute for kernel_cake_dsv4_9c3b9b4fa7cba6e44ffc failed: " + << cudaGetErrorString(smem_status); + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(512u, 1u, 1u); + + // Extended direct-source launch for a linked CUDA kernel symbol. + cudaLaunchAttribute attrs[2]{}; + int n = 0; + attrs[n].id = cudaLaunchAttributeClusterDimension; + attrs[n].val.clusterDim.x = 2u; + attrs[n].val.clusterDim.y = 1u; + attrs[n].val.clusterDim.z = 1u; + ++n; + attrs[n].id = cudaLaunchAttributeClusterSchedulingPolicyPreference; + attrs[n].val.clusterSchedulingPolicyPreference = cudaClusterSchedulingPolicySpread; + ++n; + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes = 190976u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; + cudaError_t launch_status = cudaLaunchKernelExC( + &config, reinterpret_cast(kernel_cake_dsv4_9c3b9b4fa7cba6e44ffc), kargs); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernelExC for kernel_cake_dsv4_9c3b9b4fa7cba6e44ffc failed: " + << cudaGetErrorString(launch_status); +} + +} // namespace cake_host_shim_42a0596262e73abf + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_42a0596262e73abf::Run); diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_9c3b9b4fa7cba6e44ffc_kernel.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_9c3b9b4fa7cba6e44ffc_kernel.cu new file mode 100644 index 00000000000..a623a6872fa --- /dev/null +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_9c3b9b4fa7cba6e44ffc_kernel.cu @@ -0,0 +1,2826 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +#define CAKE_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_SCRATCH_OFFSET 0 +#define NUM_K_PIPE_STAGES 2 +#define NUM_V_PIPE_STAGES 2 +#define NUM_KV_PIPE_STAGES 1 +#define NUM_INDEX_PIPE_STAGES 6 +#define SMEM_SMEM_Q_OFF 1024 +#define SMEM_SMEM_Q_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_STRIDE 8192 +#define SMEM_SMEM_KV_OFF 33792 +#define SMEM_SMEM_KV_STAGE_BYTES 16384 +#define SMEM_SMEM_KV_STRIDE 16384 +#define SMEM_SMEM_V_OFF 33792 +#define SMEM_SMEM_V_STAGE_BYTES 16384 +#define SMEM_SMEM_V_STRIDE 16384 +#define SMEM_SMEM_Q_FULL_OFF 1024 +#define SMEM_SMEM_Q_FULL_STAGE_BYTES 32768 +#define SMEM_SMEM_Q_FULL_STRIDE 32768 +#define SMEM_SMEM_K_FULL_OFF 33792 +#define SMEM_SMEM_K_FULL_STAGE_BYTES 65536 +#define SMEM_SMEM_K_FULL_STRIDE 65536 +#define SMEM_SMEM_V_FULL_OFF 33792 +#define SMEM_SMEM_V_FULL_STAGE_BYTES 65536 +#define SMEM_SMEM_V_FULL_STRIDE 65536 +#define SMEM_SMEM_STATS_MAX_OFF 164864 +#define SMEM_SMEM_STATS_MAX_STAGE_BYTES 1024 +#define SMEM_SMEM_STATS_MAX_STRIDE 1024 +#define SMEM_SMEM_STATS_SUM_OFF 165888 +#define SMEM_SMEM_STATS_SUM_STAGE_BYTES 512 +#define SMEM_SMEM_STATS_SUM_STRIDE 512 +#define SMEM_SMEM_STATS_FINAL_MAX_OFF 166400 +#define SMEM_SMEM_STATS_FINAL_MAX_STAGE_BYTES 512 +#define SMEM_SMEM_STATS_FINAL_MAX_STRIDE 512 +#define SMEM_SMEM_SOFTMAX_WARP_PAIR_EXCHANGE_OFF 166912 +#define SMEM_SMEM_SOFTMAX_WARP_PAIR_EXCHANGE_STAGE_BYTES 1024 +#define SMEM_SMEM_SOFTMAX_WARP_PAIR_EXCHANGE_STRIDE 1024 +#define SMEM_SMEM_CORR_WARP_PAIR_EXCHANGE_OFF 167936 +#define SMEM_SMEM_CORR_WARP_PAIR_EXCHANGE_STAGE_BYTES 512 +#define SMEM_SMEM_CORR_WARP_PAIR_EXCHANGE_STRIDE 512 +#define SMEM_SMEM_SPARSE_INDICES_OFF 168448 +#define SMEM_SMEM_SPARSE_INDICES_STAGE_BYTES 1024 +#define SMEM_SMEM_SPARSE_INDICES_STRIDE 1024 +#define SMEM_SMEM_P_FP8_OFF 174592 +#define SMEM_SMEM_P_FP8_STAGE_BYTES 8192 +#define SMEM_SMEM_P_FP8_STRIDE 8192 +#define SMEM_SMEM_MERGE_RECV_OFF 33792 +#define SMEM_SMEM_MERGE_RECV_STAGE_BYTES 65536 +#define SMEM_SMEM_MERGE_RECV_STRIDE 65536 +#define SMEM_SMEM_MERGE_SEND_OFF 99328 +#define SMEM_SMEM_MERGE_SEND_STAGE_BYTES 65536 +#define SMEM_SMEM_MERGE_SEND_STRIDE 65536 +#define SMEM_SMEM_MERGE_STATS_RECV_OFF 174592 +#define SMEM_SMEM_MERGE_STATS_RECV_STAGE_BYTES 512 +#define SMEM_SMEM_MERGE_STATS_RECV_STRIDE 512 +#define SMEM_SMEM_MERGE_STATS_SEND_OFF 182784 +#define SMEM_SMEM_MERGE_STATS_SEND_STAGE_BYTES 512 +#define SMEM_SMEM_MERGE_STATS_SEND_STRIDE 512 +#define SMEM_TOTAL 190976 +#define THREADS 512 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count) : "memory"); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f8f6f4( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x4(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x4.b32" + " {%0, %1, %2, %3}, [%4];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); +} + + +__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + " [%0], {%1, %2, %3, %4};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ float2 ex2_emulation_f32x2_value(float2 value) { + const float c0 = 1.0f, c1 = 0.695146143436431884765625f; + const float c2 = 0.227564394474029541015625f, c3 = 0.077119089663028717041015625f; + const float magic = 12582912.0f; + float x0 = max_noftz(value.x, -127.0f), x1 = max_noftz(value.y, -127.0f); + float2 xc2 = make_float2(x0, x1), magic2 = make_float2(magic, magic); + float2 xr2; + asm("add.rm.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xr2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&magic2)); + float2 c3_2 = make_float2(c3, c3), c2_2 = make_float2(c2, c2); + float2 c1_2 = make_float2(c1, c1), c0_2 = make_float2(c0, c0); + float2 xrb2, xfrac2; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xrb2) + : "l"(*(unsigned long long*)&xr2), "l"(*(unsigned long long*)&magic2)); + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xfrac2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&xrb2)); + float2 poly2; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&c3_2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c2_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c1_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c0_2)); + int x0r_i, x1r_i, p0_i, p1_i; + asm("mov.b64 {%0, %1}, %2;" : "=r"(x0r_i), "=r"(x1r_i) : "l"(*(unsigned long long*)&xr2)); + asm("mov.b64 {%0, %1}, %2;" : "=r"(p0_i), "=r"(p1_i) : "l"(*(unsigned long long*)&poly2)); + float r0, r1; + asm("mov.b32 %0, %1;" : "=f"(r0) : "r"((x0r_i << 23) + p0_i)); + asm("mov.b32 %0, %1;" : "=f"(r1) : "r"((x1r_i << 23) + p1_i)); + return make_float2(r0, r1); +} + +__device__ __forceinline__ void ex2_emulation_f32x2(float* x0_ptr, float* x1_ptr) { + float2 result = ex2_emulation_f32x2_value(make_float2(*x0_ptr, *x1_ptr)); + *x0_ptr = result.x; *x1_ptr = result.y; +} + +__device__ __forceinline__ void softmax_frag_exp2_cast( + float* sv, uint32_t* pv, int use_emu) +{ + #pragma unroll + for (int j = 0; j < 16; j++) { + if (use_emu && j >= 12) + ex2_emulation_f32x2(&sv[j*2], &sv[j*2+1]); + else { + sv[j*2] = approx_exp2(sv[j*2]); + sv[j*2+1] = approx_exp2(sv[j*2+1]); + } + } + #pragma unroll + for (int j = 0; j < 16; j++) { + __nv_bfloat162 bf = __float22bfloat162_rn({sv[j*2], sv[j*2+1]}); + pv[j] = reinterpret_cast(bf); + } +} + + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_2d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3}], [%4];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr) { + // Canonical .shared::cta form for non-multicast gather4, matching + // trtllm-gen / cuda_ptx and the PTX ISA qualifier order + // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, + // .shared::cluster is reserved for the multicast variant (ctaMask). + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_st_x8_u32(int addr, const uint32_t* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x8.b32" + " [%0], {%1,%2,%3,%4,%5,%6,%7,%8};" + :: "r"(addr), + "r"(src[0]), "r"(src[1]), "r"(src[2]), "r"(src[3]), + "r"(src[4]), "r"(src[5]), "r"(src[6]), "r"(src[7])); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(512, 1) __cluster_dims__(2,1,1) void +kernel_cake_dsv4_9c3b9b4fa7cba6e44ffc(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); + + const int mbar_base = smem; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define k_full_addr (mbar_base + 16) + #define k_empty_addr (mbar_base + 32) + #define v_full_addr (mbar_base + 48) + #define v_empty_addr (mbar_base + 64) + #define index_full_addr (mbar_base + 80) + #define index_empty_addr (mbar_base + 128) + #define s_full_addr (mbar_base + 176) + #define p_full_addr (mbar_base + 192) + #define corr_done_addr (mbar_base + 208) + #define stats_addr (mbar_base + 224) + #define sum_ready_addr (mbar_base + 240) + #define o_done_addr (mbar_base + 248) + #define pv_done_addr (mbar_base + 256) + #define s_seeded_addr (mbar_base + 264) + #define q_pair_ready_addr (mbar_base + 272) + #define kv_pair_ready_addr (mbar_base + 280) + #define pv_pair_ready_addr (mbar_base + 288) + #define tmem_dealloc_addr (mbar_base + 304) + #define tmem_dealloc_peer_addr (mbar_base + 312) + #define merge_ready_addr (mbar_base + 320) + #define merge_full_addr (mbar_base + 328) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 2; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + + // Kernel setup ops + uint8_t* smem_q = reinterpret_cast(smem_raw + 1024); + const int smem_q_addr = smem + 1024; + uint8_t* smem_kv = reinterpret_cast(smem_raw + 33792); + const int smem_kv_addr = smem + 33792; + uint8_t* smem_v = reinterpret_cast(smem_raw + 33792); + const int smem_v_addr = smem + 33792; + uint8_t* smem_q_full = reinterpret_cast(smem_raw + 1024); + const int smem_q_full_addr = smem + 1024; + uint8_t* smem_k_full = reinterpret_cast(smem_raw + 33792); + const int smem_k_full_addr = smem + 33792; + uint8_t* smem_v_full = reinterpret_cast(smem_raw + 33792); + const int smem_v_full_addr = smem + 33792; + float* smem_stats_max = reinterpret_cast(smem_raw + 164864); + const int smem_stats_max_addr = smem + 164864; + float* smem_stats_sum = reinterpret_cast(smem_raw + 165888); + const int smem_stats_sum_addr = smem + 165888; + float* smem_stats_final_max = reinterpret_cast(smem_raw + 166400); + const int smem_stats_final_max_addr = smem + 166400; + float* smem_softmax_warp_pair_exchange = reinterpret_cast(smem_raw + 166912); + const int smem_softmax_warp_pair_exchange_addr = smem + 166912; + float* smem_corr_warp_pair_exchange = reinterpret_cast(smem_raw + 167936); + const int smem_corr_warp_pair_exchange_addr = smem + 167936; + int* smem_sparse_indices = reinterpret_cast(smem_raw + 168448); + const int smem_sparse_indices_addr = smem + 168448; + uint8_t* smem_p_fp8 = reinterpret_cast(smem_raw + 174592); + const int smem_p_fp8_addr = smem + 174592; + float* smem_merge_recv = reinterpret_cast(smem_raw + 33792); + const int smem_merge_recv_addr = smem + 33792; + float* smem_merge_send = reinterpret_cast(smem_raw + 99328); + const int smem_merge_send_addr = smem + 99328; + float* smem_merge_stats_recv = reinterpret_cast(smem_raw + 174592); + const int smem_merge_stats_recv_addr = smem + 174592; + float* smem_merge_stats_send = reinterpret_cast(smem_raw + 182784); + const int smem_merge_stats_send_addr = smem + 182784; + { + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + } + + // Mbarrier init (23 pipeline groups, 0 ordered-sequence groups, 42 barriers) + // Mbarriers at smem_raw[0..336) + + if (warp == 12) { + uint32_t leader = elect_sync(); + if (leader) { + // q_full: 1 barriers, init_count=1 + mbarrier_init(smem + 0, 1); + // q_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 8, 1); + // --- pipeline 'k_pipe' --- + // k_full: 2 barriers, init_count=1 + mbarrier_init(smem + 16, 1); + mbarrier_init(smem + 24, 1); + // k_empty: 2 barriers, init_count=1 + mbarrier_init(smem + 32, 1); + mbarrier_init(smem + 40, 1); + // --- pipeline 'v_pipe' --- + // v_full: 2 barriers, init_count=2 + mbarrier_init(smem + 48, 2); + mbarrier_init(smem + 56, 2); + // v_empty: 2 barriers, init_count=1 + mbarrier_init(smem + 64, 1); + mbarrier_init(smem + 72, 1); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + if (warp == 9) { + uint32_t leader = elect_sync(); + if (leader) { + // --- pipeline 'index_pipe' --- + // index_full: 6 barriers, init_count=32 + mbarrier_init(smem + 80, 32); + mbarrier_init(smem + 88, 32); + mbarrier_init(smem + 96, 32); + mbarrier_init(smem + 104, 32); + mbarrier_init(smem + 112, 32); + mbarrier_init(smem + 120, 32); + // index_empty: 6 barriers, init_count=192 + mbarrier_init(smem + 128, 192); + mbarrier_init(smem + 136, 192); + mbarrier_init(smem + 144, 192); + mbarrier_init(smem + 152, 192); + mbarrier_init(smem + 160, 192); + mbarrier_init(smem + 168, 192); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // s_full: stages (0,), init_count=1 + mbarrier_init(smem + 176, 1); + // p_full: stages (0,), init_count=64 + mbarrier_init(smem + 192, 64); + // o_done: 1 barriers, init_count=1 + mbarrier_init(smem + 248, 1); + // pv_done: 1 barriers, init_count=1 + mbarrier_init(smem + 256, 1); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + if (warp == 1) { + uint32_t leader = elect_sync(); + if (leader) { + // s_full: stages (1,), init_count=1 + mbarrier_init(smem + 184, 1); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + if (warp == 2) { + uint32_t leader = elect_sync(); + if (leader) { + // p_full: stages (1,), init_count=64 + mbarrier_init(smem + 200, 64); + // s_seeded: 1 barriers, init_count=256 + mbarrier_init(smem + 264, 256); + // q_pair_ready: 1 barriers, init_count=64 + mbarrier_init(smem + 272, 64); + // --- pipeline 'kv_pipe' --- + // kv_pair_ready: 1 barriers, init_count=64 + mbarrier_init(smem + 280, 64); + // pv_pair_ready: stages (0,), init_count=64 + mbarrier_init(smem + 288, 64); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + if (warp == 4) { + uint32_t leader = elect_sync(); + if (leader) { + // corr_done: stages (0,), init_count=64 + mbarrier_init(smem + 208, 64); + // stats: stages (0,), init_count=64 + mbarrier_init(smem + 224, 64); + // sum_ready: 1 barriers, init_count=64 + mbarrier_init(smem + 240, 64); + // pv_pair_ready: stages (1,), init_count=64 + mbarrier_init(smem + 296, 64); + // tmem_dealloc: 1 barriers, init_count=288 + mbarrier_init(smem + 304, 288); + // merge_ready: 1 barriers, init_count=64 + mbarrier_init(smem + 320, 64); + // merge_full: 1 barriers, init_count=1 + mbarrier_init(smem + 328, 1); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + if (warp == 6) { + uint32_t leader = elect_sync(); + if (leader) { + // corr_done: stages (1,), init_count=64 + mbarrier_init(smem + 216, 64); + // stats: stages (1,), init_count=64 + mbarrier_init(smem + 232, 64); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + if (warp == 5) { + uint32_t leader = elect_sync(); + if (leader) { + // tmem_dealloc_peer: 1 barriers, init_count=32 + mbarrier_init(smem + 312, 32); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncwarp(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 336); + if (warp == 0) { + int _tmem_hold = smem + 336; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } + + __syncthreads(); + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_tmem_scratch = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Inc phase consumes the registers released above. + if (warp >= 0 && warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 152;"); + } else if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 144;"); + } else if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 136;"); + } + + // ---- Role: index_warp ---- + if (warp == 9) { + { // index_warp_main + const int index_dummy = 0; + unsigned int _phase_index_empty = 1; + unsigned int _phase_q_full_0 = 0; + { + int all_num_kv_tiles = (sparse_topk + 128 - 1) / 128; + unsigned int index_prod_stage = 0; + #pragma unroll + for (unsigned int static_work = 0; static_work < 1; static_work++) { + int split_idx = 0; + int query_idx = bid / 2; + int tiles_per_split = all_num_kv_tiles + 1 - 1; + int first_tile = split_idx * tiles_per_split; + int num_index_passes = (tiles_per_split + 1) / 2; + #pragma unroll 1 + for (int index_pass = 0; index_pass < num_index_passes; index_pass++) { + mbarrier_wait(index_empty_addr + (index_prod_stage) * 8, _phase_index_empty); + int index_stage_base = smem_sparse_indices_addr + index_prod_stage * 1024; + int stage_row_base = index_pass * 256; + #pragma unroll + for (int index_half = 0; index_half < 2; index_half++) { + int index_offset = index_half * 128 + lane * 4; + int tile_col_base = first_tile * 128 + stage_row_base + index_half * 128; + int* row_tile_ptr = ((tile_col_base < 128) ? swa_indices + (query_idx * swa_index_stride) : compressed_indices + (query_idx * compressed_index_stride + (tile_col_base - 128))); + int lane_col = tile_col_base + lane * 4; + if (lane_col + 4 <= sparse_topk) { + asm volatile("cp.async.cg.shared::cta.global.L2::128B [%0], [%1], 16;" + :: "r"(index_stage_base + index_offset * 4), "l"(row_tile_ptr + (lane * 4))); + } else { + int tail_rows[4]; + #pragma unroll + for (int row_i = 0; row_i < 4; row_i++) { + tail_rows[row_i] = -1; + } + #pragma unroll + for (int row_i_1 = 0; row_i_1 < 4; row_i_1++) { + if (lane_col + row_i_1 < sparse_topk) { + tail_rows[row_i_1] = row_tile_ptr[lane * 4 + row_i_1]; + } + } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(smem_sparse_indices_addr + ((unsigned int)index_stage_base - smem_sparse_indices_addr + (unsigned int)(index_offset * 4))), "r"(tail_rows[0]), "r"(tail_rows[1]), "r"(tail_rows[2]), "r"(tail_rows[3]) : "memory"); + } + } + asm volatile("cp.async.commit_group;"); + asm volatile("cp.async.wait_group 0;"); + mbarrier_arrive(index_full_addr + (index_prod_stage) * 8); + index_prod_stage += 1; + if (index_prod_stage == 6) { index_prod_stage = 0; _phase_index_empty ^= 1; } + } + } + } + } + } + // ---- Role: load_warp ---- + if (warp >= 12 && warp <= 15) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 72;"); + { // load_warp_main + const int wg2_dummy = 0; + const int load_warp_rank = warp - ((1) ? 12 : 9); + int all_num_kv_tiles_1 = (sparse_topk + 128 - 1) / 128; + unsigned int load_k_stage = 0; + unsigned int load_v_stage = 0; + unsigned int load_kv_stage = 0; + unsigned int load_k_index_stage = 0; + unsigned int load_v_index_stage = 0; + int k_cta_offset = bid % 2 * 64; + unsigned int _phase_q_empty_0 = 1; + unsigned int _phase_index_full = 0; + unsigned int _phase_k_empty = 1; + #pragma unroll + for (unsigned int work_idx = 0; work_idx < 1; work_idx++) { + int split_idx_1 = 0; + int query_idx_1 = bid / 2 >> 1; + int v_chunk = bid / 2 & 1; + { + split_idx_1 = 0; + query_idx_1 = bid / 2; + v_chunk = bid % 2; + } + int tiles_per_split_1 = all_num_kv_tiles_1 + 1 - 1; + int first_tile_1 = split_idx_1 * tiles_per_split_1; + int num_kv_tiles = tiles_per_split_1; + int sparse_extent = num_kv_tiles * 128; + int wt_num_local_tiles = (num_kv_tiles - bid % 2 + 1) / 2; + mbarrier_wait(q_empty_addr, _phase_q_empty_0); + _phase_q_empty_0 ^= 1; + if (load_warp_rank == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(q_full_addr, 32768); + #pragma unroll + for (int q_stage = 0; q_stage < 4; q_stage++) { + tma_4d_gmem2smem(smem_q_full_addr + (unsigned int)(q_stage * 8192), (&tmap_q), 0, 0, q_stage, query_idx_1, q_full_addr); + } + } + } + #pragma unroll 1 + for (int wt_tile = 0; wt_tile < wt_num_local_tiles; wt_tile++) { + int wt_global_tile = first_tile_1 + wt_tile * 2 + bid % 2; + mbarrier_wait(index_full_addr + (load_k_index_stage) * 8, _phase_index_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + int wt_index_base = smem_sparse_indices_addr + load_k_index_stage * 1024; + mbarrier_wait(k_empty_addr + (load_k_stage) * 8, _phase_k_empty); + if (load_warp_rank == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(k_full_addr + (load_k_stage) * 8, 65536); + } + } + int wt_dst = smem_k_full_addr + load_k_stage * 65536 + (unsigned int)(load_warp_rank * 16384); + int wt_group = lane; + int wt_group_offset = bid % 2 * 128 + wt_group * 4; + int wt_rows[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&wt_rows[0])), "=r"(*reinterpret_cast(&wt_rows[(0) + 1])), "=r"(*reinterpret_cast(&wt_rows[(0) + 2])), "=r"(*reinterpret_cast(&wt_rows[(0) + 3])) + : "r"(wt_index_base + wt_group_offset * 4)); + int wt_row0 = ((wt_rows[0] >= 0) ? wt_rows[0] : 0); + int wt_row1 = ((wt_rows[1] >= 0) ? wt_rows[1] : 0); + int wt_row2 = ((wt_rows[2] >= 0) ? wt_rows[2] : 0); + int wt_row3 = ((wt_rows[3] >= 0) ? wt_rows[3] : 0); + if (wt_global_tile == 0) { + tma_gather4_gmem2smem(wt_dst + wt_group * 512, (&tmap_swa_kv), load_warp_rank * 128, wt_row0, wt_row1, wt_row2, wt_row3, k_full_addr + (load_k_stage) * 8); + } else { + tma_gather4_gmem2smem(wt_dst + wt_group * 512, (&tmap_compressed_kv), load_warp_rank * 128, wt_row0, wt_row1, wt_row2, wt_row3, k_full_addr + (load_k_stage) * 8); + } + load_k_stage += 1; + if (load_k_stage == 2) { load_k_stage = 0; _phase_k_empty ^= 1; } + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(index_empty_addr + (load_k_index_stage) * 8); + load_k_index_stage += 1; + if (load_k_index_stage == 6) { load_k_index_stage = 0; _phase_index_full ^= 1; } + } + } + mbarrier_arrive(tmem_dealloc_addr); + } + } + // ---- Role: softmax_wg ---- + if (warp == 0 || warp == 2) { + { // softmax_wg_main + float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; + const int wg_dummy_inc = 0; + int all_num_kv_tiles_2 = (sparse_topk + 128 - 1) / 128; + const int tmem_row_base = ((1) ? warp % 2 * 32 : warp % 4 * 32); + const int tmem_score_row_base = ((1) ? (int)(warp % 4 * 32) : tmem_row_base); + const int n_half = ((1) ? (int)(warp % 4 / 2) : 0); + const int my_row = tmem_row_base + lane; + const int stats_row = n_half * 64 + my_row; + int softmax_tile_cursor = 0; + unsigned int softmax_index_stage = 0; + unsigned int _phase_q_full_0_1 = 0; + #pragma unroll + for (unsigned int work_idx_1 = 0; work_idx_1 < 1; work_idx_1++) { + int split_idx_2 = 0; + int query_idx_2 = bid / 2 >> 1; + { + split_idx_2 = 0; + query_idx_2 = bid / 2; + } + int tiles_per_split_2 = all_num_kv_tiles_2 + 1 - 1; + int first_tile_2 = split_idx_2 * tiles_per_split_2; + int num_kv_tiles_1 = tiles_per_split_2; + int _max_0 = ((sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) : (0)); + int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); + int active_topk = _min_0; + float row_max_val = -CAKE_INF; + float row_sum_val = 0.0f; + int sink_head = ((1) ? my_row : my_row); + if (has_sinks != 0 && sink_head < num_heads && split_idx_2 + bid % 2 == 0) { + row_max_val = sinks[sink_head] * 1.4426950408889634f / softmax_scale_log2; + row_sum_val = 1.0f; + } + #pragma unroll 1 + for (int tile = 0; tile < (num_kv_tiles_1 - bid % 2 + 1) / 2; tile++) { + int pipeline_tile = softmax_tile_cursor + tile; + int phase = pipeline_tile & 1; + int s_wait_phase = pipeline_tile >> 1 & 1; + mbarrier_wait(s_full_addr + (phase) * 8, s_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int s_off = ((phase != 0) ? 128 : 0); + int s_base = taddr + (unsigned int)s_off + (unsigned int)(tmem_score_row_base << 16); + float new_max = row_max_val; + int valid_sparse_cols = ((active_topk < sparse_topk) ? active_topk : sparse_topk); + valid_sparse_cols = valid_sparse_cols - (first_tile_2 + tile * 2 + bid % 2) * 128 - n_half * 64; + if (valid_sparse_cols < 0) { + valid_sparse_cols = 0; + } + if (valid_sparse_cols > 64) { + valid_sparse_cols = 64; + } + int index_tile_addr = smem_sparse_indices_addr + softmax_index_stage * 1024 + (unsigned int)(bid % 2 * 512); + int index_lane_addr = index_tile_addr + (n_half * 64 + lane) * 4; + int staged_index[2]; + unsigned int invalid_cols[2]; + unsigned int any_invalid = 0; + #pragma unroll + for (int w = 0; w < 2; w++) { + asm volatile("ld.shared.b32 %0, [%1];" : "=r"(*reinterpret_cast(&staged_index[w])) : "r"(index_lane_addr + w * 128)); + unsigned int _vote_0 = __ballot_sync(0xFFFFFFFF, staged_index[w] < 0); + invalid_cols[w] = _vote_0; + any_invalid = any_invalid | invalid_cols[w]; + } + if (tile >= 0) { + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(index_empty_addr + (softmax_index_stage) * 8); + softmax_index_stage = softmax_index_stage + 1; + if (softmax_index_stage == 6) { + softmax_index_stage = 0; + } + } + float _tmem_load_0[4]; + tmem_ld_x4(&_tmem_load_0[0], s_base); + float _tmem_load_1[64]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]) + : "r"(s_base)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]) + : "r"(s_base + 32)); + { + uint32_t _slice_lo_mask_0; + { + int _lim_0 = valid_sparse_cols; + if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_0 & (1u << 0))) _tmem_load_1[0] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 1))) _tmem_load_1[1] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 2))) _tmem_load_1[2] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 3))) _tmem_load_1[3] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 4))) _tmem_load_1[4] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 5))) _tmem_load_1[5] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 6))) _tmem_load_1[6] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 7))) _tmem_load_1[7] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 8))) _tmem_load_1[8] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 9))) _tmem_load_1[9] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 10))) _tmem_load_1[10] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 11))) _tmem_load_1[11] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 12))) _tmem_load_1[12] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 13))) _tmem_load_1[13] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 14))) _tmem_load_1[14] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 15))) _tmem_load_1[15] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 16))) _tmem_load_1[16] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 17))) _tmem_load_1[17] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 18))) _tmem_load_1[18] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 19))) _tmem_load_1[19] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 20))) _tmem_load_1[20] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 21))) _tmem_load_1[21] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 22))) _tmem_load_1[22] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 23))) _tmem_load_1[23] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 24))) _tmem_load_1[24] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 25))) _tmem_load_1[25] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 26))) _tmem_load_1[26] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 27))) _tmem_load_1[27] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 28))) _tmem_load_1[28] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 29))) _tmem_load_1[29] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 30))) _tmem_load_1[30] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 31))) _tmem_load_1[31] = -CAKE_INF; + uint32_t _slice_lo_mask_1; + { + int _lim_1 = valid_sparse_cols - 32; + if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) _tmem_load_1[32] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) _tmem_load_1[33] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) _tmem_load_1[34] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) _tmem_load_1[35] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) _tmem_load_1[36] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) _tmem_load_1[37] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) _tmem_load_1[38] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) _tmem_load_1[39] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) _tmem_load_1[40] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) _tmem_load_1[41] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) _tmem_load_1[42] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) _tmem_load_1[43] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) _tmem_load_1[44] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) _tmem_load_1[45] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) _tmem_load_1[46] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) _tmem_load_1[47] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) _tmem_load_1[48] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) _tmem_load_1[49] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) _tmem_load_1[50] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) _tmem_load_1[51] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) _tmem_load_1[52] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) _tmem_load_1[53] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) _tmem_load_1[54] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) _tmem_load_1[55] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) _tmem_load_1[56] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) _tmem_load_1[57] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) _tmem_load_1[58] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) _tmem_load_1[59] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) _tmem_load_1[60] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) _tmem_load_1[61] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) _tmem_load_1[62] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) _tmem_load_1[63] = -CAKE_INF; + if (any_invalid != 0) { + #pragma unroll + for (int w_1 = 0; w_1 < 2; w_1++) { + unsigned int mask_word = invalid_cols[w_1]; + _tmem_load_1[w_1 * 32] = (((mask_word & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32]); + _tmem_load_1[w_1 * 32 + 1] = (((mask_word >> 1 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 1]); + _tmem_load_1[w_1 * 32 + 2] = (((mask_word >> 2 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 2]); + _tmem_load_1[w_1 * 32 + 3] = (((mask_word >> 3 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 3]); + _tmem_load_1[w_1 * 32 + 4] = (((mask_word >> 4 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 4]); + _tmem_load_1[w_1 * 32 + 5] = (((mask_word >> 5 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 5]); + _tmem_load_1[w_1 * 32 + 6] = (((mask_word >> 6 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 6]); + _tmem_load_1[w_1 * 32 + 7] = (((mask_word >> 7 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 7]); + _tmem_load_1[w_1 * 32 + 8] = (((mask_word >> 8 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 8]); + _tmem_load_1[w_1 * 32 + 9] = (((mask_word >> 9 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 9]); + _tmem_load_1[w_1 * 32 + 10] = (((mask_word >> 10 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 10]); + _tmem_load_1[w_1 * 32 + 11] = (((mask_word >> 11 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 11]); + _tmem_load_1[w_1 * 32 + 12] = (((mask_word >> 12 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 12]); + _tmem_load_1[w_1 * 32 + 13] = (((mask_word >> 13 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 13]); + _tmem_load_1[w_1 * 32 + 14] = (((mask_word >> 14 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 14]); + _tmem_load_1[w_1 * 32 + 15] = (((mask_word >> 15 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 15]); + _tmem_load_1[w_1 * 32 + 16] = (((mask_word >> 16 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 16]); + _tmem_load_1[w_1 * 32 + 17] = (((mask_word >> 17 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 17]); + _tmem_load_1[w_1 * 32 + 18] = (((mask_word >> 18 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 18]); + _tmem_load_1[w_1 * 32 + 19] = (((mask_word >> 19 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 19]); + _tmem_load_1[w_1 * 32 + 20] = (((mask_word >> 20 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 20]); + _tmem_load_1[w_1 * 32 + 21] = (((mask_word >> 21 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 21]); + _tmem_load_1[w_1 * 32 + 22] = (((mask_word >> 22 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 22]); + _tmem_load_1[w_1 * 32 + 23] = (((mask_word >> 23 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 23]); + _tmem_load_1[w_1 * 32 + 24] = (((mask_word >> 24 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 24]); + _tmem_load_1[w_1 * 32 + 25] = (((mask_word >> 25 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 25]); + _tmem_load_1[w_1 * 32 + 26] = (((mask_word >> 26 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 26]); + _tmem_load_1[w_1 * 32 + 27] = (((mask_word >> 27 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 27]); + _tmem_load_1[w_1 * 32 + 28] = (((mask_word >> 28 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 28]); + _tmem_load_1[w_1 * 32 + 29] = (((mask_word >> 29 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 29]); + _tmem_load_1[w_1 * 32 + 30] = (((mask_word >> 30 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 30]); + _tmem_load_1[w_1 * 32 + 31] = (((mask_word >> 31 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 31]); + } + } + float2 _reg_reduce_max2_2 = {-CAKE_INF, -CAKE_INF}; + row_max_x32_accum(&_tmem_load_1[0], _reg_reduce_max2_2); + row_max_x32_accum(&_tmem_load_1[32], _reg_reduce_max2_2); + float _tmem_load_1_max = row_max_reduce(_reg_reduce_max2_2); + float _max_1 = max_noftz(new_max, _tmem_load_1_max); + new_max = _max_1; + smem_softmax_warp_pair_exchange[phase * 128 + stats_row] = new_max; + asm volatile("barrier.sync %0, 64;" :: "r"(2 + warp % 2) : "memory"); + float _max_2 = max_noftz(new_max, smem_softmax_warp_pair_exchange[phase * 128 + (stats_row ^ 64)]); + new_max = _max_2; + } + float _fma_0 = __fmaf_rn(row_max_val, softmax_scale_log2, (-new_max) * softmax_scale_log2); + float delta = _fma_0; + float _exp2_0 = approx_exp2(delta); + float exp_delta = _exp2_0; + float acc_scale = ((row_max_val > -CAKE_INF) ? exp_delta : 1.0f); + smem_stats_max[phase * 128 + stats_row] = acc_scale; + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(stats_addr + (phase) * 8); + row_max_val = new_max; + float safe_max = ((new_max == -CAKE_INF) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + float block_sum = 0.0f; + { + { + const float2 _fma_b2_3 = {softmax_scale_log2, softmax_scale_log2}; + const float2 _fma_c2_4 = {-max_scaled, -max_scaled}; + #pragma unroll + for (int _lf = 0; _lf < 32; _lf++) + fma_f32x2_inplace(&reinterpret_cast(_tmem_load_1)[_lf], _fma_b2_3, _fma_c2_4); + #pragma unroll + for (int _le = 0; _le < 64; _le++) { + _tmem_load_1[_le] = approx_exp2(_tmem_load_1[_le]); + } + { + const float2* _reg_reduce_src2_5 = reinterpret_cast(&_tmem_load_1[0]); + float2 _reg_reduce_sum2_5_0 = make_float2(0.0f, 0.0f); + float2 _reg_reduce_sum2_5_1 = make_float2(0.0f, 0.0f); + float2 _reg_reduce_sum2_5_2 = make_float2(0.0f, 0.0f); + float2 _reg_reduce_sum2_5_3 = make_float2(0.0f, 0.0f); + #pragma unroll + for (int _rr = 0; _rr < 8; _rr++) { + add_f32x2_inplace(&_reg_reduce_sum2_5_0, _reg_reduce_src2_5[_rr * 4 + 0]); + add_f32x2_inplace(&_reg_reduce_sum2_5_1, _reg_reduce_src2_5[_rr * 4 + 1]); + add_f32x2_inplace(&_reg_reduce_sum2_5_2, _reg_reduce_src2_5[_rr * 4 + 2]); + add_f32x2_inplace(&_reg_reduce_sum2_5_3, _reg_reduce_src2_5[_rr * 4 + 3]); + } + add_f32x2_inplace(&_reg_reduce_sum2_5_0, _reg_reduce_sum2_5_1); + add_f32x2_inplace(&_reg_reduce_sum2_5_2, _reg_reduce_sum2_5_3); + add_f32x2_inplace(&_reg_reduce_sum2_5_0, _reg_reduce_sum2_5_2); + float _tmem_load_1_sum = _reg_reduce_sum2_5_0.x + _reg_reduce_sum2_5_0.y; + block_sum = _tmem_load_1_sum; + } + uint32_t _fp8_0[16]; + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[0]), "f"(_tmem_load_1[1]), + "f"(_tmem_load_1[2]), "f"(_tmem_load_1[3])); + _fp8_0[0] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[4]), "f"(_tmem_load_1[5]), + "f"(_tmem_load_1[6]), "f"(_tmem_load_1[7])); + _fp8_0[1] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[8]), "f"(_tmem_load_1[9]), + "f"(_tmem_load_1[10]), "f"(_tmem_load_1[11])); + _fp8_0[2] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[12]), "f"(_tmem_load_1[13]), + "f"(_tmem_load_1[14]), "f"(_tmem_load_1[15])); + _fp8_0[3] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[16]), "f"(_tmem_load_1[17]), + "f"(_tmem_load_1[18]), "f"(_tmem_load_1[19])); + _fp8_0[4] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[20]), "f"(_tmem_load_1[21]), + "f"(_tmem_load_1[22]), "f"(_tmem_load_1[23])); + _fp8_0[5] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[24]), "f"(_tmem_load_1[25]), + "f"(_tmem_load_1[26]), "f"(_tmem_load_1[27])); + _fp8_0[6] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[28]), "f"(_tmem_load_1[29]), + "f"(_tmem_load_1[30]), "f"(_tmem_load_1[31])); + _fp8_0[7] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[32]), "f"(_tmem_load_1[33]), + "f"(_tmem_load_1[34]), "f"(_tmem_load_1[35])); + _fp8_0[8] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[36]), "f"(_tmem_load_1[37]), + "f"(_tmem_load_1[38]), "f"(_tmem_load_1[39])); + _fp8_0[9] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[40]), "f"(_tmem_load_1[41]), + "f"(_tmem_load_1[42]), "f"(_tmem_load_1[43])); + _fp8_0[10] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[44]), "f"(_tmem_load_1[45]), + "f"(_tmem_load_1[46]), "f"(_tmem_load_1[47])); + _fp8_0[11] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[48]), "f"(_tmem_load_1[49]), + "f"(_tmem_load_1[50]), "f"(_tmem_load_1[51])); + _fp8_0[12] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[52]), "f"(_tmem_load_1[53]), + "f"(_tmem_load_1[54]), "f"(_tmem_load_1[55])); + _fp8_0[13] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[56]), "f"(_tmem_load_1[57]), + "f"(_tmem_load_1[58]), "f"(_tmem_load_1[59])); + _fp8_0[14] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_1[60]), "f"(_tmem_load_1[61]), + "f"(_tmem_load_1[62]), "f"(_tmem_load_1[63])); + _fp8_0[15] = _packed; + } + int p_stage_base = smem_p_fp8_addr + (unsigned int)(phase * 8192); + #pragma unroll + for (int p_vec = 0; p_vec < 4; p_vec++) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_stage_base + (my_row * 128 + (n_half * 4 + p_vec) * 16 ^ (my_row * 128 + (n_half * 4 + p_vec) * 16 >> 7 & 7) << 4))), "r"(_fp8_0[p_vec * 4]), "r"(_fp8_0[p_vec * 4 + 1]), "r"(_fp8_0[p_vec * 4 + 2]), "r"(_fp8_0[p_vec * 4 + 3]) : "memory"); + } + } + } + { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + { + asm volatile("" ::: "memory"); + } + mbarrier_arrive(p_full_addr + (phase) * 8); + { + asm volatile("" ::: "memory"); + float _fma_3 = __fmaf_rn(row_sum_val, acc_scale, block_sum); + row_sum_val = _fma_3; + } + } + smem_stats_sum[stats_row] = row_sum_val; + smem_stats_final_max[stats_row] = row_max_val; + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(sum_ready_addr); + softmax_tile_cursor = softmax_tile_cursor + (num_kv_tiles_1 - bid % 2 + 1) / 2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_dealloc_addr); + } + } + // ---- Role: softmax_padding ---- + if (warp == 1 || warp == 3) { + { // softmax_padding_main + const int lowhead_softmax_padding = 0; + } + } + // ---- Role: correction_wg ---- + if (warp == 4 || warp == 6) { + { // correction_wg_main + float softmax_scale_log2_1 = bmm1_scale[0] * 1.4426950408889634f; + float output_scale = bmm2_scale[0]; + const int wg_dummy_inc_1 = 0; + int all_num_kv_tiles_3 = (sparse_topk + 128 - 1) / 128; + const int tmem_row_base_1 = ((1) ? warp % 2 * 32 : warp % 4 * 32); + const int n_half_1 = ((1) ? (int)(warp % 4 / 2) : 0); + const int my_row_1 = tmem_row_base_1 + lane; + const int stats_row_1 = n_half_1 * 64 + my_row_1; + const int corr_row = tmem_row_base_1 << 16; + int correction_tile_cursor = 0; + unsigned int _phase_q_full_0_2 = 0; + unsigned int _phase_pv_done_0 = 0; + unsigned int _phase_o_done_0 = 0; + unsigned int _phase_sum_ready_0 = 0; + unsigned int _phase_merge_ready_0 = 0; + unsigned int _phase_merge_full_0 = 0; + #pragma unroll + for (unsigned int work_idx_2 = 0; work_idx_2 < 1; work_idx_2++) { + int split_idx_3 = 0; + int query_idx_3 = bid / 2 >> 1; + int v_chunk_1 = bid / 2 & 1; + { + split_idx_3 = 0; + query_idx_3 = bid / 2; + v_chunk_1 = bid % 2; + } + int tiles_per_split_3 = all_num_kv_tiles_3 + 1 - 1; + int first_tile_3 = split_idx_3 * tiles_per_split_3; + int num_kv_tiles_2 = tiles_per_split_3; + #pragma unroll 1 + for (int tile_1 = 0; tile_1 < (num_kv_tiles_2 - bid % 2 + 1) / 2; tile_1++) { + int pipeline_tile_1 = correction_tile_cursor + tile_1; + int phase_1 = pipeline_tile_1 & 1; + int stats_wait_phase = pipeline_tile_1 >> 1 & 1; + mbarrier_wait(stats_addr + (phase_1) * 8, stats_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + float acc_scale_1 = smem_stats_max[phase_1 * 128 + stats_row_1]; + if (tile_1 > 0) { + mbarrier_wait(pv_done_addr, _phase_pv_done_0); + _phase_pv_done_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + } + if (tile_1 > 0) { + int _vote_1 = __any_sync(0xFFFFFFFF, acc_scale_1 < 1.0f); + int any_rescale = _vote_1; + if (any_rescale != 0) { + #pragma unroll + for (int vs = 0; vs < 2; vs++) { + int o_base = taddr + 256 + (unsigned int)(vs * 128) + (unsigned int)corr_row; + #pragma unroll + for (int c = 0; c < 128; c += 16) { + float _tmem_load_4[16]; + tmem_ld_x16(&_tmem_load_4[0], o_base + c); + const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + mul_f32x2_inplace(&reinterpret_cast(_tmem_load_4)[_ls], _scale2_0); + tmem_st_x16_f32(o_base + c, _tmem_load_4); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + } + } + } + } + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(corr_done_addr + (phase_1) * 8); + } + mbarrier_wait(o_done_addr, _phase_o_done_0); + _phase_o_done_0 ^= 1; + mbarrier_wait(sum_ready_addr, _phase_sum_ready_0); + _phase_sum_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + float wt_sum = smem_stats_sum[stats_row_1]; + float wt_max = smem_stats_final_max[stats_row_1]; + smem_corr_warp_pair_exchange[stats_row_1] = wt_sum; + asm volatile("barrier.sync %0, 64;" :: "r"(4 + warp % 2) : "memory"); + wt_sum = wt_sum + smem_corr_warp_pair_exchange[stats_row_1 ^ 64]; + int wt_peer_rank = bid % 2 ^ 1; + int wt_quadrant = warp - 4; + if (warp == 4) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(merge_full_addr, 66048); + } + } + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cluster.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(merge_ready_addr), "r"(wt_peer_rank) : "memory"); + if (n_half_1 == 0) { + smem_merge_stats_send[my_row_1] = wt_max; + smem_merge_stats_send[64 + my_row_1] = wt_sum; + } + int wt_send_block0 = (bid % 2 ^ 1) * 2; + #pragma unroll + for (int wt_b = 0; wt_b < 2; wt_b++) { + #pragma unroll + for (int wt_c = 0; wt_c < 64; wt_c += 32) { + float _tmem_load_5[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_5[0]), "=f"(_tmem_load_5[1]), "=f"(_tmem_load_5[2]), "=f"(_tmem_load_5[3]), "=f"(_tmem_load_5[4]), "=f"(_tmem_load_5[5]), "=f"(_tmem_load_5[6]), "=f"(_tmem_load_5[7]), "=f"(_tmem_load_5[8]), "=f"(_tmem_load_5[9]), "=f"(_tmem_load_5[10]), "=f"(_tmem_load_5[11]), "=f"(_tmem_load_5[12]), "=f"(_tmem_load_5[13]), "=f"(_tmem_load_5[14]), "=f"(_tmem_load_5[15]), "=f"(_tmem_load_5[16]), "=f"(_tmem_load_5[17]), "=f"(_tmem_load_5[18]), "=f"(_tmem_load_5[19]), "=f"(_tmem_load_5[20]), "=f"(_tmem_load_5[21]), "=f"(_tmem_load_5[22]), "=f"(_tmem_load_5[23]), "=f"(_tmem_load_5[24]), "=f"(_tmem_load_5[25]), "=f"(_tmem_load_5[26]), "=f"(_tmem_load_5[27]), "=f"(_tmem_load_5[28]), "=f"(_tmem_load_5[29]), "=f"(_tmem_load_5[30]), "=f"(_tmem_load_5[31]) + : "r"(taddr + 256 + (unsigned int)((wt_send_block0 + wt_b) * 64) + (unsigned int)wt_c + (unsigned int)corr_row)); + #pragma unroll + for (int wt_k = 0; wt_k < 32; wt_k += 4) { + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(smem_merge_send_addr + (unsigned int)((((wt_b * 64 + wt_c + wt_k) / 4 * 4 + wt_quadrant) * 32 + lane) * 16)), "f"(_tmem_load_5[wt_k]), "f"(_tmem_load_5[wt_k + 1]), "f"(_tmem_load_5[wt_k + 2]), "f"(_tmem_load_5[wt_k + 3]) : "memory"); + } + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("barrier.sync 6, 64;" ::: "memory"); + if (warp == 4) { + if (elect_sync()) { + mbarrier_wait(merge_ready_addr, _phase_merge_ready_0); + _phase_merge_ready_0 ^= 1; + uint32_t _mapa_0; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mapa_0) : "r"(smem_merge_recv_addr), "r"(wt_peer_rank)); + uint32_t _mapa_1; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mapa_1) : "r"(merge_full_addr), "r"(wt_peer_rank)); + asm volatile( + "cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes" + " [%0], [%1], %2, [%3];" + :: "r"(_mapa_0), "r"(smem_merge_send_addr), "r"((uint32_t)(65536)), "r"(_mapa_1) + : "memory"); + uint32_t _mapa_2; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mapa_2) : "r"(smem_merge_stats_recv_addr), "r"(wt_peer_rank)); + uint32_t _mapa_3; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mapa_3) : "r"(merge_full_addr), "r"(wt_peer_rank)); + asm volatile( + "cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes" + " [%0], [%1], %2, [%3];" + :: "r"(_mapa_2), "r"(smem_merge_stats_send_addr), "r"((uint32_t)(512)), "r"(_mapa_3) + : "memory"); + } + } + mbarrier_wait(merge_full_addr, _phase_merge_full_0); + _phase_merge_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + float wt_peer_max = smem_merge_stats_recv[my_row_1]; + float wt_peer_sum = smem_merge_stats_recv[64 + my_row_1]; + float _max_5 = max_noftz(wt_max, wt_peer_max); + float wt_merged_max = _max_5; + float _exp2_1 = approx_exp2((wt_max - wt_merged_max) * softmax_scale_log2_1); + float wt_s_own = ((wt_max > -CAKE_INF) ? _exp2_1 : 0.0f); + float _exp2_2 = approx_exp2((wt_peer_max - wt_merged_max) * softmax_scale_log2_1); + float wt_s_peer = ((wt_peer_max > -CAKE_INF) ? _exp2_2 : 0.0f); + float wt_merged_sum = wt_sum * wt_s_own + wt_peer_sum * wt_s_peer; + float _rcp_0 = approx_rcp(wt_merged_sum); + float wt_inv_sum = ((wt_merged_sum > 0.0f) ? _rcp_0 : 0.0f); + int wt_head_idx = my_row_1; + int wt_o_base = (query_idx_3 * num_heads + wt_head_idx) * 512; + if (bid % 2 == 0 && n_half_1 == 0 && wt_head_idx < num_heads) { + float _log2_0; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(wt_merged_sum)); + partial_lse[query_idx_3 * num_heads + wt_head_idx + split_idx_3] = ((wt_merged_sum > 0.0f) ? wt_merged_max * softmax_scale_log2_1 + _log2_0 : -CAKE_INF); + } + int wt_keep_block0 = bid % 2 * 2; + #pragma unroll + for (int wt_b_1 = 0; wt_b_1 < 2; wt_b_1++) { + #pragma unroll + for (int wt_c_1 = 0; wt_c_1 < 64; wt_c_1 += 32) { + float _tmem_load_6[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(_tmem_load_6[0]), "=f"(_tmem_load_6[1]), "=f"(_tmem_load_6[2]), "=f"(_tmem_load_6[3]), "=f"(_tmem_load_6[4]), "=f"(_tmem_load_6[5]), "=f"(_tmem_load_6[6]), "=f"(_tmem_load_6[7]), "=f"(_tmem_load_6[8]), "=f"(_tmem_load_6[9]), "=f"(_tmem_load_6[10]), "=f"(_tmem_load_6[11]), "=f"(_tmem_load_6[12]), "=f"(_tmem_load_6[13]), "=f"(_tmem_load_6[14]), "=f"(_tmem_load_6[15]), "=f"(_tmem_load_6[16]), "=f"(_tmem_load_6[17]), "=f"(_tmem_load_6[18]), "=f"(_tmem_load_6[19]), "=f"(_tmem_load_6[20]), "=f"(_tmem_load_6[21]), "=f"(_tmem_load_6[22]), "=f"(_tmem_load_6[23]), "=f"(_tmem_load_6[24]), "=f"(_tmem_load_6[25]), "=f"(_tmem_load_6[26]), "=f"(_tmem_load_6[27]), "=f"(_tmem_load_6[28]), "=f"(_tmem_load_6[29]), "=f"(_tmem_load_6[30]), "=f"(_tmem_load_6[31]) + : "r"(taddr + 256 + (unsigned int)((wt_keep_block0 + wt_b_1) * 64) + (unsigned int)wt_c_1 + (unsigned int)corr_row)); + float wt_peer[32]; + #pragma unroll + for (int wt_k_1 = 0; wt_k_1 < 32; wt_k_1 += 4) { + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&wt_peer[wt_k_1])), "=r"(*reinterpret_cast(&wt_peer[(wt_k_1) + 1])), "=r"(*reinterpret_cast(&wt_peer[(wt_k_1) + 2])), "=r"(*reinterpret_cast(&wt_peer[(wt_k_1) + 3])) + : "r"(smem_merge_recv_addr + (unsigned int)((((wt_b_1 * 64 + wt_c_1 + wt_k_1) / 4 * 4 + wt_quadrant) * 32 + lane) * 16))); + } + const float2 _scale2_1 = {wt_s_own, wt_s_own}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + mul_f32x2_inplace(&reinterpret_cast(_tmem_load_6)[_ls], _scale2_1); + const float2 _scale2_2 = {wt_s_peer, wt_s_peer}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + mul_f32x2_inplace(&reinterpret_cast(wt_peer)[_ls], _scale2_2); + #pragma unroll + for (int _la = 0; _la < 16; _la++) + add_f32x2_inplace(&reinterpret_cast(_tmem_load_6)[_la], reinterpret_cast(wt_peer)[_la]); + int wt_gmem = wt_o_base + (wt_keep_block0 + wt_b_1) * 128 + n_half_1 * 64 + wt_c_1; + if (wt_head_idx < num_heads) { + #pragma unroll + for (int wt_j = 0; wt_j < 32; wt_j += 8) { + { + const float2 _prescale2_3 = {wt_inv_sum * output_scale, wt_inv_sum * output_scale}; + #if __CUDA_ARCH__ >= 1000 + #pragma unroll + for (int _ps = 0; _ps < 4; _ps++) + mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_6[wt_j])[_ps], _prescale2_3); + #else + #pragma unroll + for (int _ps = 0; _ps < 8; _ps++) + _tmem_load_6[wt_j + _ps] *= wt_inv_sum * output_scale; + #endif + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_tmem_load_6[wt_j + 0], _tmem_load_6[wt_j + 1]); + _pk[1] = __floats2bfloat162_rn(_tmem_load_6[wt_j + 2], _tmem_load_6[wt_j + 3]); + _pk[2] = __floats2bfloat162_rn(_tmem_load_6[wt_j + 4], _tmem_load_6[wt_j + 5]); + _pk[3] = __floats2bfloat162_rn(_tmem_load_6[wt_j + 6], _tmem_load_6[wt_j + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(O + (wt_gmem + wt_j)))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + } + } + correction_tile_cursor = correction_tile_cursor + (num_kv_tiles_2 - bid % 2 + 1) / 2; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(tmem_dealloc_addr); + } + } + // ---- Role: correction_padding ---- + if (warp == 5 || warp == 7) { + { // correction_padding_main + const int lowhead_correction_padding = 0; + } + } + // ---- Role: mma_warp ---- + if (warp == 8) { + { // mma_warp_main + const int wg2_dummy_1 = 0; + int all_num_kv_tiles_4 = (sparse_topk + 128 - 1) / 128; + unsigned int mma_k_stage = 0; + unsigned int mma_v_stage = 0; + unsigned int mma_kv_stage = 0; + int mma_tile_cursor = 0; + unsigned int _phase_q_full_0_3 = 0; + unsigned int _phase_k_full = 0; + #pragma unroll + for (unsigned int work_idx_3 = 0; work_idx_3 < 1; work_idx_3++) { + int split_idx_4 = 0; + { + split_idx_4 = 0; + } + int tiles_per_split_4 = all_num_kv_tiles_4 + 1 - 1; + int first_tile_4 = split_idx_4 * tiles_per_split_4; + int num_kv_tiles_3 = tiles_per_split_4; + int wt_num_local_tiles_1 = (num_kv_tiles_3 - bid % 2 + 1) / 2; + mbarrier_wait(q_full_addr, _phase_q_full_0_3); + _phase_q_full_0_3 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int wt_first_pv = 1; + #pragma unroll 1 + for (int wt_tile_1 = 0; wt_tile_1 < wt_num_local_tiles_1; wt_tile_1++) { + int wt_pipeline_tile = mma_tile_cursor + wt_tile_1; + int wt_phase = wt_pipeline_tile & 1; + int wt_score_col = ((wt_phase != 0) ? 128 : 0); + mbarrier_wait(k_full_addr + (mma_k_stage) * 8, _phase_k_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_0 = make_warp_uniform((((smem_q_full_addr) >> 4) & 0x3FFF) + (0) * 2048); + int _mma_b_lo_0 = make_warp_uniform((((smem_k_full_addr) >> 4) & 0x3FFF) + (mma_k_stage) * 4096); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 69206032;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (wt_score_col))), "r"(0)); + int _mma_a_lo_1 = make_warp_uniform((((smem_q_full_addr + 8192) >> 4) & 0x3FFF) + (0) * 2048); + int _mma_b_lo_1 = make_warp_uniform((((smem_k_full_addr + 16384) >> 4) & 0x3FFF) + (mma_k_stage) * 4096); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 69206032;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_tmem_scratch + (wt_score_col))), "r"(1)); + int _mma_a_lo_2 = make_warp_uniform((((smem_q_full_addr + 16384) >> 4) & 0x3FFF) + (0) * 2048); + int _mma_b_lo_2 = make_warp_uniform((((smem_k_full_addr + 32768) >> 4) & 0x3FFF) + (mma_k_stage) * 4096); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 69206032;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (wt_score_col))), "r"(1)); + int _mma_a_lo_3 = make_warp_uniform((((smem_q_full_addr + 24576) >> 4) & 0x3FFF) + (0) * 2048); + int _mma_b_lo_3 = make_warp_uniform((((smem_k_full_addr + 49152) >> 4) & 0x3FFF) + (mma_k_stage) * 4096); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 69206032;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"((tmem_tmem_scratch + (wt_score_col))), "r"(1)); + elect_commit(s_full_addr + (wt_phase) * 8); + if (wt_tile_1 > 0) { + int wt_prev_pipeline_tile = wt_pipeline_tile - 1; + int wt_prev_phase = wt_prev_pipeline_tile & 1; + int wt_pv_wait_phase = wt_prev_pipeline_tile >> 1 & 1; + mbarrier_wait(p_full_addr + (wt_prev_phase) * 8, wt_pv_wait_phase); + mbarrier_wait(corr_done_addr + (wt_prev_phase) * 8, wt_pv_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + unsigned int wt_prev_kv_stage = mma_k_stage ^ 1; + int _mma_a_lo_4 = make_warp_uniform((((smem_p_fp8_addr) >> 4) & 0x3FFF) + (wt_prev_phase) * 512); + int _mma_b_lo_4 = make_warp_uniform(((((smem_v_full_addr) >> 4) & 0x3FFF) | 0x4000000) + (wt_prev_kv_stage) * 4096); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 69271568;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"((tmem_tmem_scratch + (256))), "r"(((wt_first_pv) ? 0 : 1))); + int _mma_a_lo_5 = make_warp_uniform((((smem_p_fp8_addr) >> 4) & 0x3FFF) + (wt_prev_phase) * 512); + int _mma_b_lo_5 = make_warp_uniform(((((smem_v_full_addr + 16384) >> 4) & 0x3FFF) | 0x4000000) + (wt_prev_kv_stage) * 4096); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 69271568;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_5), "r"(_mma_b_lo_5), "r"((tmem_tmem_scratch + (320))), "r"(((wt_first_pv) ? 0 : 1))); + int _mma_a_lo_6 = make_warp_uniform((((smem_p_fp8_addr) >> 4) & 0x3FFF) + (wt_prev_phase) * 512); + int _mma_b_lo_6 = make_warp_uniform(((((smem_v_full_addr + 32768) >> 4) & 0x3FFF) | 0x4000000) + (wt_prev_kv_stage) * 4096); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 69271568;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_6), "r"(_mma_b_lo_6), "r"((tmem_tmem_scratch + (384))), "r"(((wt_first_pv) ? 0 : 1))); + int _mma_a_lo_7 = make_warp_uniform((((smem_p_fp8_addr) >> 4) & 0x3FFF) + (wt_prev_phase) * 512); + int _mma_b_lo_7 = make_warp_uniform(((((smem_v_full_addr + 49152) >> 4) & 0x3FFF) | 0x4000000) + (wt_prev_kv_stage) * 4096); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 69271568;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_7), "r"(_mma_b_lo_7), "r"((tmem_tmem_scratch + (448))), "r"(((wt_first_pv) ? 0 : 1))); + wt_first_pv = 0; + elect_commit(k_empty_addr + (wt_prev_kv_stage) * 8); + elect_commit(pv_done_addr); + } + mma_k_stage += 1; + if (mma_k_stage == 2) { mma_k_stage = 0; _phase_k_full ^= 1; } + } + if (wt_num_local_tiles_1 > 0) { + int wt_last_pipeline_tile = mma_tile_cursor + wt_num_local_tiles_1 - 1; + int wt_last_phase = wt_last_pipeline_tile & 1; + int wt_drain_wait_phase = wt_last_pipeline_tile >> 1 & 1; + mbarrier_wait(p_full_addr + (wt_last_phase) * 8, wt_drain_wait_phase); + mbarrier_wait(corr_done_addr + (wt_last_phase) * 8, wt_drain_wait_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + unsigned int wt_last_kv_stage = mma_k_stage ^ 1; + int _mma_a_lo_8 = make_warp_uniform((((smem_p_fp8_addr) >> 4) & 0x3FFF) + (wt_last_phase) * 512); + int _mma_b_lo_8 = make_warp_uniform(((((smem_v_full_addr) >> 4) & 0x3FFF) | 0x4000000) + (wt_last_kv_stage) * 4096); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 69271568;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_8), "r"(_mma_b_lo_8), "r"((tmem_tmem_scratch + (256))), "r"(((wt_first_pv) ? 0 : 1))); + int _mma_a_lo_9 = make_warp_uniform((((smem_p_fp8_addr) >> 4) & 0x3FFF) + (wt_last_phase) * 512); + int _mma_b_lo_9 = make_warp_uniform(((((smem_v_full_addr + 16384) >> 4) & 0x3FFF) | 0x4000000) + (wt_last_kv_stage) * 4096); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 69271568;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_9), "r"(_mma_b_lo_9), "r"((tmem_tmem_scratch + (320))), "r"(((wt_first_pv) ? 0 : 1))); + int _mma_a_lo_10 = make_warp_uniform((((smem_p_fp8_addr) >> 4) & 0x3FFF) + (wt_last_phase) * 512); + int _mma_b_lo_10 = make_warp_uniform(((((smem_v_full_addr + 32768) >> 4) & 0x3FFF) | 0x4000000) + (wt_last_kv_stage) * 4096); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 69271568;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_10), "r"(_mma_b_lo_10), "r"((tmem_tmem_scratch + (384))), "r"(((wt_first_pv) ? 0 : 1))); + int _mma_a_lo_11 = make_warp_uniform((((smem_p_fp8_addr) >> 4) & 0x3FFF) + (wt_last_phase) * 512); + int _mma_b_lo_11 = make_warp_uniform(((((smem_v_full_addr + 49152) >> 4) & 0x3FFF) | 0x4000000) + (wt_last_kv_stage) * 4096); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 69271568;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.ws.cta_group::1.kind::f8f6f4 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_11), "r"(_mma_b_lo_11), "r"((tmem_tmem_scratch + (448))), "r"(((wt_first_pv) ? 0 : 1))); + elect_commit(k_empty_addr + (wt_last_kv_stage) * 8); + } + elect_commit(q_empty_addr); + elect_commit(o_done_addr); + mma_tile_cursor = mma_tile_cursor + wt_num_local_tiles_1; + } + mbarrier_arrive(tmem_dealloc_addr); + unsigned int _phase_tmem_dealloc_0 = 0; + mbarrier_wait(tmem_dealloc_addr, _phase_tmem_dealloc_0); + _phase_tmem_dealloc_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + } + // ---- Role: empty1 ---- + if (warp >= 10 && warp <= 11) { + { // empty1_main + const int wg2_dummy_2 = 0; + unsigned int _phase_q_full_0_4 = 0; + } + } + + // Cleanup +} + +} // extern "C" diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_bab2852cc0ea4a6d4860_binding.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_bab2852cc0ea4a6d4860_binding.cu new file mode 100644 index 00000000000..a3213ad3a76 --- /dev/null +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_bab2852cc0ea4a6d4860_binding.cu @@ -0,0 +1,441 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +// Generated by the Cake source exporter — do not edit. +// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_bab2852cc0ea4a6d4860'. +#include +#include +#include +#include + +#include "tvm_ffi_utils.h" + +#include +#include +#include +#include +#include +#include +#include + +extern "C" __global__ void kernel_cake_dsv4_bab2852cc0ea4a6d4860(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ cum_seq_lens_q, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); + + +namespace cake_host_shim_66b5cf782c79bc55 { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t CakeDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +// 4D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + int64_t d3 = t.size(t.ndim() - 3); + TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + TVM_FFI_CHECK(d2 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d2 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(d1), (uint64_t)(64), (uint64_t)((d2 / 64)), (uint64_t)(d3)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (128, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = (64 * d1); + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + (uint64_t)((carrier_stride_1 * 8) / 8), + (uint64_t)((carrier_stride_2 * 8) / 8), + }; + uint32_t box_dim[4] = {128u, 64u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_L2_128B, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_swa_kv' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_swa_kv' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; + TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (128, 1) exceeds resolved global dims for 'tmap_swa_kv'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + }; + uint32_t box_dim[2] = {128u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_compressed_kv' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_compressed_kv(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_compressed_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_compressed_kv' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_compressed_kv' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_compressed_kv"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_compressed_kv' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved a non-positive global dim"; + TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (128, 1) exceeds resolved global dims for 'tmap_compressed_kv'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 8) / 8), + }; + uint32_t box_dim[2] = {128u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_kv') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_compressed_kv, TensorView arg_O, TensorView arg_cum_seq_lens_q, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 1, 8, 1); + CheckContiguous(arg_tmap_q, "tmap_q"); + CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); + CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 1, 8, 1); + CheckCudaTensor(arg_tmap_compressed_kv, "tmap_compressed_kv"); + CheckDtype(arg_tmap_compressed_kv, "tmap_compressed_kv", 1, 8, 1); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_cum_seq_lens_q, "cum_seq_lens_q"); + CheckDtype(arg_cum_seq_lens_q, "cum_seq_lens_q", 0, 32, 1); + CheckContiguous(arg_cum_seq_lens_q, "cum_seq_lens_q"); + CheckCudaTensor(arg_swa_indices, "swa_indices"); + CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); + CheckContiguous(arg_swa_indices, "swa_indices"); + CheckCudaTensor(arg_compressed_indices, "compressed_indices"); + CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); + CheckContiguous(arg_compressed_indices, "compressed_indices"); + CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); + CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); + CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); + CheckContiguous(arg_bmm1_scale, "bmm1_scale"); + CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); + CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); + CheckContiguous(arg_bmm2_scale, "bmm2_scale"); + TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) + << "scalar 'num_heads' value " << arg_num_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) + << "scalar 'swa_index_stride' value " << arg_swa_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) + << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) + << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) + << "scalar 'sparse_topk' value " << arg_sparse_topk + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) + << "scalar 'has_sinks' value " << arg_has_sinks + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); + CheckSameCudaDevice(arg_tmap_compressed_kv, arg_tmap_q, "tmap_compressed_kv", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_cum_seq_lens_q, arg_tmap_q, "cum_seq_lens_q", "tmap_q"); + CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); + CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); + CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); + CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + + cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); + CUtensorMap p_tmap_compressed_kv = EncodeTma_tmap_compressed_kv(arg_tmap_compressed_kv); + __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); + int* p_cum_seq_lens_q = static_cast(arg_cum_seq_lens_q.data_ptr()); + int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); + int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); + int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); + float* p_sinks = static_cast(arg_sinks.data_ptr()); + float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); + float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); + int32_t v_num_heads = (int32_t)arg_num_heads; + int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; + int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; + int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; + int32_t v_sparse_topk = (int32_t)arg_sparse_topk; + int32_t v_has_sinks = (int32_t)arg_has_sinks; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_tmap_compressed_kv, &p_O, &p_cum_seq_lens_q, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_has_sinks, &v_total_work_items}; + + static std::mutex smem_mu; + static auto* smem_status_by_device = + new std::unordered_map(); + cudaError_t smem_status = cudaSuccess; + { + std::lock_guard lock(smem_mu); + auto it = smem_status_by_device->find(dev.device_id); + if (it == smem_status_by_device->end()) { + int current_device = -1; + smem_status = cudaGetDevice(¤t_device); + if (smem_status == cudaSuccess && current_device != dev.device_id) { + smem_status = cudaErrorInvalidDevice; + } + if (smem_status == cudaSuccess) { + smem_status = cudaFuncSetAttribute( + reinterpret_cast(kernel_cake_dsv4_bab2852cc0ea4a6d4860), + cudaFuncAttributeMaxDynamicSharedMemorySize, + 203136); + } + smem_status_by_device->emplace(dev.device_id, smem_status); + } else { + smem_status = it->second; + } + } + TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) + << "cudaFuncSetAttribute for kernel_cake_dsv4_bab2852cc0ea4a6d4860 failed: " + << cudaGetErrorString(smem_status); + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(512u, 1u, 1u); + + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes = 203136u; + config.stream = stream; + cudaError_t launch_status = cudaLaunchKernelExC( + &config, reinterpret_cast(kernel_cake_dsv4_bab2852cc0ea4a6d4860), kargs); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernelExC for kernel_cake_dsv4_bab2852cc0ea4a6d4860 failed: " + << cudaGetErrorString(launch_status); +} + +} // namespace cake_host_shim_66b5cf782c79bc55 + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_66b5cf782c79bc55::Run); diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_bab2852cc0ea4a6d4860_kernel.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_bab2852cc0ea4a6d4860_kernel.cu new file mode 100644 index 00000000000..cb8f8e0f09e --- /dev/null +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_bab2852cc0ea4a6d4860_kernel.cu @@ -0,0 +1,2968 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; +struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); +#include +#include + +#define CAKE_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_TMEM_OFFSET 0 +#define NUM_Q_PIPE_STAGES 2 +#define NUM_KV_PIPE_STAGES 4 +#define NUM_PAGE_PIPE_STAGES 6 +#define NUM_SCORE_PIPE_STAGES 2 +#define NUM_STATS_PIPE_STAGES 2 +#define NUM_WORK_PIPE_STAGES 2 +#define NUM_THROTTLE_PIPE_STAGES 2 +#define SMEM_SMEM_Q_OFF 0 +#define SMEM_SMEM_Q_STAGE_BYTES 32768 +#define SMEM_SMEM_Q_STRIDE 32768 +#define SMEM_SMEM_K_OFF 65536 +#define SMEM_SMEM_K_STAGE_BYTES 32768 +#define SMEM_SMEM_K_STRIDE 32768 +#define SMEM_SMEM_V_OFF 65536 +#define SMEM_SMEM_V_STAGE_BYTES 32768 +#define SMEM_SMEM_V_STRIDE 32768 +#define SMEM_SMEM_PAGE_OFFSETS_OFF 196608 +#define SMEM_SMEM_PAGE_OFFSETS_STAGE_BYTES 1024 +#define SMEM_SMEM_PAGE_OFFSETS_STRIDE 1024 +#define SMEM_WORK_RESPONSE_OFF 202752 +#define SMEM_WORK_RESPONSE_STAGE_BYTES 16 +#define SMEM_WORK_RESPONSE_STRIDE 16 +#define SMEM_TOTAL 203136 +#define THREADS 512 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { + asm volatile("mbarrier.init.b64 [%0], %1;" + :: "l"(mbar_addr), "r"(count) : "memory"); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Source-faithful relaxed CTA wait used only by a typed protocol that does +// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. +__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, 10000000;\n\t" + "@P1 bra.uni DONE_RELAXED;\n\t" + "bra.uni LAB_WAIT_RELAXED;\n\t" + "DONE_RELAXED:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +// Exact source ports may request the PTX suspendTimeHint operand explicitly. +// The hint is expressed in nanoseconds and is kept separate from the canonical +// no-hint CTA helper so unrelated schedules retain their existing retry path. +__device__ __forceinline__ void mbarrier_wait_suspend( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_SUSPEND:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_SUSPEND;\n\t" + "bra.uni LAB_WAIT_SUSPEND;\n\t" + "DONE_SUSPEND:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + ".reg .u32 WAIT_ADDR;\n\t" + "mov.u32 WAIT_ADDR, %0;\n\t" + "LAB_WAIT_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [WAIT_ADDR], %1, %2;\n\t" + "@P1 bra.uni DONE_HINT;\n\t" + "bra.uni LAB_WAIT_HINT;\n\t" + "DONE_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +// Exact unqualified CTA wait used by source schedules whose PTX intentionally +// omits the acquire qualifier while retaining a typed suspendTimeHint operand. +__device__ __forceinline__ void mbarrier_wait_relaxed_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_RELAXED_HINT:\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra DONE_RELAXED_HINT;\n\t" + "bra LAB_WAIT_RELAXED_HINT;\n\t" + "DONE_RELAXED_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_suspend( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( + int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { + if (token == 0) { + mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f8f6f4( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d) + : "memory"); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + +__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { + MmaSmemDesc tmp; + tmp.u64 = smem_desc; + tmp.u32[0] += offset; + smem_desc = tmp.u64; +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float warp_reduce_max(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + return val; +} + + +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ float2 ex2_emulation_f32x2_value(float2 value) { + const float c0 = 1.0f, c1 = 0.695146143436431884765625f; + const float c2 = 0.227564394474029541015625f, c3 = 0.077119089663028717041015625f; + const float magic = 12582912.0f; + float x0 = max_noftz(value.x, -127.0f), x1 = max_noftz(value.y, -127.0f); + float2 xc2 = make_float2(x0, x1), magic2 = make_float2(magic, magic); + float2 xr2; + asm("add.rm.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xr2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&magic2)); + float2 c3_2 = make_float2(c3, c3), c2_2 = make_float2(c2, c2); + float2 c1_2 = make_float2(c1, c1), c0_2 = make_float2(c0, c0); + float2 xrb2, xfrac2; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xrb2) + : "l"(*(unsigned long long*)&xr2), "l"(*(unsigned long long*)&magic2)); + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xfrac2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&xrb2)); + float2 poly2; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&c3_2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c2_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c1_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c0_2)); + int x0r_i, x1r_i, p0_i, p1_i; + asm("mov.b64 {%0, %1}, %2;" : "=r"(x0r_i), "=r"(x1r_i) : "l"(*(unsigned long long*)&xr2)); + asm("mov.b64 {%0, %1}, %2;" : "=r"(p0_i), "=r"(p1_i) : "l"(*(unsigned long long*)&poly2)); + float r0, r1; + asm("mov.b32 %0, %1;" : "=f"(r0) : "r"((x0r_i << 23) + p0_i)); + asm("mov.b32 %0, %1;" : "=f"(r1) : "r"((x1r_i << 23) + p1_i)); + return make_float2(r0, r1); +} + +__device__ __forceinline__ void ex2_emulation_f32x2(float* x0_ptr, float* x1_ptr) { + float2 result = ex2_emulation_f32x2_value(make_float2(*x0_ptr, *x1_ptr)); + *x0_ptr = result.x; *x1_ptr = result.y; +} + +__device__ __forceinline__ void softmax_frag_exp2_cast( + float* sv, uint32_t* pv, int use_emu) +{ + #pragma unroll + for (int j = 0; j < 16; j++) { + if (use_emu && j >= 12) + ex2_emulation_f32x2(&sv[j*2], &sv[j*2+1]); + else { + sv[j*2] = approx_exp2(sv[j*2]); + sv[j*2+1] = approx_exp2(sv[j*2+1]); + } + } + #pragma unroll + for (int j = 0; j < 16; j++) { + __nv_bfloat162 bf = __float22bfloat162_rn({sv[j*2], sv[j*2+1]}); + pv[j] = reinterpret_cast(bf); + } +} + + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { + asm("add.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { + asm("sub.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("add.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("sub.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { + float2 r; + asm("mul.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("add.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("add.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("add.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("add.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("add.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("add.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("add.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rz.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rm.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { + float2 r; + asm("mul.rp.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_2d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3}], [%4];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_gather4_gmem2smem( + int dst, const void *tmap_ptr, + int col_idx, int row0, int row1, int row2, int row3, + int mbar_addr) { + // Canonical .shared::cta form for non-multicast gather4, matching + // trtllm-gen / cuda_ptx and the PTX ISA qualifier order + // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, + // .shared::cluster is reserved for the multicast variant (ctaMask). + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" + :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), + "r"(row0), "r"(row1), "r"(row2), "r"(row3), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + + +template +struct CakePipelineState { + uint32_t index_; + uint32_t phase_; + uint32_t count_; + + __device__ __forceinline__ uint32_t index() const { return index_; } + __device__ __forceinline__ uint32_t phase() const { return phase_; } + __device__ __forceinline__ uint32_t count() const { return count_; } + __device__ __forceinline__ CakePipelineState& operator++() { + ++count_; + ++index_; + if (index_ == Stages) { + index_ = 0; + phase_ ^= 1; + } + return *this; + } +}; + +template +__device__ __forceinline__ CakePipelineState +cake_pipeline_state_at(CakePipelineState state) { + #pragma unroll + for (int step = 0; step < Offset; ++step) { + ++state; + } + return state; +} + +extern "C" { + +__global__ __launch_bounds__(512, 1) void +kernel_cake_dsv4_bab2852cc0ea4a6d4860(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ cum_seq_lens_q, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); + + const int mbar_base = smem + 202784; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 16) + #define kv_full_addr (mbar_base + 32) + #define kv_empty_addr (mbar_base + 64) + #define page_full_addr (mbar_base + 96) + #define page_empty_addr (mbar_base + 144) + #define score_full_addr (mbar_base + 192) + #define score_empty_addr (mbar_base + 208) + #define stats_full_addr (mbar_base + 224) + #define stats_empty_addr (mbar_base + 240) + #define o_full_addr (mbar_base + 256) + #define o_empty_addr (mbar_base + 264) + #define work_full_addr (mbar_base + 272) + #define work_empty_addr (mbar_base + 288) + #define throttle_full_addr (mbar_base + 304) + #define throttle_empty_addr (mbar_base + 320) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + + // Kernel setup ops + uint8_t* smem_q = reinterpret_cast(smem_raw + 0); + const int smem_q_addr = smem + 0; + uint8_t* smem_k = reinterpret_cast(smem_raw + 65536); + const int smem_k_addr = smem + 65536; + uint8_t* smem_v = reinterpret_cast(smem_raw + 65536); + const int smem_v_addr = smem + 65536; + int* smem_page_offsets = reinterpret_cast(smem_raw + 196608); + const int smem_page_offsets_addr = smem + 196608; + unsigned int* work_response = reinterpret_cast(smem_raw + 202752); + const int work_response_addr = smem + 202752; + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_q))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_swa_kv))) : "memory"); + asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_compressed_kv))) : "memory"); + + // Mbarrier init (16 pipeline groups, 0 ordered-sequence groups, 42 barriers) + // Mbarriers at smem_raw[202784..203120) + + if (warp == 12) { + uint32_t leader = elect_sync(); + if (leader) { + // --- pipeline 'q_pipe' --- + // q_full: 2 barriers, init_count=1 + mbarrier_init(smem + 202784, 1); + mbarrier_init(smem + 202792, 1); + // q_empty: 2 barriers, init_count=1 + mbarrier_init(smem + 202800, 1); + mbarrier_init(smem + 202808, 1); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // --- pipeline 'kv_pipe' --- + // kv_full: 4 barriers, init_count=4 + mbarrier_init(smem + 202816, 4); + mbarrier_init(smem + 202824, 4); + mbarrier_init(smem + 202832, 4); + mbarrier_init(smem + 202840, 4); + // kv_empty: 4 barriers, init_count=1 + mbarrier_init(smem + 202848, 1); + mbarrier_init(smem + 202856, 1); + mbarrier_init(smem + 202864, 1); + mbarrier_init(smem + 202872, 1); + // --- pipeline 'page_pipe' --- + // page_full: 6 barriers, init_count=32 + mbarrier_init(smem + 202880, 32); + mbarrier_init(smem + 202888, 32); + mbarrier_init(smem + 202896, 32); + mbarrier_init(smem + 202904, 32); + mbarrier_init(smem + 202912, 32); + mbarrier_init(smem + 202920, 32); + // page_empty: 6 barriers, init_count=128 + mbarrier_init(smem + 202928, 128); + mbarrier_init(smem + 202936, 128); + mbarrier_init(smem + 202944, 128); + mbarrier_init(smem + 202952, 128); + mbarrier_init(smem + 202960, 128); + mbarrier_init(smem + 202968, 128); + // --- pipeline 'score_pipe' --- + // score_full: 2 barriers, init_count=1 + mbarrier_init(smem + 202976, 1); + mbarrier_init(smem + 202984, 1); + // score_empty: 2 barriers, init_count=128 + mbarrier_init(smem + 202992, 128); + mbarrier_init(smem + 203000, 128); + // --- pipeline 'stats_pipe' --- + // stats_full: 2 barriers, init_count=128 + mbarrier_init(smem + 203008, 128); + mbarrier_init(smem + 203016, 128); + // stats_empty: 2 barriers, init_count=128 + mbarrier_init(smem + 203024, 128); + mbarrier_init(smem + 203032, 128); + // o_full: 1 barriers, init_count=1 + mbarrier_init(smem + 203040, 1); + // o_empty: 1 barriers, init_count=128 + mbarrier_init(smem + 203048, 128); + // --- pipeline 'work_pipe' --- + // work_full: 2 barriers, init_count=1 + mbarrier_init(smem + 203056, 1); + mbarrier_init(smem + 203064, 1); + // work_empty: 2 barriers, init_count=512 + mbarrier_init(smem + 203072, 512); + mbarrier_init(smem + 203080, 512); + // --- pipeline 'throttle_pipe' --- + // throttle_full: 2 barriers, init_count=128 + mbarrier_init(smem + 203088, 128); + mbarrier_init(smem + 203096, 128); + // throttle_empty: 2 barriers, init_count=32 + mbarrier_init(smem + 203104, 32); + mbarrier_init(smem + 203112, 32); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncwarp(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 203120); + if (warp == 0) { + int _tmem_hold = smem + 203120; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_tmem = taddr; + + // ---- Ordered hardware-WG register redistribution ---- + // Inc phase consumes the registers released above. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 136;"); + } + + // ---- Role: softmax ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 152;"); + { // softmax_main + const int warp_rank = warp; + const int logical_row_origin = warp_rank * 16; + const int my_row = logical_row_origin + lane % 16; + const int col_half = lane / 16; + int score_cursor = 0; + int stats_cursor = 0; + unsigned int softmax_work_stage = 0; + unsigned int softmax_q = blockIdx.x; + unsigned int softmax_head = blockIdx.y; + unsigned int softmax_batch = blockIdx.z; + unsigned int _phase_work_full = 0; + #pragma unroll 1 + for (unsigned int _work = 0; _work < total_work_items; _work++) { + int q_begin = cum_seq_lens_q[(int)softmax_batch]; + int q_end = cum_seq_lens_q[(int)(softmax_batch + 1)]; + int q_local = (int)softmax_q; + int work_valid = q_local < q_end - q_begin; + asm volatile("griddepcontrol.wait;" ::: "memory"); + float softmax_rescale_log2 = bmm1_scale[0] * 1.4426950408889634f; + asm volatile("griddepcontrol.wait;" ::: "memory"); + float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; + if (work_valid != 0) { + int query_idx = q_begin + q_local; + int _max_4 = ((sparse_topk_lens[query_idx] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx] + sparse_topk_lens_offset) : (0)); + int _min_3 = ((_max_4) < (sparse_topk) ? (_max_4) : (sparse_topk)); + int active_topk = _min_3; + int num_tiles = (active_topk + 128 - 1) / 128; + float row_max = -CAKE_INF; + float row_sum = 0.0f; + #pragma unroll 1 + for (int tile = 0; tile < num_tiles; tile++) { + int pipeline_tile = score_cursor + tile; + int score_stage = pipeline_tile & 1; + int score_phase = pipeline_tile >> 1 & 1; + mbarrier_wait(score_full_addr + (score_stage) * 8, score_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int score_col = ((score_stage != 0) ? 128 : 0); + float _tmem_load_0[64]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_0[0])), "=r"(*reinterpret_cast(&_tmem_load_0[1])), "=r"(*reinterpret_cast(&_tmem_load_0[2])), "=r"(*reinterpret_cast(&_tmem_load_0[3])), "=r"(*reinterpret_cast(&_tmem_load_0[4])), "=r"(*reinterpret_cast(&_tmem_load_0[5])), "=r"(*reinterpret_cast(&_tmem_load_0[6])), "=r"(*reinterpret_cast(&_tmem_load_0[7])), "=r"(*reinterpret_cast(&_tmem_load_0[8])), "=r"(*reinterpret_cast(&_tmem_load_0[9])), "=r"(*reinterpret_cast(&_tmem_load_0[10])), "=r"(*reinterpret_cast(&_tmem_load_0[11])), "=r"(*reinterpret_cast(&_tmem_load_0[12])), "=r"(*reinterpret_cast(&_tmem_load_0[13])), "=r"(*reinterpret_cast(&_tmem_load_0[14])), "=r"(*reinterpret_cast(&_tmem_load_0[15])), "=r"(*reinterpret_cast(&_tmem_load_0[16])), "=r"(*reinterpret_cast(&_tmem_load_0[17])), "=r"(*reinterpret_cast(&_tmem_load_0[18])), "=r"(*reinterpret_cast(&_tmem_load_0[19])), "=r"(*reinterpret_cast(&_tmem_load_0[20])), "=r"(*reinterpret_cast(&_tmem_load_0[21])), "=r"(*reinterpret_cast(&_tmem_load_0[22])), "=r"(*reinterpret_cast(&_tmem_load_0[23])), "=r"(*reinterpret_cast(&_tmem_load_0[24])), "=r"(*reinterpret_cast(&_tmem_load_0[25])), "=r"(*reinterpret_cast(&_tmem_load_0[26])), "=r"(*reinterpret_cast(&_tmem_load_0[27])), "=r"(*reinterpret_cast(&_tmem_load_0[28])), "=r"(*reinterpret_cast(&_tmem_load_0[29])), "=r"(*reinterpret_cast(&_tmem_load_0[30])), "=r"(*reinterpret_cast(&_tmem_load_0[31])) + : "r"(taddr + (unsigned int)score_col)); + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" + : "=r"(*reinterpret_cast(&_tmem_load_0[32])), "=r"(*reinterpret_cast(&_tmem_load_0[33])), "=r"(*reinterpret_cast(&_tmem_load_0[34])), "=r"(*reinterpret_cast(&_tmem_load_0[35])), "=r"(*reinterpret_cast(&_tmem_load_0[36])), "=r"(*reinterpret_cast(&_tmem_load_0[37])), "=r"(*reinterpret_cast(&_tmem_load_0[38])), "=r"(*reinterpret_cast(&_tmem_load_0[39])), "=r"(*reinterpret_cast(&_tmem_load_0[40])), "=r"(*reinterpret_cast(&_tmem_load_0[41])), "=r"(*reinterpret_cast(&_tmem_load_0[42])), "=r"(*reinterpret_cast(&_tmem_load_0[43])), "=r"(*reinterpret_cast(&_tmem_load_0[44])), "=r"(*reinterpret_cast(&_tmem_load_0[45])), "=r"(*reinterpret_cast(&_tmem_load_0[46])), "=r"(*reinterpret_cast(&_tmem_load_0[47])), "=r"(*reinterpret_cast(&_tmem_load_0[48])), "=r"(*reinterpret_cast(&_tmem_load_0[49])), "=r"(*reinterpret_cast(&_tmem_load_0[50])), "=r"(*reinterpret_cast(&_tmem_load_0[51])), "=r"(*reinterpret_cast(&_tmem_load_0[52])), "=r"(*reinterpret_cast(&_tmem_load_0[53])), "=r"(*reinterpret_cast(&_tmem_load_0[54])), "=r"(*reinterpret_cast(&_tmem_load_0[55])), "=r"(*reinterpret_cast(&_tmem_load_0[56])), "=r"(*reinterpret_cast(&_tmem_load_0[57])), "=r"(*reinterpret_cast(&_tmem_load_0[58])), "=r"(*reinterpret_cast(&_tmem_load_0[59])), "=r"(*reinterpret_cast(&_tmem_load_0[60])), "=r"(*reinterpret_cast(&_tmem_load_0[61])), "=r"(*reinterpret_cast(&_tmem_load_0[62])), "=r"(*reinterpret_cast(&_tmem_load_0[63])) + : "r"(taddr + (unsigned int)score_col + 32)); + int valid_cols = active_topk - tile * 128 - col_half * 64; + if (valid_cols < 0) { + valid_cols = 0; + } + if (valid_cols > 64) { + valid_cols = 64; + } + uint32_t _slice_lo_mask_0; + { + int _lim_0 = valid_cols; + if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } + else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); + } + } + if (!(_slice_lo_mask_0 & (1u << 0))) _tmem_load_0[0] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 1))) _tmem_load_0[1] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 2))) _tmem_load_0[2] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 3))) _tmem_load_0[3] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 4))) _tmem_load_0[4] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 5))) _tmem_load_0[5] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 6))) _tmem_load_0[6] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 7))) _tmem_load_0[7] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 8))) _tmem_load_0[8] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 9))) _tmem_load_0[9] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 10))) _tmem_load_0[10] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 11))) _tmem_load_0[11] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 12))) _tmem_load_0[12] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 13))) _tmem_load_0[13] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 14))) _tmem_load_0[14] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 15))) _tmem_load_0[15] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 16))) _tmem_load_0[16] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 17))) _tmem_load_0[17] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 18))) _tmem_load_0[18] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 19))) _tmem_load_0[19] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 20))) _tmem_load_0[20] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 21))) _tmem_load_0[21] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 22))) _tmem_load_0[22] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 23))) _tmem_load_0[23] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 24))) _tmem_load_0[24] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 25))) _tmem_load_0[25] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 26))) _tmem_load_0[26] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 27))) _tmem_load_0[27] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 28))) _tmem_load_0[28] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 29))) _tmem_load_0[29] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 30))) _tmem_load_0[30] = -CAKE_INF; + if (!(_slice_lo_mask_0 & (1u << 31))) _tmem_load_0[31] = -CAKE_INF; + uint32_t _slice_lo_mask_1; + { + int _lim_1 = valid_cols - 32; + if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } + else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); + } + } + if (!(_slice_lo_mask_1 & (1u << 0))) _tmem_load_0[32] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 1))) _tmem_load_0[33] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 2))) _tmem_load_0[34] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 3))) _tmem_load_0[35] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 4))) _tmem_load_0[36] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 5))) _tmem_load_0[37] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 6))) _tmem_load_0[38] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 7))) _tmem_load_0[39] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 8))) _tmem_load_0[40] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 9))) _tmem_load_0[41] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 10))) _tmem_load_0[42] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 11))) _tmem_load_0[43] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 12))) _tmem_load_0[44] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 13))) _tmem_load_0[45] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 14))) _tmem_load_0[46] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 15))) _tmem_load_0[47] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 16))) _tmem_load_0[48] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 17))) _tmem_load_0[49] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 18))) _tmem_load_0[50] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 19))) _tmem_load_0[51] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 20))) _tmem_load_0[52] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 21))) _tmem_load_0[53] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 22))) _tmem_load_0[54] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 23))) _tmem_load_0[55] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 24))) _tmem_load_0[56] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 25))) _tmem_load_0[57] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 26))) _tmem_load_0[58] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 27))) _tmem_load_0[59] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 28))) _tmem_load_0[60] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 29))) _tmem_load_0[61] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 30))) _tmem_load_0[62] = -CAKE_INF; + if (!(_slice_lo_mask_1 & (1u << 31))) _tmem_load_0[63] = -CAKE_INF; + float2 _reg_reduce_max2_2 = {-CAKE_INF, -CAKE_INF}; + row_max_x32_accum(&_tmem_load_0[0], _reg_reduce_max2_2); + row_max_x32_accum(&_tmem_load_0[32], _reg_reduce_max2_2); + float _tmem_load_0_max = row_max_reduce(_reg_reduce_max2_2); + float _max_5 = max_noftz(row_max, _tmem_load_0_max); + float new_max = _max_5; + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, new_max, 16); + float _max_6 = max_noftz(new_max, _shfl_xor_0); + new_max = _max_6; + float _fma_0 = __fmaf_rn(row_max, softmax_rescale_log2, (-new_max) * softmax_rescale_log2); + float delta = _fma_0; + float _exp2_0 = approx_exp2(delta); + float acc_scale = ((row_max > -CAKE_INF) ? _exp2_0 : 1.0f); + int stats_tile = stats_cursor + tile; + int stats_stage = stats_tile & 1; + int stats_phase = stats_tile >> 1 & 1; + mbarrier_wait(stats_empty_addr + (stats_stage) * 8, 1 - stats_phase); + float stats_pair[2]; + stats_pair[0] = row_max; + stats_pair[1] = new_max; + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(taddr + (unsigned int)(stats_stage * 128)), "f"(stats_pair[0]), "f"(stats_pair[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(stats_full_addr + (stats_stage) * 8); + float safe_max = ((new_max == -CAKE_INF) ? 0.0f : new_max); + float max_scaled = safe_max * softmax_scale_log2; + const float2 _fma_b2_3 = {softmax_scale_log2, softmax_scale_log2}; + const float2 _fma_c2_4 = {-max_scaled + 8.8073549f, -max_scaled + 8.8073549f}; + #pragma unroll + for (int _lf = 0; _lf < 32; _lf++) + fma_f32x2_inplace(&reinterpret_cast(_tmem_load_0)[_lf], _fma_b2_3, _fma_c2_4); + #pragma unroll + for (int _le = 0; _le < 64; _le++) { + _tmem_load_0[_le] = approx_exp2(_tmem_load_0[_le]); + } + float2 _reg_reduce_sum2_5 = make_float2(0.0f, 0.0f); + softmax_block_sum(&_tmem_load_0[0], &_reg_reduce_sum2_5); + softmax_block_sum(&_tmem_load_0[32], &_reg_reduce_sum2_5); + float _tmem_load_0_sum = _reg_reduce_sum2_5.x + _reg_reduce_sum2_5.y; + float block_sum = _tmem_load_0_sum; + float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, block_sum, 16); + block_sum = block_sum + _shfl_xor_1; + float _fma_1 = __fmaf_rn(row_sum, acc_scale, block_sum); + row_sum = _fma_1; + uint32_t _fp8_0[16]; + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_0[0]), "f"(_tmem_load_0[1]), + "f"(_tmem_load_0[2]), "f"(_tmem_load_0[3])); + _fp8_0[0] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_0[4]), "f"(_tmem_load_0[5]), + "f"(_tmem_load_0[6]), "f"(_tmem_load_0[7])); + _fp8_0[1] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_0[8]), "f"(_tmem_load_0[9]), + "f"(_tmem_load_0[10]), "f"(_tmem_load_0[11])); + _fp8_0[2] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_0[12]), "f"(_tmem_load_0[13]), + "f"(_tmem_load_0[14]), "f"(_tmem_load_0[15])); + _fp8_0[3] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_0[16]), "f"(_tmem_load_0[17]), + "f"(_tmem_load_0[18]), "f"(_tmem_load_0[19])); + _fp8_0[4] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_0[20]), "f"(_tmem_load_0[21]), + "f"(_tmem_load_0[22]), "f"(_tmem_load_0[23])); + _fp8_0[5] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_0[24]), "f"(_tmem_load_0[25]), + "f"(_tmem_load_0[26]), "f"(_tmem_load_0[27])); + _fp8_0[6] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_0[28]), "f"(_tmem_load_0[29]), + "f"(_tmem_load_0[30]), "f"(_tmem_load_0[31])); + _fp8_0[7] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_0[32]), "f"(_tmem_load_0[33]), + "f"(_tmem_load_0[34]), "f"(_tmem_load_0[35])); + _fp8_0[8] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_0[36]), "f"(_tmem_load_0[37]), + "f"(_tmem_load_0[38]), "f"(_tmem_load_0[39])); + _fp8_0[9] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_0[40]), "f"(_tmem_load_0[41]), + "f"(_tmem_load_0[42]), "f"(_tmem_load_0[43])); + _fp8_0[10] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_0[44]), "f"(_tmem_load_0[45]), + "f"(_tmem_load_0[46]), "f"(_tmem_load_0[47])); + _fp8_0[11] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_0[48]), "f"(_tmem_load_0[49]), + "f"(_tmem_load_0[50]), "f"(_tmem_load_0[51])); + _fp8_0[12] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_0[52]), "f"(_tmem_load_0[53]), + "f"(_tmem_load_0[54]), "f"(_tmem_load_0[55])); + _fp8_0[13] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_0[56]), "f"(_tmem_load_0[57]), + "f"(_tmem_load_0[58]), "f"(_tmem_load_0[59])); + _fp8_0[14] = _packed; + } + { + uint32_t _packed; + asm volatile("{\n\t" + ".reg .b16 _lo;\n\t" + ".reg .b16 _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}" + : "=r"(_packed) : "f"(_tmem_load_0[60]), "f"(_tmem_load_0[61]), + "f"(_tmem_load_0[62]), "f"(_tmem_load_0[63])); + _fp8_0[15] = _packed; + } + int p_col = ((score_stage != 0) ? 160 : 32); + asm volatile( + "tcgen05.st.sync.aligned.16x32bx2.x16.b32" + " [%0], 16, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(taddr + (unsigned int)p_col), "r"(*reinterpret_cast(&_fp8_0[0])), "r"(*reinterpret_cast(&_fp8_0[1])), "r"(*reinterpret_cast(&_fp8_0[2])), "r"(*reinterpret_cast(&_fp8_0[3])), "r"(*reinterpret_cast(&_fp8_0[4])), "r"(*reinterpret_cast(&_fp8_0[5])), "r"(*reinterpret_cast(&_fp8_0[6])), "r"(*reinterpret_cast(&_fp8_0[7])), "r"(*reinterpret_cast(&_fp8_0[8])), "r"(*reinterpret_cast(&_fp8_0[9])), "r"(*reinterpret_cast(&_fp8_0[10])), "r"(*reinterpret_cast(&_fp8_0[11])), "r"(*reinterpret_cast(&_fp8_0[12])), "r"(*reinterpret_cast(&_fp8_0[13])), "r"(*reinterpret_cast(&_fp8_0[14])), "r"(*reinterpret_cast(&_fp8_0[15]))); + asm volatile( + "tcgen05.st.sync.aligned.16x32bx2.x16.b32" + " [%0], 16, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(taddr + 1048576 + (unsigned int)p_col), "r"(*reinterpret_cast(&_fp8_0[0])), "r"(*reinterpret_cast(&_fp8_0[1])), "r"(*reinterpret_cast(&_fp8_0[2])), "r"(*reinterpret_cast(&_fp8_0[3])), "r"(*reinterpret_cast(&_fp8_0[4])), "r"(*reinterpret_cast(&_fp8_0[5])), "r"(*reinterpret_cast(&_fp8_0[6])), "r"(*reinterpret_cast(&_fp8_0[7])), "r"(*reinterpret_cast(&_fp8_0[8])), "r"(*reinterpret_cast(&_fp8_0[9])), "r"(*reinterpret_cast(&_fp8_0[10])), "r"(*reinterpret_cast(&_fp8_0[11])), "r"(*reinterpret_cast(&_fp8_0[12])), "r"(*reinterpret_cast(&_fp8_0[13])), "r"(*reinterpret_cast(&_fp8_0[14])), "r"(*reinterpret_cast(&_fp8_0[15]))); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + if (num_tiles > tile + 1) { + mbarrier_arrive(score_empty_addr + (score_stage) * 8); + } + row_max = new_max; + } + int final_stats_tile = stats_cursor + num_tiles; + int final_stats_stage = final_stats_tile & 1; + int final_stats_phase = final_stats_tile >> 1 & 1; + mbarrier_wait(stats_empty_addr + (final_stats_stage) * 8, 1 - final_stats_phase); + float final_stats[2]; + final_stats[0] = row_sum; + final_stats[1] = row_max; + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x2.b32" + " [%0], {%1, %2};" + :: "r"(taddr + (unsigned int)(final_stats_stage * 128)), "f"(final_stats[0]), "f"(final_stats[1])); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(stats_full_addr + (final_stats_stage) * 8); + int tail_stats_tile = final_stats_tile + 1; + int tail_stats_stage = tail_stats_tile & 1; + int tail_stats_phase = tail_stats_tile >> 1 & 1; + mbarrier_wait(stats_empty_addr + (tail_stats_stage) * 8, 1 - tail_stats_phase); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(stats_full_addr + (tail_stats_stage) * 8); + int last_score_tile = score_cursor + num_tiles - 1; + mbarrier_arrive(score_empty_addr + (last_score_tile & 1) * 8); + stats_cursor = stats_cursor + num_tiles + 2; + score_cursor = score_cursor + num_tiles; + } + mbarrier_wait(work_full_addr + (softmax_work_stage) * 8, _phase_work_full); + uint32_t _clc_valid_3 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_3) + : "r"(work_response_addr + softmax_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_9 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_9) + : "r"(work_response_addr + softmax_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_10 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_10) + : "r"(work_response_addr + softmax_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_11 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_11) + : "r"(work_response_addr + softmax_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(work_empty_addr + (softmax_work_stage) * 8); + softmax_work_stage += 1; + if (softmax_work_stage == 2) { softmax_work_stage = 0; _phase_work_full ^= 1; } + if (_clc_valid_3 == 0) { + break; + } + softmax_q = _clc_ctaid_9; + softmax_head = _clc_ctaid_10; + softmax_batch = _clc_ctaid_11; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + } + } + // ---- Role: correction ---- + if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 144;"); + { // correction_main + const int warp_rank_1 = warp - 4; + const int logical_row_origin_1 = warp_rank_1 * 16; + const int my_row_1 = logical_row_origin_1 + lane % 16; + const int col_half_1 = lane / 16; + int score_cursor_1 = 0; + int stats_cursor_1 = 0; + unsigned int correction_work_stage = 0; + unsigned int correction_q = blockIdx.x; + unsigned int correction_head = blockIdx.y; + unsigned int correction_batch = blockIdx.z; + unsigned int _phase_work_full_1 = 0; + #pragma unroll 1 + for (unsigned int _work_1 = 0; _work_1 < total_work_items; _work_1++) { + int q_begin_1 = cum_seq_lens_q[(int)correction_batch]; + int q_end_1 = cum_seq_lens_q[(int)(correction_batch + 1)]; + int q_local_1 = (int)correction_q; + int work_valid_1 = q_local_1 < q_end_1 - q_begin_1; + asm volatile("griddepcontrol.wait;" ::: "memory"); + float softmax_scale_log2_1 = bmm1_scale[0] * 1.4426950408889634f; + float output_scale = bmm2_scale[0]; + if (work_valid_1 != 0) { + int query_idx_1 = q_begin_1 + q_local_1; + int head_base = (int)correction_head * 64; + int head_idx = head_base + my_row_1; + int _max_7 = ((sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) : (0)); + int _min_4 = ((_max_7) < (sparse_topk) ? (_max_7) : (sparse_topk)); + int active_topk_1 = _min_4; + int num_tiles_1 = (active_topk_1 + 128 - 1) / 128; + int first_stats_tile = stats_cursor_1; + int first_stats_stage = first_stats_tile & 1; + int first_stats_phase = first_stats_tile >> 1 & 1; + mbarrier_wait(stats_full_addr + (first_stats_stage) * 8, first_stats_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + mbarrier_arrive(stats_empty_addr + (first_stats_stage) * 8); + #pragma unroll 1 + for (int tile_1 = 1; tile_1 < num_tiles_1; tile_1++) { + int stats_tile_1 = stats_cursor_1 + tile_1; + int stats_stage_1 = stats_tile_1 & 1; + int stats_phase_1 = stats_tile_1 >> 1 & 1; + mbarrier_wait(stats_full_addr + (stats_stage_1) * 8, stats_phase_1); + asm volatile("tcgen05.fence::after_thread_sync;"); + float _tmem_load_1[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]) + : "r"(taddr + (unsigned int)(stats_stage_1 * 128))); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + mbarrier_arrive(stats_empty_addr + (stats_stage_1) * 8); + int output_tile = score_cursor_1 + tile_1 - 1; + mbarrier_wait(o_full_addr, output_tile & 1); + asm volatile("tcgen05.fence::after_thread_sync;"); + float max_delta = _tmem_load_1[0] - _tmem_load_1[1]; + float _exp2_1 = approx_exp2(max_delta * softmax_scale_log2_1); + float acc_scale_1 = ((max_delta != 0.0f) ? _exp2_1 : 1.0f); + int _vote_0 = __any_sync(0xFFFFFFFF, acc_scale_1 < 1.0f); + int any_rescale = _vote_0; + if (any_rescale != 0) { + #pragma unroll + for (int output_half = 0; output_half < 2; output_half++) { + int row_bank = ((output_half != 0) ? 1048576 : 0); + #pragma unroll + for (int col = 0; col < 128; col += 64) { + float _tmem_load_2[64]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 128;" + : "=r"(*reinterpret_cast(&_tmem_load_2[0])), "=r"(*reinterpret_cast(&_tmem_load_2[1])), "=r"(*reinterpret_cast(&_tmem_load_2[2])), "=r"(*reinterpret_cast(&_tmem_load_2[3])), "=r"(*reinterpret_cast(&_tmem_load_2[4])), "=r"(*reinterpret_cast(&_tmem_load_2[5])), "=r"(*reinterpret_cast(&_tmem_load_2[6])), "=r"(*reinterpret_cast(&_tmem_load_2[7])), "=r"(*reinterpret_cast(&_tmem_load_2[8])), "=r"(*reinterpret_cast(&_tmem_load_2[9])), "=r"(*reinterpret_cast(&_tmem_load_2[10])), "=r"(*reinterpret_cast(&_tmem_load_2[11])), "=r"(*reinterpret_cast(&_tmem_load_2[12])), "=r"(*reinterpret_cast(&_tmem_load_2[13])), "=r"(*reinterpret_cast(&_tmem_load_2[14])), "=r"(*reinterpret_cast(&_tmem_load_2[15])), "=r"(*reinterpret_cast(&_tmem_load_2[16])), "=r"(*reinterpret_cast(&_tmem_load_2[17])), "=r"(*reinterpret_cast(&_tmem_load_2[18])), "=r"(*reinterpret_cast(&_tmem_load_2[19])), "=r"(*reinterpret_cast(&_tmem_load_2[20])), "=r"(*reinterpret_cast(&_tmem_load_2[21])), "=r"(*reinterpret_cast(&_tmem_load_2[22])), "=r"(*reinterpret_cast(&_tmem_load_2[23])), "=r"(*reinterpret_cast(&_tmem_load_2[24])), "=r"(*reinterpret_cast(&_tmem_load_2[25])), "=r"(*reinterpret_cast(&_tmem_load_2[26])), "=r"(*reinterpret_cast(&_tmem_load_2[27])), "=r"(*reinterpret_cast(&_tmem_load_2[28])), "=r"(*reinterpret_cast(&_tmem_load_2[29])), "=r"(*reinterpret_cast(&_tmem_load_2[30])), "=r"(*reinterpret_cast(&_tmem_load_2[31])) + : "r"(taddr + 256 + (unsigned int)row_bank + (unsigned int)col)); + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 128;" + : "=r"(*reinterpret_cast(&_tmem_load_2[32])), "=r"(*reinterpret_cast(&_tmem_load_2[33])), "=r"(*reinterpret_cast(&_tmem_load_2[34])), "=r"(*reinterpret_cast(&_tmem_load_2[35])), "=r"(*reinterpret_cast(&_tmem_load_2[36])), "=r"(*reinterpret_cast(&_tmem_load_2[37])), "=r"(*reinterpret_cast(&_tmem_load_2[38])), "=r"(*reinterpret_cast(&_tmem_load_2[39])), "=r"(*reinterpret_cast(&_tmem_load_2[40])), "=r"(*reinterpret_cast(&_tmem_load_2[41])), "=r"(*reinterpret_cast(&_tmem_load_2[42])), "=r"(*reinterpret_cast(&_tmem_load_2[43])), "=r"(*reinterpret_cast(&_tmem_load_2[44])), "=r"(*reinterpret_cast(&_tmem_load_2[45])), "=r"(*reinterpret_cast(&_tmem_load_2[46])), "=r"(*reinterpret_cast(&_tmem_load_2[47])), "=r"(*reinterpret_cast(&_tmem_load_2[48])), "=r"(*reinterpret_cast(&_tmem_load_2[49])), "=r"(*reinterpret_cast(&_tmem_load_2[50])), "=r"(*reinterpret_cast(&_tmem_load_2[51])), "=r"(*reinterpret_cast(&_tmem_load_2[52])), "=r"(*reinterpret_cast(&_tmem_load_2[53])), "=r"(*reinterpret_cast(&_tmem_load_2[54])), "=r"(*reinterpret_cast(&_tmem_load_2[55])), "=r"(*reinterpret_cast(&_tmem_load_2[56])), "=r"(*reinterpret_cast(&_tmem_load_2[57])), "=r"(*reinterpret_cast(&_tmem_load_2[58])), "=r"(*reinterpret_cast(&_tmem_load_2[59])), "=r"(*reinterpret_cast(&_tmem_load_2[60])), "=r"(*reinterpret_cast(&_tmem_load_2[61])), "=r"(*reinterpret_cast(&_tmem_load_2[62])), "=r"(*reinterpret_cast(&_tmem_load_2[63])) + : "r"(taddr + 256 + (unsigned int)row_bank + (unsigned int)col + 32)); + const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 32; _ls++) + mul_f32x2_inplace(&reinterpret_cast(_tmem_load_2)[_ls], _scale2_0); + asm volatile( + "tcgen05.st.sync.aligned.16x32bx2.x64.b32" + " [%0], 128, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64};" + :: "r"(taddr + 256 + (unsigned int)row_bank + (unsigned int)col), "r"(*reinterpret_cast(&_tmem_load_2[0])), "r"(*reinterpret_cast(&_tmem_load_2[1])), "r"(*reinterpret_cast(&_tmem_load_2[2])), "r"(*reinterpret_cast(&_tmem_load_2[3])), "r"(*reinterpret_cast(&_tmem_load_2[4])), "r"(*reinterpret_cast(&_tmem_load_2[5])), "r"(*reinterpret_cast(&_tmem_load_2[6])), "r"(*reinterpret_cast(&_tmem_load_2[7])), "r"(*reinterpret_cast(&_tmem_load_2[8])), "r"(*reinterpret_cast(&_tmem_load_2[9])), "r"(*reinterpret_cast(&_tmem_load_2[10])), "r"(*reinterpret_cast(&_tmem_load_2[11])), "r"(*reinterpret_cast(&_tmem_load_2[12])), "r"(*reinterpret_cast(&_tmem_load_2[13])), "r"(*reinterpret_cast(&_tmem_load_2[14])), "r"(*reinterpret_cast(&_tmem_load_2[15])), "r"(*reinterpret_cast(&_tmem_load_2[16])), "r"(*reinterpret_cast(&_tmem_load_2[17])), "r"(*reinterpret_cast(&_tmem_load_2[18])), "r"(*reinterpret_cast(&_tmem_load_2[19])), "r"(*reinterpret_cast(&_tmem_load_2[20])), "r"(*reinterpret_cast(&_tmem_load_2[21])), "r"(*reinterpret_cast(&_tmem_load_2[22])), "r"(*reinterpret_cast(&_tmem_load_2[23])), "r"(*reinterpret_cast(&_tmem_load_2[24])), "r"(*reinterpret_cast(&_tmem_load_2[25])), "r"(*reinterpret_cast(&_tmem_load_2[26])), "r"(*reinterpret_cast(&_tmem_load_2[27])), "r"(*reinterpret_cast(&_tmem_load_2[28])), "r"(*reinterpret_cast(&_tmem_load_2[29])), "r"(*reinterpret_cast(&_tmem_load_2[30])), "r"(*reinterpret_cast(&_tmem_load_2[31])), "r"(*reinterpret_cast(&_tmem_load_2[32])), "r"(*reinterpret_cast(&_tmem_load_2[33])), "r"(*reinterpret_cast(&_tmem_load_2[34])), "r"(*reinterpret_cast(&_tmem_load_2[35])), "r"(*reinterpret_cast(&_tmem_load_2[36])), "r"(*reinterpret_cast(&_tmem_load_2[37])), "r"(*reinterpret_cast(&_tmem_load_2[38])), "r"(*reinterpret_cast(&_tmem_load_2[39])), "r"(*reinterpret_cast(&_tmem_load_2[40])), "r"(*reinterpret_cast(&_tmem_load_2[41])), "r"(*reinterpret_cast(&_tmem_load_2[42])), "r"(*reinterpret_cast(&_tmem_load_2[43])), "r"(*reinterpret_cast(&_tmem_load_2[44])), "r"(*reinterpret_cast(&_tmem_load_2[45])), "r"(*reinterpret_cast(&_tmem_load_2[46])), "r"(*reinterpret_cast(&_tmem_load_2[47])), "r"(*reinterpret_cast(&_tmem_load_2[48])), "r"(*reinterpret_cast(&_tmem_load_2[49])), "r"(*reinterpret_cast(&_tmem_load_2[50])), "r"(*reinterpret_cast(&_tmem_load_2[51])), "r"(*reinterpret_cast(&_tmem_load_2[52])), "r"(*reinterpret_cast(&_tmem_load_2[53])), "r"(*reinterpret_cast(&_tmem_load_2[54])), "r"(*reinterpret_cast(&_tmem_load_2[55])), "r"(*reinterpret_cast(&_tmem_load_2[56])), "r"(*reinterpret_cast(&_tmem_load_2[57])), "r"(*reinterpret_cast(&_tmem_load_2[58])), "r"(*reinterpret_cast(&_tmem_load_2[59])), "r"(*reinterpret_cast(&_tmem_load_2[60])), "r"(*reinterpret_cast(&_tmem_load_2[61])), "r"(*reinterpret_cast(&_tmem_load_2[62])), "r"(*reinterpret_cast(&_tmem_load_2[63]))); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + } + } + } + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(o_empty_addr); + } + int final_stats_tile_1 = stats_cursor_1 + num_tiles_1; + int final_stats_stage_1 = final_stats_tile_1 & 1; + int final_stats_phase_1 = final_stats_tile_1 >> 1 & 1; + mbarrier_wait(stats_full_addr + (final_stats_stage_1) * 8, final_stats_phase_1); + asm volatile("tcgen05.fence::after_thread_sync;"); + float _tmem_load_3[2]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x2.b32" + " {%0, %1}, [%2];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]) + : "r"(taddr + (unsigned int)(final_stats_stage_1 * 128))); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + mbarrier_arrive(stats_empty_addr + (final_stats_stage_1) * 8); + int tail_stats_tile_1 = final_stats_tile_1 + 1; + int tail_stats_stage_1 = tail_stats_tile_1 & 1; + int tail_stats_phase_1 = tail_stats_tile_1 >> 1 & 1; + mbarrier_wait(stats_full_addr + (tail_stats_stage_1) * 8, tail_stats_phase_1); + asm volatile("tcgen05.fence::after_thread_sync;"); + mbarrier_arrive(stats_empty_addr + (tail_stats_stage_1) * 8); + int last_output_tile = score_cursor_1 + num_tiles_1 - 1; + mbarrier_wait(o_full_addr, last_output_tile & 1); + asm volatile("tcgen05.fence::after_thread_sync;"); + float total_sum = _tmem_load_3[0]; + float final_max = _tmem_load_3[1]; + int output_in_bounds = head_idx < num_heads; + if (((int)(has_sinks != 0) & output_in_bounds) != 0) { + float _exp2_2 = approx_exp2(sinks[head_idx] * 1.4426950408889634f - final_max * softmax_scale_log2_1); + total_sum = total_sum + _exp2_2 * 448.0f; + } + float _rcp_0 = approx_rcp(total_sum); + float inv_sum = ((total_sum > 0.0f) ? output_scale * _rcp_0 : 0.0f); + int output_base = (query_idx_1 * num_heads + head_idx) * 512; + #pragma unroll + for (int output_half_1 = 0; output_half_1 < 2; output_half_1++) { + int row_bank_1 = ((output_half_1 != 0) ? 1048576 : 0); + int col_base = col_half_1 * 128; + #pragma unroll 2 + for (int wcol = 0; wcol < 128; wcol += 16) { + float _tmem_load_4[16]; + asm volatile( + "tcgen05.ld.sync.aligned.16x32bx2.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16], 128;" + : "=r"(*reinterpret_cast(&_tmem_load_4[0])), "=r"(*reinterpret_cast(&_tmem_load_4[1])), "=r"(*reinterpret_cast(&_tmem_load_4[2])), "=r"(*reinterpret_cast(&_tmem_load_4[3])), "=r"(*reinterpret_cast(&_tmem_load_4[4])), "=r"(*reinterpret_cast(&_tmem_load_4[5])), "=r"(*reinterpret_cast(&_tmem_load_4[6])), "=r"(*reinterpret_cast(&_tmem_load_4[7])), "=r"(*reinterpret_cast(&_tmem_load_4[8])), "=r"(*reinterpret_cast(&_tmem_load_4[9])), "=r"(*reinterpret_cast(&_tmem_load_4[10])), "=r"(*reinterpret_cast(&_tmem_load_4[11])), "=r"(*reinterpret_cast(&_tmem_load_4[12])), "=r"(*reinterpret_cast(&_tmem_load_4[13])), "=r"(*reinterpret_cast(&_tmem_load_4[14])), "=r"(*reinterpret_cast(&_tmem_load_4[15])) + : "r"(taddr + 256 + (unsigned int)row_bank_1 + (unsigned int)wcol)); + const float2 _scale2_1 = {inv_sum, inv_sum}; + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + mul_f32x2_inplace(&reinterpret_cast(_tmem_load_4)[_ls], _scale2_1); + uint32_t _tmem_load_4_bf16[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_4[_lp*2 + 0], _tmem_load_4[_lp*2+1 + 0])); + _tmem_load_4_bf16[_lp] = *(uint32_t*)&_bf2; + } + if (output_in_bounds != 0) { + { + const unsigned* _raw_stv8_2 = reinterpret_cast(_tmem_load_4_bf16); + asm volatile( + "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"((void*)(O + (output_base + output_half_1 * 256 + col_base + wcol))), "r"(_raw_stv8_2[0]), "r"(_raw_stv8_2[1]), "r"(_raw_stv8_2[2]), "r"(_raw_stv8_2[3]), "r"(_raw_stv8_2[4]), "r"(_raw_stv8_2[5]), "r"(_raw_stv8_2[6]), "r"(_raw_stv8_2[7]) : "memory"); + } + } + } + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(o_empty_addr); + stats_cursor_1 = stats_cursor_1 + num_tiles_1 + 2; + score_cursor_1 = score_cursor_1 + num_tiles_1; + } + mbarrier_wait(work_full_addr + (correction_work_stage) * 8, _phase_work_full_1); + uint32_t _clc_valid_4 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_4) + : "r"(work_response_addr + correction_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_12 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_12) + : "r"(work_response_addr + correction_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_13 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_13) + : "r"(work_response_addr + correction_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_14 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_14) + : "r"(work_response_addr + correction_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(work_empty_addr + (correction_work_stage) * 8); + correction_work_stage += 1; + if (correction_work_stage == 2) { correction_work_stage = 0; _phase_work_full_1 ^= 1; } + if (_clc_valid_4 == 0) { + break; + } + correction_q = _clc_ctaid_12; + correction_head = _clc_ctaid_13; + correction_batch = _clc_ctaid_14; + } + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("barrier.sync 7, 128;" ::: "memory"); + if (warp == 4) { + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); + } + } + } + // ---- Role: mma ---- + if (warp == 8) { + { // mma_main + CakePipelineState<4> _kv_cons_state_state{0u, 0u, 0u}; + unsigned int q_cons_stage = 0; + int score_prod_stage = 0; + int score_prod_phase = 0; + int p_cons_stage = 0; + int p_cons_phase = 0; + int o_prod_phase = 0; + unsigned int mma_work_stage = 0; + unsigned int mma_q = blockIdx.x; + unsigned int mma_head = blockIdx.y; + unsigned int mma_batch = blockIdx.z; + #pragma unroll + for (int score_lookahead = 0; score_lookahead < 2; score_lookahead++) { + int lookahead_stage = score_lookahead & 1; + int lookahead_phase = score_lookahead >> 1 & 1; + mbarrier_wait(score_empty_addr + (lookahead_stage) * 8, 1 - lookahead_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + } + unsigned int _phase_q_full = 0; + unsigned int _phase_work_full_2 = 0; + #pragma unroll 1 + for (unsigned int _work_2 = 0; _work_2 < total_work_items; _work_2++) { + int q_begin_2 = cum_seq_lens_q[(int)mma_batch]; + int q_end_2 = cum_seq_lens_q[(int)(mma_batch + 1)]; + int q_local_2 = (int)mma_q; + int work_valid_2 = q_local_2 < q_end_2 - q_begin_2; + if (work_valid_2 != 0) { + int query_idx_2 = q_begin_2 + q_local_2; + int _max_3 = ((sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) : (0)); + int _min_2 = ((_max_3) < (sparse_topk) ? (_max_3) : (sparse_topk)); + int active_topk_2 = _min_2; + int num_tiles_2 = (active_topk_2 + 128 - 1) / 128; + if (num_tiles_2 > 0) { + mbarrier_wait(q_full_addr + (q_cons_stage) * 8, _phase_q_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + int score_col_1 = ((score_prod_stage != 0) ? 128 : 0); + mbarrier_wait(kv_full_addr + (_kv_cons_state_state.index()) * 8, _kv_cons_state_state.phase()); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_0 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (q_cons_stage) * 2048); + int _mma_b_lo_0 = make_warp_uniform((((smem_k_addr) >> 4) & 0x3FFF) + (_kv_cons_state_state.index()) * 2048); + { + uint64_t _mma_ss_a_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_0); + uint64_t _mma_ss_b_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_0); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69206032, 0); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 506U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 1018U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69206032, 1); + } + } + elect_commit(kv_empty_addr + (_kv_cons_state_state.index()) * 8); + ++_kv_cons_state_state; + mbarrier_wait(kv_full_addr + (_kv_cons_state_state.index()) * 8, _kv_cons_state_state.phase()); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_1 = make_warp_uniform((((smem_q_addr + 16384) >> 4) & 0x3FFF) + (q_cons_stage) * 2048); + int _mma_b_lo_1 = make_warp_uniform((((smem_k_addr) >> 4) & 0x3FFF) + (_kv_cons_state_state.index()) * 2048); + { + uint64_t _mma_ss_a_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_1); + uint64_t _mma_ss_b_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_1); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 506U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 1018U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69206032, 1); + } + } + elect_commit(score_full_addr + (score_prod_stage) * 8); + elect_commit(kv_empty_addr + (_kv_cons_state_state.index()) * 8); + ++_kv_cons_state_state; + score_prod_stage ^= 1; + if (score_prod_stage == 0) { + score_prod_phase ^= 1; + } + int first_pv = 1; + #pragma unroll 1 + for (int _steady = 0; _steady < num_tiles_2 - 1; _steady++) { + score_col_1 = ((score_prod_stage != 0) ? 128 : 0); + mbarrier_wait(kv_full_addr + (_kv_cons_state_state.index()) * 8, _kv_cons_state_state.phase()); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_2 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (q_cons_stage) * 2048); + int _mma_b_lo_2 = make_warp_uniform((((smem_k_addr) >> 4) & 0x3FFF) + (_kv_cons_state_state.index()) * 2048); + { + uint64_t _mma_ss_a_desc_2 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_2); + uint64_t _mma_ss_b_desc_2 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_2); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 69206032, 0); + } + incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_2, 506U); + incr_smem_desc_lo(_mma_ss_b_desc_2, 1018U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 69206032, 1); + } + } + elect_commit(kv_empty_addr + (_kv_cons_state_state.index()) * 8); + ++_kv_cons_state_state; + mbarrier_wait(kv_full_addr + (_kv_cons_state_state.index()) * 8, _kv_cons_state_state.phase()); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_lo_3 = make_warp_uniform((((smem_q_addr + 16384) >> 4) & 0x3FFF) + (q_cons_stage) * 2048); + int _mma_b_lo_3 = make_warp_uniform((((smem_k_addr) >> 4) & 0x3FFF) + (_kv_cons_state_state.index()) * 2048); + { + uint64_t _mma_ss_a_desc_3 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_3); + uint64_t _mma_ss_b_desc_3 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_3); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_3, 506U); + incr_smem_desc_lo(_mma_ss_b_desc_3, 1018U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 69206032, 1); + } + incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); + incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); + if (elect_sync()) { + tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 69206032, 1); + } + } + elect_commit(score_full_addr + (score_prod_stage) * 8); + elect_commit(kv_empty_addr + (_kv_cons_state_state.index()) * 8); + ++_kv_cons_state_state; + score_prod_stage ^= 1; + if (score_prod_stage == 0) { + score_prod_phase ^= 1; + } + mbarrier_wait(score_empty_addr + (score_prod_stage) * 8, 1 - score_prod_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + mbarrier_wait(o_empty_addr, 1 - o_prod_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int v_half = 0; v_half < 2; v_half++) { + mbarrier_wait(kv_full_addr + (_kv_cons_state_state.index()) * 8, _kv_cons_state_state.phase()); + asm volatile("tcgen05.fence::after_thread_sync;"); + int row_bank_2 = ((v_half != 0) ? 1048576 : 0); + int p_col_1 = ((p_cons_stage != 0) ? 160 : 32) + row_bank_2; + int _mma_b_lo_4 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (_kv_cons_state_state.index()) * 2048); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 dhi, blo, ta, id;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 71368720;\n\t" + "mov.b32 ta, %2;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p0;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" + "}\n" + :: "r"((tmem_tmem + (256 + row_bank_2))), "r"(_mma_b_lo_4), "r"(tmem_tmem + p_col_1), "r"(((first_pv) ? 0 : 1))); + elect_commit(kv_empty_addr + (_kv_cons_state_state.index()) * 8); + ++_kv_cons_state_state; + } + first_pv = 0; + elect_commit(o_full_addr); + p_cons_stage ^= 1; + if (p_cons_stage == 0) { + p_cons_phase ^= 1; + } + o_prod_phase ^= 1; + } + elect_commit(q_empty_addr + (q_cons_stage) * 8); + q_cons_stage += 1; + if (q_cons_stage == 2) { q_cons_stage = 0; _phase_q_full ^= 1; } + int tail_score_stage = score_prod_stage ^ 1; + int tail_score_phase = score_prod_phase ^ score_prod_stage; + mbarrier_wait(score_empty_addr + (tail_score_stage) * 8, 1 - tail_score_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + mbarrier_wait(o_empty_addr, 1 - o_prod_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll + for (int v_half_1 = 0; v_half_1 < 2; v_half_1++) { + mbarrier_wait(kv_full_addr + (_kv_cons_state_state.index()) * 8, _kv_cons_state_state.phase()); + asm volatile("tcgen05.fence::after_thread_sync;"); + int row_bank_3 = ((v_half_1 != 0) ? 1048576 : 0); + int p_col_2 = ((p_cons_stage != 0) ? 160 : 32) + row_bank_3; + int _mma_b_lo_5 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (_kv_cons_state_state.index()) * 2048); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 dhi, blo, ta, id;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 71368720;\n\t" + "mov.b32 ta, %2;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p0;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 256;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" + "}\n" + :: "r"((tmem_tmem + (256 + row_bank_3))), "r"(_mma_b_lo_5), "r"(tmem_tmem + p_col_2), "r"(((first_pv) ? 0 : 1))); + elect_commit(kv_empty_addr + (_kv_cons_state_state.index()) * 8); + ++_kv_cons_state_state; + } + elect_commit(o_full_addr); + p_cons_stage ^= 1; + if (p_cons_stage == 0) { + p_cons_phase ^= 1; + } + o_prod_phase ^= 1; + } + } + mbarrier_wait(work_full_addr + (mma_work_stage) * 8, _phase_work_full_2); + uint32_t _clc_valid_2 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_2) + : "r"(work_response_addr + mma_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_6 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_6) + : "r"(work_response_addr + mma_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_7 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_7) + : "r"(work_response_addr + mma_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_8 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_8) + : "r"(work_response_addr + mma_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(work_empty_addr + (mma_work_stage) * 8); + mma_work_stage += 1; + if (mma_work_stage == 2) { mma_work_stage = 0; _phase_work_full_2 ^= 1; } + if (_clc_valid_2 == 0) { + break; + } + mma_q = _clc_ctaid_6; + mma_head = _clc_ctaid_7; + mma_batch = _clc_ctaid_8; + } + } + } + // ---- Role: page_loader ---- + if (warp == 9) { + { // page_loader_main + unsigned int page_prod_stage = 0; + unsigned int page_prod_phase = 1; + unsigned int page_work_stage = 0; + unsigned int page_q = blockIdx.x; + unsigned int page_head = blockIdx.y; + unsigned int page_batch = blockIdx.z; + unsigned int _phase_work_full_3 = 0; + #pragma unroll 1 + for (unsigned int _work_3 = 0; _work_3 < total_work_items; _work_3++) { + int q_begin_3 = cum_seq_lens_q[(int)page_batch]; + int q_end_3 = cum_seq_lens_q[(int)(page_batch + 1)]; + int q_local_3 = (int)page_q; + int work_valid_3 = q_local_3 < q_end_3 - q_begin_3; + if (work_valid_3 != 0) { + int query_idx_3 = q_begin_3 + q_local_3; + int _max_0 = ((sparse_topk_lens[query_idx_3] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_3] + sparse_topk_lens_offset) : (0)); + int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); + int active_topk_3 = _min_0; + int num_tiles_3 = (active_topk_3 + 128 - 1) / 128; + int num_page_passes = (num_tiles_3 + 1) / 2; + int* swa_row = swa_indices + (query_idx_3 * swa_index_stride); + int* compressed_row = compressed_indices + (query_idx_3 * compressed_index_stride); + int _max_1 = (((active_topk_3 - 1) / 4 * 4) > (0) ? ((active_topk_3 - 1) / 4 * 4) : (0)); + int last_vec_base = _max_1; + #pragma unroll 1 + for (int page_pass = 0; page_pass < num_page_passes; page_pass++) { + #pragma unroll + for (int _duplicate = 0; _duplicate < 2; _duplicate++) { + mbarrier_wait(page_empty_addr + (page_prod_stage) * 8, page_prod_phase); + int page_dst = smem_page_offsets_addr + page_prod_stage * 1024; + #pragma unroll + for (int page_half = 0; page_half < 2; page_half++) { + int page_offset = page_pass * 256 + page_half * 128 + lane * 4; + int clamped_offset = ((page_offset < active_topk_3) ? page_offset : last_vec_base); + int* page_src = ((clamped_offset < 128) ? swa_row + clamped_offset : compressed_row + (clamped_offset - 128)); + asm volatile("cp.async.cg.shared::cta.global.L2::128B [%0], [%1], 16;" + :: "r"(page_dst + (page_half * 128 + lane * 4) * 4), "l"(page_src)); + } + asm volatile("cp.async.commit_group;"); + asm volatile("cp.async.wait_group 0;"); + mbarrier_arrive(page_full_addr + (page_prod_stage) * 8); + page_prod_stage += 1; + if (page_prod_stage == 6) { page_prod_stage = 0; page_prod_phase ^= 1; } + } + } + } + mbarrier_wait(work_full_addr + (page_work_stage) * 8, _phase_work_full_3); + uint32_t _clc_valid_0 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_0) + : "r"(work_response_addr + page_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_0 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_0) + : "r"(work_response_addr + page_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_1 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_1) + : "r"(work_response_addr + page_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_2 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_2) + : "r"(work_response_addr + page_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(work_empty_addr + (page_work_stage) * 8); + page_work_stage += 1; + if (page_work_stage == 2) { page_work_stage = 0; _phase_work_full_3 ^= 1; } + if (_clc_valid_0 == 0) { + break; + } + page_q = _clc_ctaid_0; + page_head = _clc_ctaid_1; + page_batch = _clc_ctaid_2; + } + } + } + // ---- Role: scheduler ---- + if (warp == 10) { + { // scheduler_main + unsigned int scheduler_work_stage = 0; + unsigned int scheduler_throttle_stage = 0; + unsigned int _phase_throttle_full = 0; + unsigned int _phase_work_empty = 1; + unsigned int _phase_work_full_4 = 0; + #pragma unroll 1 + for (unsigned int _work_4 = 0; _work_4 < total_work_items; _work_4++) { + mbarrier_wait(throttle_full_addr + (scheduler_throttle_stage) * 8, _phase_throttle_full); + mbarrier_arrive(throttle_empty_addr + (scheduler_throttle_stage) * 8); + scheduler_throttle_stage += 1; + if (scheduler_throttle_stage == 2) { scheduler_throttle_stage = 0; _phase_throttle_full ^= 1; } + if (elect_sync()) { + mbarrier_wait(work_empty_addr + (scheduler_work_stage) * 8, _phase_work_empty); + mbarrier_arrive_expect_tx(work_full_addr + (scheduler_work_stage) * 8, 16); + asm volatile( + "fence.proxy.async.shared::cta;\n\t" + "clusterlaunchcontrol.try_cancel.async.shared::cta" + ".mbarrier::complete_tx::bytes.b128" + " [%0], [%1];" + :: "r"(work_response_addr + scheduler_work_stage * 16 + 0 * 16), "r"(work_full_addr + scheduler_work_stage * 8) + : "memory"); + } + mbarrier_wait(work_full_addr + (scheduler_work_stage) * 8, _phase_work_full_4); + uint32_t _clc_valid_5 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_5) + : "r"(work_response_addr + scheduler_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_15 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_15) + : "r"(work_response_addr + scheduler_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_16 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_16) + : "r"(work_response_addr + scheduler_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_17 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_17) + : "r"(work_response_addr + scheduler_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(work_empty_addr + (scheduler_work_stage) * 8); + scheduler_work_stage += 1; + if (scheduler_work_stage == 2) { scheduler_work_stage = 0; _phase_work_empty ^= 1; _phase_work_full_4 ^= 1; } + if (_clc_valid_5 == 0) { + break; + } + } + } + } + // ---- Role: padding ---- + if (warp == 11) { + { // padding_main + unsigned int padding_work_stage = 0; + unsigned int _phase_work_full_5 = 0; + #pragma unroll 1 + for (unsigned int _work_5 = 0; _work_5 < total_work_items; _work_5++) { + mbarrier_wait(work_full_addr + (padding_work_stage) * 8, _phase_work_full_5); + uint32_t _clc_valid_6 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_6) + : "r"(work_response_addr + padding_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_18 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_18) + : "r"(work_response_addr + padding_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_19 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_19) + : "r"(work_response_addr + padding_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_20 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_20) + : "r"(work_response_addr + padding_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(work_empty_addr + (padding_work_stage) * 8); + padding_work_stage += 1; + if (padding_work_stage == 2) { padding_work_stage = 0; _phase_work_full_5 ^= 1; } + if (_clc_valid_6 == 0) { + break; + } + } + } + } + // ---- Role: load ---- + if (warp >= 12 && warp <= 15) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 72;"); + { // load_main + const int load_rank = warp - 12; + unsigned int q_prod_stage = 0; + unsigned int kv_prod_stage = 0; + unsigned int page_k_stage = 0; + unsigned int page_k_phase = 0; + unsigned int page_v_stage = 1; + unsigned int page_v_phase = 0; + unsigned int load_work_stage = 0; + unsigned int throttle_stage = 0; + unsigned int load_q = blockIdx.x; + unsigned int load_head = blockIdx.y; + unsigned int load_batch = blockIdx.z; + unsigned int _phase_throttle_empty = 1; + unsigned int _phase_q_empty = 1; + unsigned int _phase_kv_empty = 1; + unsigned int _phase_work_full_6 = 0; + #pragma unroll 1 + for (unsigned int _work_6 = 0; _work_6 < total_work_items; _work_6++) { + mbarrier_wait(throttle_empty_addr + (throttle_stage) * 8, _phase_throttle_empty); + mbarrier_arrive(throttle_full_addr + (throttle_stage) * 8); + throttle_stage += 1; + if (throttle_stage == 2) { throttle_stage = 0; _phase_throttle_empty ^= 1; } + int q_begin_4 = cum_seq_lens_q[(int)load_batch]; + int q_end_4 = cum_seq_lens_q[(int)(load_batch + 1)]; + int q_local_4 = (int)load_q; + int work_valid_4 = q_local_4 < q_end_4 - q_begin_4; + if (work_valid_4 != 0) { + int query_idx_4 = q_begin_4 + q_local_4; + int _max_2 = ((sparse_topk_lens[query_idx_4] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_4] + sparse_topk_lens_offset) : (0)); + int _min_1 = ((_max_2) < (sparse_topk) ? (_max_2) : (sparse_topk)); + int active_topk_4 = _min_1; + int num_tiles_4 = (active_topk_4 + 128 - 1) / 128; + mbarrier_wait(q_empty_addr + (q_prod_stage) * 8, _phase_q_empty); + if (load_rank == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(q_full_addr + (q_prod_stage) * 8, 32768); + } + } + asm volatile("griddepcontrol.wait;" ::: "memory"); + if (load_rank == 0) { + if (elect_sync()) { + #pragma unroll + for (int q_chunk = 0; q_chunk < 4; q_chunk++) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global.tile.mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(smem_q_addr + q_prod_stage * 32768 + (unsigned int)(q_chunk * 8192)), "l"((&tmap_q)), "r"(q_chunk * 128), "r"(0), "r"((int)load_head), "r"(query_idx_4), "r"(q_full_addr + (q_prod_stage) * 8) : "memory"); + } + } + } + q_prod_stage += 1; + if (q_prod_stage == 2) { q_prod_stage = 0; _phase_q_empty ^= 1; } + #pragma unroll 1 + for (int tile_2 = 0; tile_2 < num_tiles_4; tile_2++) { + if ((tile_2 & 1) == 0) { + mbarrier_wait(page_full_addr + (page_k_stage) * 8, page_k_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + } + int page_k_base = smem_page_offsets_addr + page_k_stage * 1024; + #pragma unroll + for (int k_half = 0; k_half < 2; k_half++) { + mbarrier_wait(kv_empty_addr + (kv_prod_stage) * 8, _phase_kv_empty); + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (kv_prod_stage) * 8, 8192); + if (tile_2 == 0) { + #pragma unroll + for (int local_group = 0; local_group < 8; local_group++) { + int group = load_rank * 8 + local_group; + int group_offset = (tile_2 & 1) * 128 + group * 4; + int rows[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&rows[0])), "=r"(*reinterpret_cast(&rows[(0) + 1])), "=r"(*reinterpret_cast(&rows[(0) + 2])), "=r"(*reinterpret_cast(&rows[(0) + 3])) + : "r"(page_k_base + group_offset * 4)); + int k_dst = smem_k_addr + kv_prod_stage * 32768; + tma_gather4_gmem2smem(k_dst + group * 512, (&tmap_swa_kv), k_half * 256, rows[0], rows[1], rows[2], rows[3], kv_full_addr + (kv_prod_stage) * 8); + tma_gather4_gmem2smem(k_dst + 16384 + group * 512, (&tmap_swa_kv), k_half * 256 + 128, rows[0], rows[1], rows[2], rows[3], kv_full_addr + (kv_prod_stage) * 8); + } + } else { + #pragma unroll + for (int local_group_1 = 0; local_group_1 < 8; local_group_1++) { + int group_1 = load_rank * 8 + local_group_1; + int group_offset_1 = (tile_2 & 1) * 128 + group_1 * 4; + int rows_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&rows_1[0])), "=r"(*reinterpret_cast(&rows_1[(0) + 1])), "=r"(*reinterpret_cast(&rows_1[(0) + 2])), "=r"(*reinterpret_cast(&rows_1[(0) + 3])) + : "r"(page_k_base + group_offset_1 * 4)); + int k_dst_1 = smem_k_addr + kv_prod_stage * 32768; + tma_gather4_gmem2smem(k_dst_1 + group_1 * 512, (&tmap_compressed_kv), k_half * 256, rows_1[0], rows_1[1], rows_1[2], rows_1[3], kv_full_addr + (kv_prod_stage) * 8); + tma_gather4_gmem2smem(k_dst_1 + 16384 + group_1 * 512, (&tmap_compressed_kv), k_half * 256 + 128, rows_1[0], rows_1[1], rows_1[2], rows_1[3], kv_full_addr + (kv_prod_stage) * 8); + } + } + } + kv_prod_stage += 1; + if (kv_prod_stage == 4) { kv_prod_stage = 0; _phase_kv_empty ^= 1; } + } + if ((tile_2 & 1) != 0 || tile_2 == num_tiles_4 - 1) { + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(page_empty_addr + (page_k_stage) * 8); + page_k_stage += 1; + if (page_k_stage == 6) { page_k_stage = 0; page_k_phase ^= 1; } + page_k_stage += 1; + if (page_k_stage == 6) { page_k_stage = 0; page_k_phase ^= 1; } + } + if (tile_2 > 0) { + int prev_tile = tile_2 - 1; + if ((prev_tile & 1) == 0) { + mbarrier_wait(page_full_addr + (page_v_stage) * 8, page_v_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + } + int page_v_base = smem_page_offsets_addr + page_v_stage * 1024; + #pragma unroll + for (int v_half_2 = 0; v_half_2 < 2; v_half_2++) { + mbarrier_wait(kv_empty_addr + (kv_prod_stage) * 8, _phase_kv_empty); + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (kv_prod_stage) * 8, 8192); + if (prev_tile == 0) { + #pragma unroll + for (int local_group_2 = 0; local_group_2 < 8; local_group_2++) { + int group_2 = load_rank * 8 + local_group_2; + int group_offset_2 = (prev_tile & 1) * 128 + group_2 * 4; + int rows_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&rows_2[0])), "=r"(*reinterpret_cast(&rows_2[(0) + 1])), "=r"(*reinterpret_cast(&rows_2[(0) + 2])), "=r"(*reinterpret_cast(&rows_2[(0) + 3])) + : "r"(page_v_base + group_offset_2 * 4)); + int v_dst = smem_v_addr + kv_prod_stage * 32768; + tma_gather4_gmem2smem(v_dst + group_2 * 512, (&tmap_swa_kv), v_half_2 * 256, rows_2[0], rows_2[1], rows_2[2], rows_2[3], kv_full_addr + (kv_prod_stage) * 8); + tma_gather4_gmem2smem(v_dst + 16384 + group_2 * 512, (&tmap_swa_kv), v_half_2 * 256 + 128, rows_2[0], rows_2[1], rows_2[2], rows_2[3], kv_full_addr + (kv_prod_stage) * 8); + } + } else { + #pragma unroll + for (int local_group_3 = 0; local_group_3 < 8; local_group_3++) { + int group_3 = load_rank * 8 + local_group_3; + int group_offset_3 = (prev_tile & 1) * 128 + group_3 * 4; + int rows_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&rows_3[0])), "=r"(*reinterpret_cast(&rows_3[(0) + 1])), "=r"(*reinterpret_cast(&rows_3[(0) + 2])), "=r"(*reinterpret_cast(&rows_3[(0) + 3])) + : "r"(page_v_base + group_offset_3 * 4)); + int v_dst_1 = smem_v_addr + kv_prod_stage * 32768; + tma_gather4_gmem2smem(v_dst_1 + group_3 * 512, (&tmap_compressed_kv), v_half_2 * 256, rows_3[0], rows_3[1], rows_3[2], rows_3[3], kv_full_addr + (kv_prod_stage) * 8); + tma_gather4_gmem2smem(v_dst_1 + 16384 + group_3 * 512, (&tmap_compressed_kv), v_half_2 * 256 + 128, rows_3[0], rows_3[1], rows_3[2], rows_3[3], kv_full_addr + (kv_prod_stage) * 8); + } + } + } + kv_prod_stage += 1; + if (kv_prod_stage == 4) { kv_prod_stage = 0; _phase_kv_empty ^= 1; } + } + if ((prev_tile & 1) != 0) { + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(page_empty_addr + (page_v_stage) * 8); + page_v_stage += 1; + if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } + page_v_stage += 1; + if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } + } + } + } + int last_tile = num_tiles_4 - 1; + mbarrier_wait(page_full_addr + (page_v_stage) * 8, page_v_phase); + asm volatile("tcgen05.fence::after_thread_sync;"); + int page_v_base_1 = smem_page_offsets_addr + page_v_stage * 1024; + #pragma unroll + for (int v_half_3 = 0; v_half_3 < 2; v_half_3++) { + mbarrier_wait(kv_empty_addr + (kv_prod_stage) * 8, _phase_kv_empty); + if (elect_sync()) { + mbarrier_arrive_expect_tx(kv_full_addr + (kv_prod_stage) * 8, 8192); + if (last_tile == 0) { + #pragma unroll + for (int local_group_4 = 0; local_group_4 < 8; local_group_4++) { + int group_4 = load_rank * 8 + local_group_4; + int group_offset_4 = (last_tile & 1) * 128 + group_4 * 4; + int rows_4[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&rows_4[0])), "=r"(*reinterpret_cast(&rows_4[(0) + 1])), "=r"(*reinterpret_cast(&rows_4[(0) + 2])), "=r"(*reinterpret_cast(&rows_4[(0) + 3])) + : "r"(page_v_base_1 + group_offset_4 * 4)); + int v_dst_2 = smem_v_addr + kv_prod_stage * 32768; + tma_gather4_gmem2smem(v_dst_2 + group_4 * 512, (&tmap_swa_kv), v_half_3 * 256, rows_4[0], rows_4[1], rows_4[2], rows_4[3], kv_full_addr + (kv_prod_stage) * 8); + tma_gather4_gmem2smem(v_dst_2 + 16384 + group_4 * 512, (&tmap_swa_kv), v_half_3 * 256 + 128, rows_4[0], rows_4[1], rows_4[2], rows_4[3], kv_full_addr + (kv_prod_stage) * 8); + } + } else { + #pragma unroll + for (int local_group_5 = 0; local_group_5 < 8; local_group_5++) { + int group_5 = load_rank * 8 + local_group_5; + int group_offset_5 = (last_tile & 1) * 128 + group_5 * 4; + int rows_5[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&rows_5[0])), "=r"(*reinterpret_cast(&rows_5[(0) + 1])), "=r"(*reinterpret_cast(&rows_5[(0) + 2])), "=r"(*reinterpret_cast(&rows_5[(0) + 3])) + : "r"(page_v_base_1 + group_offset_5 * 4)); + int v_dst_3 = smem_v_addr + kv_prod_stage * 32768; + tma_gather4_gmem2smem(v_dst_3 + group_5 * 512, (&tmap_compressed_kv), v_half_3 * 256, rows_5[0], rows_5[1], rows_5[2], rows_5[3], kv_full_addr + (kv_prod_stage) * 8); + tma_gather4_gmem2smem(v_dst_3 + 16384 + group_5 * 512, (&tmap_compressed_kv), v_half_3 * 256 + 128, rows_5[0], rows_5[1], rows_5[2], rows_5[3], kv_full_addr + (kv_prod_stage) * 8); + } + } + } + kv_prod_stage += 1; + if (kv_prod_stage == 4) { kv_prod_stage = 0; _phase_kv_empty ^= 1; } + } + asm volatile("tcgen05.fence::before_thread_sync;"); + mbarrier_arrive(page_empty_addr + (page_v_stage) * 8); + page_v_stage += 1; + if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } + page_v_stage += 1; + if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } + } + mbarrier_wait(work_full_addr + (load_work_stage) * 8, _phase_work_full_6); + uint32_t _clc_valid_1 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "selp.u32 %0, 1, 0, p1;\n\t" + "}\n" + : "=r"(_clc_valid_1) + : "r"(work_response_addr + load_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_3 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_3) + : "r"(work_response_addr + load_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_4 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_4) + : "r"(work_response_addr + load_work_stage * 16 + 0 * 16) + : "memory"); + uint32_t _clc_ctaid_5 = 0; + asm volatile( + "{\n\t" + ".reg .pred p1;\n\t" + ".reg .b128 clc_r;\n\t" + "ld.shared.b128 clc_r, [%1];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" + "}\n" + : "=r"(_clc_ctaid_5) + : "r"(work_response_addr + load_work_stage * 16 + 0 * 16) + : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(work_empty_addr + (load_work_stage) * 8); + load_work_stage += 1; + if (load_work_stage == 2) { load_work_stage = 0; _phase_work_full_6 ^= 1; } + if (_clc_valid_1 == 0) { + break; + } + load_q = _clc_ctaid_3; + load_head = _clc_ctaid_4; + load_batch = _clc_ctaid_5; + } + } + } + + // Cleanup +} + +} // extern "C" diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_family_5383a615a7e502f6fdac_binding.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_family_5383a615a7e502f6fdac_binding.cu new file mode 100644 index 00000000000..3047ed0720b --- /dev/null +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_family_5383a615a7e502f6fdac_binding.cu @@ -0,0 +1,703 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +// Generated by cake.runtime.ffi_dispatch — do not edit. +// tvm-ffi dispatch module for kernel family 'cake_dsv4_family_5383a615a7e502f6fdac'. +// dispatch-policy-revision: unpinned +#include +#include +#include +#include + +#include "tvm_ffi_utils.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +extern "C" __global__ void kernel_cake_dsv4_d906a03bd66a7488fc29(__nv_bfloat16* __restrict__ partial_O, float* __restrict__ partial_lse, __nv_bfloat16* __restrict__ O, int num_q_heads, int num_split); +extern "C" __global__ void kernel_cake_dsv4_714df747e199d0adbe08(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, const __grid_constant__ CUtensorMap O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); + +namespace cake_host_shim_3ce874e06ea0fc76 { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t CakeDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +namespace variant_cake_dsv4_714df747e199d0adbe08_sm_100a_591fb3561ebd { + +// 4D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + int64_t d3 = t.size(t.ndim() - 3); + TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(d3)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 64u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_swa_k' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_swa_k(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_swa_k' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_swa_k' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_swa_k' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_swa_k"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_swa_k' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_swa_k' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_k'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_swa_k') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_compressed_k' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_compressed_k(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_compressed_k' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_compressed_k' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_compressed_k' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_compressed_k"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_compressed_k' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_compressed_k' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_k'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_k') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_swa_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_swa_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_swa_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_swa_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_swa_v' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_swa_v"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_swa_v' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_swa_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_v'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_swa_v') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_compressed_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_compressed_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_compressed_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_compressed_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_compressed_v' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_compressed_v"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_compressed_v' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_compressed_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_v'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_v') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'O' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_O(const TensorView& t, int64_t tp_num_heads, int64_t tp_num_query_tokens) { + TVM_FFI_CHECK(t.ndim() >= 1, ValueError) + << "TMA source 'O' must have at least 1 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'O' must have unit innermost stride, got " << t.stride(-1); + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(8), (uint64_t)(5), (uint64_t)((tp_num_query_tokens * tp_num_heads))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'O' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1] && 1u <= global_dim[2] && 32u <= global_dim[3], ValueError) + << "TMA box (64, 1, 1, 32) exceeds resolved global dims for 'O'"; + int64_t carrier_stride_0 = 64; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 512; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = 2560; + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 1u, 1u, 32u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'O') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_compressed_k, TensorView arg_tmap_swa_v, TensorView arg_tmap_compressed_v, TensorView arg_O, TensorView arg_partial_lse, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_num_query_tokens, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckContiguous(arg_tmap_q, "tmap_q"); + CheckCudaTensor(arg_tmap_swa_k, "tmap_swa_k"); + CheckDtype(arg_tmap_swa_k, "tmap_swa_k", 4, 16, 1); + CheckCudaTensor(arg_tmap_compressed_k, "tmap_compressed_k"); + CheckDtype(arg_tmap_compressed_k, "tmap_compressed_k", 4, 16, 1); + CheckCudaTensor(arg_tmap_swa_v, "tmap_swa_v"); + CheckDtype(arg_tmap_swa_v, "tmap_swa_v", 4, 16, 1); + CheckCudaTensor(arg_tmap_compressed_v, "tmap_compressed_v"); + CheckDtype(arg_tmap_compressed_v, "tmap_compressed_v", 4, 16, 1); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_partial_lse, "partial_lse"); + CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); + CheckContiguous(arg_partial_lse, "partial_lse"); + CheckCudaTensor(arg_swa_indices, "swa_indices"); + CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); + CheckContiguous(arg_swa_indices, "swa_indices"); + CheckCudaTensor(arg_compressed_indices, "compressed_indices"); + CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); + CheckContiguous(arg_compressed_indices, "compressed_indices"); + CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); + CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); + CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); + CheckContiguous(arg_bmm1_scale, "bmm1_scale"); + CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); + CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); + CheckContiguous(arg_bmm2_scale, "bmm2_scale"); + TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) + << "scalar 'num_heads' value " << arg_num_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_query_tokens >= -2147483648LL && arg_num_query_tokens <= 2147483647LL, ValueError) + << "scalar 'num_query_tokens' value " << arg_num_query_tokens + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) + << "scalar 'swa_index_stride' value " << arg_swa_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) + << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) + << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) + << "scalar 'sparse_topk' value " << arg_sparse_topk + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) + << "scalar 'has_sinks' value " << arg_has_sinks + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_swa_k, arg_tmap_q, "tmap_swa_k", "tmap_q"); + CheckSameCudaDevice(arg_tmap_compressed_k, arg_tmap_q, "tmap_compressed_k", "tmap_q"); + CheckSameCudaDevice(arg_tmap_swa_v, arg_tmap_q, "tmap_swa_v", "tmap_q"); + CheckSameCudaDevice(arg_tmap_compressed_v, arg_tmap_q, "tmap_compressed_v", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_partial_lse, arg_tmap_q, "partial_lse", "tmap_q"); + CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); + CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); + CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); + CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") must be divisible by cluster dims (2, 1, 1)"; + + + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_swa_k = EncodeTma_tmap_swa_k(arg_tmap_swa_k); + CUtensorMap p_tmap_compressed_k = EncodeTma_tmap_compressed_k(arg_tmap_compressed_k); + CUtensorMap p_tmap_swa_v = EncodeTma_tmap_swa_v(arg_tmap_swa_v); + CUtensorMap p_tmap_compressed_v = EncodeTma_tmap_compressed_v(arg_tmap_compressed_v); + CUtensorMap p_O = EncodeTma_O(arg_O, (int64_t)arg_num_heads, (int64_t)arg_num_query_tokens); + float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); + int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); + int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); + int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); + float* p_sinks = static_cast(arg_sinks.data_ptr()); + float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); + float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); + int32_t v_num_heads = (int32_t)arg_num_heads; + int32_t v_num_query_tokens = (int32_t)arg_num_query_tokens; + int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; + int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; + int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; + int32_t v_sparse_topk = (int32_t)arg_sparse_topk; + int32_t v_has_sinks = (int32_t)arg_has_sinks; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + void* kargs[] = {&p_tmap_q, &p_tmap_swa_k, &p_tmap_compressed_k, &p_tmap_swa_v, &p_tmap_compressed_v, &p_O, &p_partial_lse, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_num_query_tokens, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_has_sinks, &v_total_work_items}; + + static std::mutex smem_mu; + static auto* smem_status_by_device = + new std::unordered_map(); + cudaError_t smem_status = cudaSuccess; + { + std::lock_guard lock(smem_mu); + auto it = smem_status_by_device->find(dev.device_id); + if (it == smem_status_by_device->end()) { + int current_device = -1; + smem_status = cudaGetDevice(¤t_device); + if (smem_status == cudaSuccess && current_device != dev.device_id) { + smem_status = cudaErrorInvalidDevice; + } + if (smem_status == cudaSuccess) { + smem_status = cudaFuncSetAttribute( + reinterpret_cast(kernel_cake_dsv4_714df747e199d0adbe08), + cudaFuncAttributeMaxDynamicSharedMemorySize, + 231808); + } + smem_status_by_device->emplace(dev.device_id, smem_status); + } else { + smem_status = it->second; + } + } + TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) + << "cudaFuncSetAttribute for kernel_cake_dsv4_714df747e199d0adbe08 failed: " + << cudaGetErrorString(smem_status); + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(512u, 1u, 1u); + + // Extended direct-source launch for a linked CUDA kernel symbol. + cudaLaunchAttribute attrs[1]{}; + int n = 0; + attrs[n].id = cudaLaunchAttributeClusterDimension; + attrs[n].val.clusterDim.x = 2u; + attrs[n].val.clusterDim.y = 1u; + attrs[n].val.clusterDim.z = 1u; + ++n; + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes = 231808u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; + cudaError_t launch_status = cudaLaunchKernelExC( + &config, reinterpret_cast(kernel_cake_dsv4_714df747e199d0adbe08), kargs); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernelExC for kernel_cake_dsv4_714df747e199d0adbe08 failed: " + << cudaGetErrorString(launch_status); +} +} // namespace variant_cake_dsv4_714df747e199d0adbe08_sm_100a_591fb3561ebd + +namespace variant_cake_dsv4_d906a03bd66a7488fc29_sm_100a_cd6fb31da15f { + +void Run(TensorView arg_partial_O, TensorView arg_partial_lse, TensorView arg_O, int64_t arg_num_q_heads, int64_t arg_num_split, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_partial_O.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_partial_O, "partial_O"); + CheckDtype(arg_partial_O, "partial_O", 4, 16, 1); + CheckContiguous(arg_partial_O, "partial_O"); + CheckCudaTensor(arg_partial_lse, "partial_lse"); + CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); + CheckContiguous(arg_partial_lse, "partial_lse"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + TVM_FFI_CHECK(arg_num_q_heads >= -2147483648LL && arg_num_q_heads <= 2147483647LL, ValueError) + << "scalar 'num_q_heads' value " << arg_num_q_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_split >= -2147483648LL && arg_num_split <= 2147483647LL, ValueError) + << "scalar 'num_split' value " << arg_num_split + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_partial_lse, arg_partial_O, "partial_lse", "partial_O"); + CheckSameCudaDevice(arg_O, arg_partial_O, "O", "partial_O"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + + + __nv_bfloat16* p_partial_O = static_cast<__nv_bfloat16*>(arg_partial_O.data_ptr()); + float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); + __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); + int32_t v_num_q_heads = (int32_t)arg_num_q_heads; + int32_t v_num_split = (int32_t)arg_num_split; + void* kargs[] = {&p_partial_O, &p_partial_lse, &p_O, &v_num_q_heads, &v_num_split}; + + + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(128u, 1u, 1u); + + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes = 0u; + config.stream = stream; + cudaError_t launch_status = cudaLaunchKernelExC( + &config, reinterpret_cast(kernel_cake_dsv4_d906a03bd66a7488fc29), kargs); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernelExC for kernel_cake_dsv4_d906a03bd66a7488fc29 failed: " + << cudaGetErrorString(launch_status); +} +} // namespace variant_cake_dsv4_d906a03bd66a7488fc29_sm_100a_cd6fb31da15f + +int SelectVariant(int32_t num_heads, int32_t num_query_tokens, int32_t swa_index_stride, int32_t compressed_index_stride, int32_t sparse_topk_lens_offset, int32_t sparse_topk, int32_t has_sinks, int32_t total_work_items, int32_t num_split) { + if ((num_split == 5)) return 0; + return -1; +} + +void Dispatch(tvm::ffi::TensorView arg_Q, tvm::ffi::TensorView arg_SWA_cache, tvm::ffi::TensorView arg_compressed_KV_cache, tvm::ffi::TensorView arg_partial_O, tvm::ffi::TensorView arg_partial_lse, tvm::ffi::TensorView arg_swa_indices, tvm::ffi::TensorView arg_compressed_indices, tvm::ffi::TensorView arg_sparse_topk_lens, tvm::ffi::TensorView arg_sinks, tvm::ffi::TensorView arg_bmm1_scale, tvm::ffi::TensorView arg_bmm2_scale, tvm::ffi::TensorView arg_O, int32_t num_heads, int32_t num_query_tokens, int32_t swa_index_stride, int32_t compressed_index_stride, int32_t sparse_topk_lens_offset, int32_t sparse_topk, int32_t has_sinks, int32_t total_work_items, int32_t num_split, int64_t cuda_stream_ptr) { + CheckCudaTensor(arg_Q, "Q"); + CheckCudaTensor(arg_SWA_cache, "SWA_cache"); + CheckCudaTensor(arg_compressed_KV_cache, "compressed_KV_cache"); + CheckCudaTensor(arg_partial_O, "partial_O"); + CheckCudaTensor(arg_partial_lse, "partial_lse"); + CheckCudaTensor(arg_swa_indices, "swa_indices"); + CheckCudaTensor(arg_compressed_indices, "compressed_indices"); + CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); + CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); + CheckCudaTensor(arg_O, "O"); + cudaStream_t stream = reinterpret_cast(static_cast(cuda_stream_ptr)); + int v = SelectVariant(num_heads, num_query_tokens, swa_index_stride, compressed_index_stride, sparse_topk_lens_offset, sparse_topk, has_sinks, total_work_items, num_split); + switch (v) { + case 0: { + { + int64_t gx = (total_work_items * 2); + int64_t gy = 1; + int64_t gz = 1; + cake_host_shim_3ce874e06ea0fc76::variant_cake_dsv4_714df747e199d0adbe08_sm_100a_591fb3561ebd::Run(arg_Q, arg_SWA_cache, arg_compressed_KV_cache, arg_SWA_cache, arg_compressed_KV_cache, arg_partial_O, arg_partial_lse, arg_swa_indices, arg_compressed_indices, arg_sparse_topk_lens, arg_sinks, arg_bmm1_scale, arg_bmm2_scale, num_heads, num_query_tokens, swa_index_stride, compressed_index_stride, sparse_topk_lens_offset, sparse_topk, has_sinks, total_work_items, gx, gy, gz, stream); + } + { + int64_t gx = num_query_tokens; + int64_t gy = num_heads; + int64_t gz = 1; + cake_host_shim_3ce874e06ea0fc76::variant_cake_dsv4_d906a03bd66a7488fc29_sm_100a_cd6fb31da15f::Run(arg_partial_O, arg_partial_lse, arg_O, num_heads, num_split, gx, gy, gz, stream); + } + return; + } + default: + TVM_FFI_CHECK(false, ValueError) << "no dispatch route for cake_dsv4_family_5383a615a7e502f6fdac shape (num_heads=" << num_heads << ", num_query_tokens=" << num_query_tokens << ", swa_index_stride=" << swa_index_stride << ", compressed_index_stride=" << compressed_index_stride << ", sparse_topk_lens_offset=" << sparse_topk_lens_offset << ", sparse_topk=" << sparse_topk << ", has_sinks=" << has_sinks << ", total_work_items=" << total_work_items << ", num_split=" << num_split << ")"; + } +} + +int64_t RouteOf(int32_t num_heads, int32_t num_query_tokens, int32_t swa_index_stride, int32_t compressed_index_stride, int32_t sparse_topk_lens_offset, int32_t sparse_topk, int32_t has_sinks, int32_t total_work_items, int32_t num_split) { return SelectVariant(num_heads, num_query_tokens, swa_index_stride, compressed_index_stride, sparse_topk_lens_offset, sparse_topk, has_sinks, total_work_items, num_split); } + +} // namespace cake_host_shim_3ce874e06ea0fc76 + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_3ce874e06ea0fc76::Dispatch); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(route_of, cake_host_shim_3ce874e06ea0fc76::RouteOf); diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_family_9c96528827a52c42bbd7_binding.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_family_9c96528827a52c42bbd7_binding.cu new file mode 100644 index 00000000000..d224bcffd84 --- /dev/null +++ b/csrc/cake_dsv4/sm_100a/cake_dsv4_family_9c96528827a52c42bbd7_binding.cu @@ -0,0 +1,703 @@ +/* + * Copyright (c) 2023 by FlashInfer 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. + */ + +// Generated by cake.runtime.ffi_dispatch — do not edit. +// tvm-ffi dispatch module for kernel family 'cake_dsv4_family_9c96528827a52c42bbd7'. +// dispatch-policy-revision: unpinned +#include +#include +#include +#include + +#include "tvm_ffi_utils.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +extern "C" __global__ void kernel_cake_dsv4_d50dd820272ad24f7027(__nv_bfloat16* __restrict__ partial_O, float* __restrict__ partial_lse, __nv_bfloat16* __restrict__ O, int num_q_heads, int num_split); +extern "C" __global__ void kernel_cake_dsv4_42fa8e98da5fa46d1a94(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, const __grid_constant__ CUtensorMap O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); + +namespace cake_host_shim_6534755d2bf16a2e { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +class ScopedCudaDevice { + public: + explicit ScopedCudaDevice(int device_id) { + cudaError_t error = cudaGetDevice(&previous_device_); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed before host-shim launch: cudaError=" + << static_cast(error); + if (previous_device_ != device_id) { + error = cudaSetDevice(device_id); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaSetDevice failed before host-shim launch for cuda:" + << device_id << ": cudaError=" << static_cast(error); + restore_ = true; + } + } + + ScopedCudaDevice(const ScopedCudaDevice&) = delete; + ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; + + ~ScopedCudaDevice() noexcept { + if (restore_) { + (void)cudaSetDevice(previous_device_); + } + } + + private: + int previous_device_ = -1; + bool restore_ = false; +}; + +inline int64_t CakeDeviceMultiprocessorCount(int device_id) { + constexpr int kMaxCachedCudaDevices = 64; + TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) + << "physical-SM-count cache does not cover cuda:" << device_id; + static std::atomic count_by_device[kMaxCachedCudaDevices]{}; + int cached = count_by_device[device_id].load(std::memory_order_acquire); + if (cached > 0) return cached; + + int count = 0; + cudaError_t error = cudaDeviceGetAttribute( + &count, cudaDevAttrMultiProcessorCount, device_id); + TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) + << "querying multiProcessorCount failed for cuda:" << device_id + << ": cudaError=" << static_cast(error) << ", count=" << count; + // Concurrent first launches may repeat the immutable device query, but all + // publication is atomic and every later hot-path lookup is one acquire load. + count_by_device[device_id].store(count, std::memory_order_release); + return count; +} + +inline void CheckCudaTensor(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; +} + +inline void CheckSameCudaDevice( + const TensorView& t, + const TensorView& reference, + const char* name, + const char* reference_name) { + TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) + << name << " must be on the same CUDA device as " << reference_name + << ": got cuda:" << t.device().device_id + << " versus cuda:" << reference.device().device_id; +} + +inline void CheckCurrentCudaDevice( + const TensorView& reference, + const char* reference_name) { + int current_device = -1; + cudaError_t error = cudaGetDevice(¤t_device); + TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) + << "cudaGetDevice failed while validating " << reference_name + << ": cudaError=" << static_cast(error); + TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) + << "current CUDA device must match " << reference_name + << ": current=cuda:" << current_device + << ", tensor=cuda:" << reference.device().device_id; +} + +inline void CheckContiguous(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits + << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits + << ", lanes=" << (int)d.lanes << ")"; +} + +// A logical axis.outer(trailing) folds every source dim above the trailing +// dimensions. Shape products are independent of physical strides, so verify +// the leading dimensions form one dense row-major chain instead of inventing +// a "folded stride". The descriptor reads its exact adjacent physical step +// separately through stride[-(trailing + 1)]. +inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing + << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) + << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) + << ", expected " << expected; + } + } +} + +namespace variant_cake_dsv4_42fa8e98da5fa46d1a94_sm_100a_6276393afd19 { + +// 4D TMA descriptor for buffer 'tmap_q' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + int64_t d2 = t.size(t.ndim() - 2); + int64_t d3 = t.size(t.ndim() - 3); + TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) + << "TMA source 'tmap_q' trailing dims must be positive"; + TVM_FFI_CHECK(d1 % 64 == 0, ValueError) + << "TMA source 'tmap_q' extent " << d1 + << " must divide exactly by " << 64; + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(d3)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) + << "TMA box (64, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; + int64_t carrier_stride_0 = d1; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 64; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = (d2 * d1); + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 64u, 1u, 1u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_swa_k' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_swa_k(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_swa_k' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_swa_k' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_swa_k' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_swa_k"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_swa_k' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_swa_k' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_k'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_swa_k') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_compressed_k' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_compressed_k(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_compressed_k' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_compressed_k' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_compressed_k' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_compressed_k"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_compressed_k' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_compressed_k' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_k'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_k') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_swa_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_swa_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_swa_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_swa_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_swa_v' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_swa_v"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_swa_v' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_swa_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_v'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_swa_v') failed: CUresult=" << (int)r; + return tm; +} + +// 2D TMA descriptor for buffer 'tmap_compressed_v' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_tmap_compressed_v(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 2, ValueError) + << "TMA source 'tmap_compressed_v' must have at least 2 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'tmap_compressed_v' must have unit innermost stride, got " << t.stride(-1); + int64_t d1 = t.size(t.ndim() - 1); + TVM_FFI_CHECK(d1 > 0, ValueError) + << "TMA source 'tmap_compressed_v' trailing dims must be positive"; + int64_t outer1 = t.numel() / (d1); + CheckDenseLeadingFold(t, 1, "tmap_compressed_v"); + int64_t s2 = t.stride(t.ndim() - 2) * 1; + TVM_FFI_CHECK(s2 > 0, ValueError) + << "TMA source 'tmap_compressed_v' physical strides must be positive"; + uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) + << "TMA descriptor for 'tmap_compressed_v' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) + << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_v'"; + int64_t carrier_stride_0 = s2; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 to a non-whole-byte offset"; + uint64_t global_strides[1] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + }; + uint32_t box_dim[2] = {64u, 1u}; + uint32_t elem_strides[2] = {1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_v') failed: CUresult=" << (int)r; + return tm; +} + +// 4D TMA descriptor for buffer 'O' — compiled from the +// descriptor's std.Expr global_dim/global_strides/checks record. +inline CUtensorMap EncodeTma_O(const TensorView& t, int64_t tp_num_heads, int64_t tp_num_query_tokens) { + TVM_FFI_CHECK(t.ndim() >= 1, ValueError) + << "TMA source 'O' must have at least 1 dimensions, got ndim=" << t.ndim(); + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << "TMA source 'O' must have unit innermost stride, got " << t.stride(-1); + uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(8), (uint64_t)(4), (uint64_t)((tp_num_query_tokens * tp_num_heads))}; + TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) + << "TMA descriptor for 'O' resolved a non-positive global dim"; + TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1] && 1u <= global_dim[2] && 32u <= global_dim[3], ValueError) + << "TMA box (64, 1, 1, 32) exceeds resolved global dims for 'O'"; + int64_t carrier_stride_0 = 64; + TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 negative"; + TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 zero while global dimension 1 is not 1"; + TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 1 to a non-whole-byte offset"; + int64_t carrier_stride_1 = 512; + TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 negative"; + TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 zero while global dimension 2 is not 1"; + TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 2 to a non-whole-byte offset"; + int64_t carrier_stride_2 = 2048; + TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 negative"; + TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 zero while global dimension 3 is not 1"; + TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) + << "TMA descriptor for 'O' resolved global stride 3 to a non-whole-byte offset"; + uint64_t global_strides[3] = { + (uint64_t)((carrier_stride_0 * 16) / 8), + (uint64_t)((carrier_stride_1 * 16) / 8), + (uint64_t)((carrier_stride_2 * 16) / 8), + }; + uint32_t box_dim[4] = {64u, 1u, 1u, 32u}; + uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + CUtensorMap tm{}; + const void* tensor_base = static_cast(t.data_ptr()) + 0u; + CUresult r = cuTensorMapEncodeTiled( + &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) + << "cuTensorMapEncodeTiled (4D, 'O') failed: CUresult=" << (int)r; + return tm; +} + +void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_compressed_k, TensorView arg_tmap_swa_v, TensorView arg_tmap_compressed_v, TensorView arg_O, TensorView arg_partial_lse, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_num_query_tokens, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_tmap_q.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_tmap_q, "tmap_q"); + CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); + CheckContiguous(arg_tmap_q, "tmap_q"); + CheckCudaTensor(arg_tmap_swa_k, "tmap_swa_k"); + CheckDtype(arg_tmap_swa_k, "tmap_swa_k", 4, 16, 1); + CheckCudaTensor(arg_tmap_compressed_k, "tmap_compressed_k"); + CheckDtype(arg_tmap_compressed_k, "tmap_compressed_k", 4, 16, 1); + CheckCudaTensor(arg_tmap_swa_v, "tmap_swa_v"); + CheckDtype(arg_tmap_swa_v, "tmap_swa_v", 4, 16, 1); + CheckCudaTensor(arg_tmap_compressed_v, "tmap_compressed_v"); + CheckDtype(arg_tmap_compressed_v, "tmap_compressed_v", 4, 16, 1); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + CheckCudaTensor(arg_partial_lse, "partial_lse"); + CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); + CheckContiguous(arg_partial_lse, "partial_lse"); + CheckCudaTensor(arg_swa_indices, "swa_indices"); + CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); + CheckContiguous(arg_swa_indices, "swa_indices"); + CheckCudaTensor(arg_compressed_indices, "compressed_indices"); + CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); + CheckContiguous(arg_compressed_indices, "compressed_indices"); + CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); + CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckDtype(arg_sinks, "sinks", 2, 32, 1); + CheckContiguous(arg_sinks, "sinks"); + CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); + CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); + CheckContiguous(arg_bmm1_scale, "bmm1_scale"); + CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); + CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); + CheckContiguous(arg_bmm2_scale, "bmm2_scale"); + TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) + << "scalar 'num_heads' value " << arg_num_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_query_tokens >= -2147483648LL && arg_num_query_tokens <= 2147483647LL, ValueError) + << "scalar 'num_query_tokens' value " << arg_num_query_tokens + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) + << "scalar 'swa_index_stride' value " << arg_swa_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) + << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) + << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) + << "scalar 'sparse_topk' value " << arg_sparse_topk + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) + << "scalar 'has_sinks' value " << arg_has_sinks + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) + << "scalar 'total_work_items' value " << arg_total_work_items + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_tmap_swa_k, arg_tmap_q, "tmap_swa_k", "tmap_q"); + CheckSameCudaDevice(arg_tmap_compressed_k, arg_tmap_q, "tmap_compressed_k", "tmap_q"); + CheckSameCudaDevice(arg_tmap_swa_v, arg_tmap_q, "tmap_swa_v", "tmap_q"); + CheckSameCudaDevice(arg_tmap_compressed_v, arg_tmap_q, "tmap_compressed_v", "tmap_q"); + CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); + CheckSameCudaDevice(arg_partial_lse, arg_tmap_q, "partial_lse", "tmap_q"); + CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); + CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); + CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); + CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); + CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); + CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") must be divisible by cluster dims (2, 1, 1)"; + + + CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); + CUtensorMap p_tmap_swa_k = EncodeTma_tmap_swa_k(arg_tmap_swa_k); + CUtensorMap p_tmap_compressed_k = EncodeTma_tmap_compressed_k(arg_tmap_compressed_k); + CUtensorMap p_tmap_swa_v = EncodeTma_tmap_swa_v(arg_tmap_swa_v); + CUtensorMap p_tmap_compressed_v = EncodeTma_tmap_compressed_v(arg_tmap_compressed_v); + CUtensorMap p_O = EncodeTma_O(arg_O, (int64_t)arg_num_heads, (int64_t)arg_num_query_tokens); + float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); + int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); + int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); + int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); + float* p_sinks = static_cast(arg_sinks.data_ptr()); + float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); + float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); + int32_t v_num_heads = (int32_t)arg_num_heads; + int32_t v_num_query_tokens = (int32_t)arg_num_query_tokens; + int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; + int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; + int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; + int32_t v_sparse_topk = (int32_t)arg_sparse_topk; + int32_t v_has_sinks = (int32_t)arg_has_sinks; + int32_t v_total_work_items = (int32_t)arg_total_work_items; + void* kargs[] = {&p_tmap_q, &p_tmap_swa_k, &p_tmap_compressed_k, &p_tmap_swa_v, &p_tmap_compressed_v, &p_O, &p_partial_lse, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_num_query_tokens, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_has_sinks, &v_total_work_items}; + + static std::mutex smem_mu; + static auto* smem_status_by_device = + new std::unordered_map(); + cudaError_t smem_status = cudaSuccess; + { + std::lock_guard lock(smem_mu); + auto it = smem_status_by_device->find(dev.device_id); + if (it == smem_status_by_device->end()) { + int current_device = -1; + smem_status = cudaGetDevice(¤t_device); + if (smem_status == cudaSuccess && current_device != dev.device_id) { + smem_status = cudaErrorInvalidDevice; + } + if (smem_status == cudaSuccess) { + smem_status = cudaFuncSetAttribute( + reinterpret_cast(kernel_cake_dsv4_42fa8e98da5fa46d1a94), + cudaFuncAttributeMaxDynamicSharedMemorySize, + 220032); + } + smem_status_by_device->emplace(dev.device_id, smem_status); + } else { + smem_status = it->second; + } + } + TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) + << "cudaFuncSetAttribute for kernel_cake_dsv4_42fa8e98da5fa46d1a94 failed: " + << cudaGetErrorString(smem_status); + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(512u, 1u, 1u); + + // Extended direct-source launch for a linked CUDA kernel symbol. + cudaLaunchAttribute attrs[1]{}; + int n = 0; + attrs[n].id = cudaLaunchAttributeClusterDimension; + attrs[n].val.clusterDim.x = 2u; + attrs[n].val.clusterDim.y = 1u; + attrs[n].val.clusterDim.z = 1u; + ++n; + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes = 220032u; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = n; + cudaError_t launch_status = cudaLaunchKernelExC( + &config, reinterpret_cast(kernel_cake_dsv4_42fa8e98da5fa46d1a94), kargs); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernelExC for kernel_cake_dsv4_42fa8e98da5fa46d1a94 failed: " + << cudaGetErrorString(launch_status); +} +} // namespace variant_cake_dsv4_42fa8e98da5fa46d1a94_sm_100a_6276393afd19 + +namespace variant_cake_dsv4_d50dd820272ad24f7027_sm_100a_aa3105f3b9de { + +void Run(TensorView arg_partial_O, TensorView arg_partial_lse, TensorView arg_O, int64_t arg_num_q_heads, int64_t arg_num_split, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { + DLDevice dev = arg_partial_O.device(); + ScopedCudaDevice device_guard(dev.device_id); + CheckCudaTensor(arg_partial_O, "partial_O"); + CheckDtype(arg_partial_O, "partial_O", 4, 16, 1); + CheckContiguous(arg_partial_O, "partial_O"); + CheckCudaTensor(arg_partial_lse, "partial_lse"); + CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); + CheckContiguous(arg_partial_lse, "partial_lse"); + CheckCudaTensor(arg_O, "O"); + CheckDtype(arg_O, "O", 4, 16, 1); + CheckContiguous(arg_O, "O"); + TVM_FFI_CHECK(arg_num_q_heads >= -2147483648LL && arg_num_q_heads <= 2147483647LL, ValueError) + << "scalar 'num_q_heads' value " << arg_num_q_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_split >= -2147483648LL && arg_num_split <= 2147483647LL, ValueError) + << "scalar 'num_split' value " << arg_num_split + << " is outside i32 range [-2147483648, 2147483647]"; + CheckSameCudaDevice(arg_partial_lse, arg_partial_O, "partial_lse", "partial_O"); + CheckSameCudaDevice(arg_O, arg_partial_O, "O", "partial_O"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y + << ", " << grid_z << ")"; + + + __nv_bfloat16* p_partial_O = static_cast<__nv_bfloat16*>(arg_partial_O.data_ptr()); + float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); + __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); + int32_t v_num_q_heads = (int32_t)arg_num_q_heads; + int32_t v_num_split = (int32_t)arg_num_split; + void* kargs[] = {&p_partial_O, &p_partial_lse, &p_O, &v_num_q_heads, &v_num_split}; + + + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(128u, 1u, 1u); + + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes = 0u; + config.stream = stream; + cudaError_t launch_status = cudaLaunchKernelExC( + &config, reinterpret_cast(kernel_cake_dsv4_d50dd820272ad24f7027), kargs); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernelExC for kernel_cake_dsv4_d50dd820272ad24f7027 failed: " + << cudaGetErrorString(launch_status); +} +} // namespace variant_cake_dsv4_d50dd820272ad24f7027_sm_100a_aa3105f3b9de + +int SelectVariant(int32_t num_heads, int32_t num_query_tokens, int32_t swa_index_stride, int32_t compressed_index_stride, int32_t sparse_topk_lens_offset, int32_t sparse_topk, int32_t has_sinks, int32_t total_work_items, int32_t num_split) { + if ((num_split == 4)) return 0; + return -1; +} + +void Dispatch(tvm::ffi::TensorView arg_Q, tvm::ffi::TensorView arg_SWA_cache, tvm::ffi::TensorView arg_compressed_KV_cache, tvm::ffi::TensorView arg_partial_O, tvm::ffi::TensorView arg_partial_lse, tvm::ffi::TensorView arg_swa_indices, tvm::ffi::TensorView arg_compressed_indices, tvm::ffi::TensorView arg_sparse_topk_lens, tvm::ffi::TensorView arg_sinks, tvm::ffi::TensorView arg_bmm1_scale, tvm::ffi::TensorView arg_bmm2_scale, tvm::ffi::TensorView arg_O, int32_t num_heads, int32_t num_query_tokens, int32_t swa_index_stride, int32_t compressed_index_stride, int32_t sparse_topk_lens_offset, int32_t sparse_topk, int32_t has_sinks, int32_t total_work_items, int32_t num_split, int64_t cuda_stream_ptr) { + CheckCudaTensor(arg_Q, "Q"); + CheckCudaTensor(arg_SWA_cache, "SWA_cache"); + CheckCudaTensor(arg_compressed_KV_cache, "compressed_KV_cache"); + CheckCudaTensor(arg_partial_O, "partial_O"); + CheckCudaTensor(arg_partial_lse, "partial_lse"); + CheckCudaTensor(arg_swa_indices, "swa_indices"); + CheckCudaTensor(arg_compressed_indices, "compressed_indices"); + CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); + CheckCudaTensor(arg_sinks, "sinks"); + CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); + CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); + CheckCudaTensor(arg_O, "O"); + cudaStream_t stream = reinterpret_cast(static_cast(cuda_stream_ptr)); + int v = SelectVariant(num_heads, num_query_tokens, swa_index_stride, compressed_index_stride, sparse_topk_lens_offset, sparse_topk, has_sinks, total_work_items, num_split); + switch (v) { + case 0: { + { + int64_t gx = (total_work_items * 2); + int64_t gy = 1; + int64_t gz = 1; + cake_host_shim_6534755d2bf16a2e::variant_cake_dsv4_42fa8e98da5fa46d1a94_sm_100a_6276393afd19::Run(arg_Q, arg_SWA_cache, arg_compressed_KV_cache, arg_SWA_cache, arg_compressed_KV_cache, arg_partial_O, arg_partial_lse, arg_swa_indices, arg_compressed_indices, arg_sparse_topk_lens, arg_sinks, arg_bmm1_scale, arg_bmm2_scale, num_heads, num_query_tokens, swa_index_stride, compressed_index_stride, sparse_topk_lens_offset, sparse_topk, has_sinks, total_work_items, gx, gy, gz, stream); + } + { + int64_t gx = num_query_tokens; + int64_t gy = num_heads; + int64_t gz = 1; + cake_host_shim_6534755d2bf16a2e::variant_cake_dsv4_d50dd820272ad24f7027_sm_100a_aa3105f3b9de::Run(arg_partial_O, arg_partial_lse, arg_O, num_heads, num_split, gx, gy, gz, stream); + } + return; + } + default: + TVM_FFI_CHECK(false, ValueError) << "no dispatch route for cake_dsv4_family_9c96528827a52c42bbd7 shape (num_heads=" << num_heads << ", num_query_tokens=" << num_query_tokens << ", swa_index_stride=" << swa_index_stride << ", compressed_index_stride=" << compressed_index_stride << ", sparse_topk_lens_offset=" << sparse_topk_lens_offset << ", sparse_topk=" << sparse_topk << ", has_sinks=" << has_sinks << ", total_work_items=" << total_work_items << ", num_split=" << num_split << ")"; + } +} + +int64_t RouteOf(int32_t num_heads, int32_t num_query_tokens, int32_t swa_index_stride, int32_t compressed_index_stride, int32_t sparse_topk_lens_offset, int32_t sparse_topk, int32_t has_sinks, int32_t total_work_items, int32_t num_split) { return SelectVariant(num_heads, num_query_tokens, swa_index_stride, compressed_index_stride, sparse_topk_lens_offset, sparse_topk, has_sinks, total_work_items, num_split); } + +} // namespace cake_host_shim_6534755d2bf16a2e + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_6534755d2bf16a2e::Dispatch); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(route_of, cake_host_shim_6534755d2bf16a2e::RouteOf); diff --git a/flashinfer/jit/cake_dsv4.py b/flashinfer/jit/cake_dsv4.py index 6c1525f8151..10b77c359f2 100644 --- a/flashinfer/jit/cake_dsv4.py +++ b/flashinfer/jit/cake_dsv4.py @@ -85,7 +85,7 @@ }, "sources": ["cake_dsv4_family_06111202fc1bd2f72060_binding.cu"], }, - "cake_dsv4_family_068c73b21d7cbcd958b8": { + "cake_dsv4_family_5383a615a7e502f6fdac": { "arch": "sm_100a", "compile_flags": ["-std=c++17"], "dependencies": [ @@ -94,7 +94,7 @@ ], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "567dd4bd885f3cb2e808cfe4aa2012ef8bbc25e3698b5a7febf987f09df63eb9", + "identity": "78bbedd13dc7248f5985a1e029ffb2bd8e7be9402c95fda93a698cb73e806ace", "route_entry": "route_of", "routes": [ { @@ -159,15 +159,15 @@ "workspace_dtypes": {}, "workspace_keys": [], }, - "sources": ["cake_dsv4_family_068c73b21d7cbcd958b8_binding.cu"], + "sources": ["cake_dsv4_family_5383a615a7e502f6fdac_binding.cu"], }, - "cake_dsv4_family_a84dbf9944ef3cfb088b": { + "cake_dsv4_family_9c96528827a52c42bbd7": { "arch": "sm_100a", "compile_flags": ["-std=c++17"], "dependencies": ["bf16_h128_topk128x_split4_sm100", "split_reduce"], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "7a1c48125738408e45898acfc24926b83b872852cfb31556b4a0dc5cf6b30b8c", + "identity": "048575f321b03f7bd49bbb1f639d8f3790b65d4044648b653e005f9ba51de109", "route_entry": "route_of", "routes": [ { @@ -232,7 +232,7 @@ "workspace_dtypes": {}, "workspace_keys": [], }, - "sources": ["cake_dsv4_family_a84dbf9944ef3cfb088b_binding.cu"], + "sources": ["cake_dsv4_family_9c96528827a52c42bbd7_binding.cu"], }, "cake_dsv4_family_c6fc162d8e0b1752b4f7": { "arch": "sm_100a", @@ -391,7 +391,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "b3715376b8c6e27f09d941cd1d259130d7a54170640d87e0b534d83708ead0fd", + "identity": "bddbb08633a8c728dfaafad7d7d2bebf224379cf83d5b0ba3637eb6520917fb7", "sources": [ "cake_dsv4_314d943a39ab7eade2c8_kernel.cu", "cake_dsv4_314d943a39ab7eade2c8_binding.cu", @@ -429,7 +429,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "24344094f11b37ec4b591eed3a095f38440fc3d66ab7d03357d47093633b9cbb", + "identity": "5f0ec3e7b912f6be4f30e9622a5f2ac879da9c5f09963122bcb02004f32e0f3e", "sources": [ "cake_dsv4_d994492f2dfb135a7dbf_kernel.cu", "cake_dsv4_d994492f2dfb135a7dbf_binding.cu", @@ -451,7 +451,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "1b1cc19373e52f67981ef2cff298918c5579ee69afa319f8c1ed8c1c5cf4ee7c", + "identity": "056b13094b2c64be626f377c14517e7bf67014b79476caae010d9e927d8ce4c7", "sources": [ "cake_dsv4_d906a03bd66a7488fc29_kernel.cu", "cake_dsv4_d906a03bd66a7488fc29_binding.cu", @@ -484,10 +484,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "82976f52e6a753f3e2fe9a76552492077e07345543038a013d52d92ae7d037a2", + "identity": "4c903082c53d323a454998a3b5305b7bcf38cc671c2661ab6426801494b3731d", "sources": [ - "cake_dsv4_57e7b93def70f2be356f_kernel.cu", - "cake_dsv4_57e7b93def70f2be356f_binding.cu", + "cake_dsv4_0c8ab955ac7b62e7bfee_kernel.cu", + "cake_dsv4_0c8ab955ac7b62e7bfee_binding.cu", ], "tma_workspace_bytes": 0, }, @@ -522,7 +522,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "f2200c038a3de20ae07fe52f1a3db30d71941feff6d47fe4db8cd1a45e9d9d44", + "identity": "1bbf6b3a268b2084c0832b70c721618ad84d91c6feb75adba665af2238a7da5a", "sources": [ "cake_dsv4_e6d521df78ecb146fd49_kernel.cu", "cake_dsv4_e6d521df78ecb146fd49_binding.cu", @@ -537,7 +537,7 @@ ["tma_buffer", "tmap_compressed_k"], ["tma_buffer", "tmap_swa_v"], ["tma_buffer", "tmap_compressed_v"], - ["buffer", "O"], + ["tma_buffer", "O"], ["buffer", "partial_lse"], ["buffer", "swa_indices"], ["buffer", "compressed_indices"], @@ -560,10 +560,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "5609e5563448b459fc62af27735acd289478d73d4dfa97cc276009a78b472705", + "identity": "175b037118c3f8d346bc89148513a38f643f541121ed3485f18b9d8c58694b44", "sources": [ - "cake_dsv4_e1660a7f44d156286fea_kernel.cu", - "cake_dsv4_e1660a7f44d156286fea_binding.cu", + "cake_dsv4_42fa8e98da5fa46d1a94_kernel.cu", + "cake_dsv4_42fa8e98da5fa46d1a94_binding.cu", ], "tma_workspace_bytes": 0, }, @@ -575,7 +575,7 @@ ["tma_buffer", "tmap_compressed_k"], ["tma_buffer", "tmap_swa_v"], ["tma_buffer", "tmap_compressed_v"], - ["buffer", "O"], + ["tma_buffer", "O"], ["buffer", "partial_lse"], ["buffer", "swa_indices"], ["buffer", "compressed_indices"], @@ -598,10 +598,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "11243ea83b1c9138a3751dda4de6de1d9ba24712eddce8a30c46fa7027f09e4f", + "identity": "6af64a72ba546c5afa01185851bbd377cb0ee1b41e85a30292124934a9518da3", "sources": [ - "cake_dsv4_ba9ed6a6233761682cc3_kernel.cu", - "cake_dsv4_ba9ed6a6233761682cc3_binding.cu", + "cake_dsv4_714df747e199d0adbe08_kernel.cu", + "cake_dsv4_714df747e199d0adbe08_binding.cu", ], "tma_workspace_bytes": 0, }, @@ -630,7 +630,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "dce42787024ee775d92b9917c6aef6c75769aa31d20cc420a5986533fae9dc29", + "identity": "d68dac00e6e7cf66a1d5ada72032d13c8f4d11349ccbba93b4dc9d69bcbedb6c", "sources": [ "cake_dsv4_a96a3de25becc6c0a422_kernel.cu", "cake_dsv4_a96a3de25becc6c0a422_binding.cu", @@ -668,7 +668,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "7a08a9d6d26f32e66a62402142d1fdb47a1358542fe79be3c76445e4826c6824", + "identity": "ac092d7073e4b4a91c2d48a58a124dfb9e682342c050bac9688017681469e0ee", "sources": [ "cake_dsv4_c0e3d0af9a266850de48_kernel.cu", "cake_dsv4_c0e3d0af9a266850de48_binding.cu", @@ -703,7 +703,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "fe3d4502bb2dfe17a4a4d94fe8f02215e9ee7ae48fd4c70ef3e29326fe7e3640", + "identity": "a910d48878523c7afc2a69bda42a34e7b36b2731d2c4a633592bad70e22551fd", "sources": [ "cake_dsv4_422eafcc723b03b8e044_kernel.cu", "cake_dsv4_422eafcc723b03b8e044_binding.cu", @@ -725,7 +725,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "2e14912ba8d257e634816d04ff5cce6abc380b94e8b771da784bb8e3d97e2b53", + "identity": "1c7462a3ce0fef598330a20566a419bf922d70f655b6493501807217ffac4232", "sources": [ "cake_dsv4_697c8662eb01c8d6bd50_kernel.cu", "cake_dsv4_697c8662eb01c8d6bd50_binding.cu", @@ -763,10 +763,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "e1b220a2d15a79b981a5e717f4d92d99f416852bcea16cb3c6d377e39c9d9671", + "identity": "427e28dadac3a6dfe75d9810c7687601813da18412b5afe0fc94a074b9f09ca9", "sources": [ - "cake_dsv4_3b3789d6356f15f9cbe3_kernel.cu", - "cake_dsv4_3b3789d6356f15f9cbe3_binding.cu", + "cake_dsv4_09e14f1d9cbd6739db61_kernel.cu", + "cake_dsv4_09e14f1d9cbd6739db61_binding.cu", ], "tma_workspace_bytes": 0, }, @@ -797,10 +797,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "53b1e34a7eaa5c285ff343627132fc0cf9a943a2f7860e2bc10cb45797d989f6", + "identity": "cf8478ea922b580842f74f3ea666062aa418264360e9d65577a18fdd4f8bdedc", "sources": [ - "cake_dsv4_9526f80c89f2b5ef4c6c_kernel.cu", - "cake_dsv4_9526f80c89f2b5ef4c6c_binding.cu", + "cake_dsv4_15eabf5d2a5f50db201b_kernel.cu", + "cake_dsv4_15eabf5d2a5f50db201b_binding.cu", ], "tma_workspace_bytes": 0, }, @@ -833,7 +833,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "14152b76d08f7df3fc1e57ac1a9c48ee800539333b92cd59c65ec148f1025558", + "identity": "f923f90544f3007aa8e2db40fa8998257b8dc655c2fe04fe04b4d86a2bcc7853", "sources": [ "cake_dsv4_bf2bffa3152d5dce396f_kernel.cu", "cake_dsv4_bf2bffa3152d5dce396f_binding.cu", @@ -865,7 +865,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "d9b38ca3f8c4bf622b0e37322c11c66c33889b7444792e64a4fbb398f0671e71", + "identity": "cf95eea83e8bdabd58514f45be73fa45bb47398dc513f838d56a6dd5baf873c4", "sources": [ "cake_dsv4_88d5662b6147fe74ed0c_kernel.cu", "cake_dsv4_88d5662b6147fe74ed0c_binding.cu", @@ -897,10 +897,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "db3a18b88f9a22de1bbaf8e18f45b1281ce7ed5f69a7a102ea55af5e8ca5cf24", + "identity": "b6d2b855936ced61ba3dfa126235e6c1fde09b383afa90c887e6c26b7f036592", "sources": [ - "cake_dsv4_4ac65f00a4a032c10ecb_kernel.cu", - "cake_dsv4_4ac65f00a4a032c10ecb_binding.cu", + "cake_dsv4_12ac507a3c73f2784621_kernel.cu", + "cake_dsv4_12ac507a3c73f2784621_binding.cu", ], "tma_workspace_bytes": 0, }, @@ -941,7 +941,7 @@ ], "entry": "run", "host_linkage_flags": [], - "identity": "2c813be0fac81518383683a7faa658c9e452fb1f1966745bdbc4306327d3dee1", + "identity": "5001ca9df604adc9b0728459de36a9f4fd2f544d95d5b6363eb8dd740e4ce4c2", "sources": [ "cake_dsv4_2dd31513440dbd1a0ef2_kernel.cu", "cake_dsv4_2dd31513440dbd1a0ef2_binding.cu", @@ -982,7 +982,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "4dce24ad4f95460dfa2ecc045bf899c67c4e7927a1bb362b2f3cf1d5d24ceb26", + "identity": "8968f0f75b48b5e140c407a3610819fcf62395d53fb2d97d25c896f52b60dfd8", "sources": [ "cake_dsv4_77c162685d381298b786_kernel.cu", "cake_dsv4_77c162685d381298b786_binding.cu", @@ -1022,7 +1022,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "ca35d52a3fbfe936fa85d8f8e4c62b2f624c9b1043519640803738047d5c17c3", + "identity": "17011dd54760db735a19a3b3f6ac9312dd9ab0baa8828fd98ce06cd66deaa64b", "sources": [ "cake_dsv4_119728d98fe411f18a22_kernel.cu", "cake_dsv4_119728d98fe411f18a22_binding.cu", @@ -1057,10 +1057,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "1dc5dc743aca6bb61cb8956a6906a9cb03997d40adfc53f5eb5308dc16fd279a", + "identity": "ec7155bae85cda09aa4c0a15c0b804ffbe91284f388c51a7e60971bae78adde3", "sources": [ - "cake_dsv4_5e5d6648b7bd42bcd0d6_kernel.cu", - "cake_dsv4_5e5d6648b7bd42bcd0d6_binding.cu", + "cake_dsv4_bab2852cc0ea4a6d4860_kernel.cu", + "cake_dsv4_bab2852cc0ea4a6d4860_binding.cu", ], "tma_workspace_bytes": 0, }, @@ -1093,10 +1093,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "7777f8fda7eda670a2733c023b36fd5c05a0cef563e051a8a715ed22d5cc8b79", + "identity": "b4d0530f81c68e3ce8c5d71100116cd46594116ec0c3baa5a71435b99ef03033", "sources": [ - "cake_dsv4_fbf138c989dfde9f764f_kernel.cu", - "cake_dsv4_fbf138c989dfde9f764f_binding.cu", + "cake_dsv4_21a93d40caefcf5b252d_kernel.cu", + "cake_dsv4_21a93d40caefcf5b252d_binding.cu", ], "tma_workspace_bytes": 0, }, @@ -1129,10 +1129,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "3640ee1967a674a977c4ce19a91cfc1c0095d43d59c6ea1fea8cd6c9157331cc", + "identity": "b49f25ffcfc867c0feb10b2184dc867b5e9538760d711e94b08ac337f028099f", "sources": [ - "cake_dsv4_7771bcc5c13a5ff83890_kernel.cu", - "cake_dsv4_7771bcc5c13a5ff83890_binding.cu", + "cake_dsv4_5703afd5b398f2b6c980_kernel.cu", + "cake_dsv4_5703afd5b398f2b6c980_binding.cu", ], "tma_workspace_bytes": 0, }, @@ -1168,10 +1168,10 @@ ], "entry": "run", "host_linkage_flags": [], - "identity": "ef787d9bc5f14ef9eab0e77e05e3e8a3a350341db6b57f23816515a16a92b7d0", + "identity": "c9c07f67f19855f4b65189ccf12c9f9154c53c4fe8b771f198aa2b6ebe9b3948", "sources": [ - "cake_dsv4_fb6409904b35ed0f7a1a_kernel.cu", - "cake_dsv4_fb6409904b35ed0f7a1a_binding.cu", + "cake_dsv4_9c3b9b4fa7cba6e44ffc_kernel.cu", + "cake_dsv4_9c3b9b4fa7cba6e44ffc_binding.cu", ], "tma_workspace_bytes": 0, }, @@ -1204,10 +1204,10 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": [], - "identity": "2a1fc4a24d65057c1db942da73a0d62ca1b02b8a1a1002b05532990638dd2aca", + "identity": "7f336bc7088175237e4e51cbd4b8e212eb7e9024aaa080e7d655d611bbd6679e", "sources": [ - "cake_dsv4_75827bd38b443de66f67_kernel.cu", - "cake_dsv4_75827bd38b443de66f67_binding.cu", + "cake_dsv4_530fe7e6cb8c7d09744d_kernel.cu", + "cake_dsv4_530fe7e6cb8c7d09744d_binding.cu", ], "tma_workspace_bytes": 0, }, @@ -1226,7 +1226,7 @@ "compile_flags": ["--use_fast_math"], "entry": "run", "host_linkage_flags": ["--device-entity-has-hidden-visibility=false"], - "identity": "014512652fd83c055323b34e9c4807829313cd35fbcde3a1d44cd87e89404265", + "identity": "f3b805d486f48af9a2e1e8bc5d443a18f1509b90ce0326a93cb92e64a8fa2589", "sources": [ "cake_dsv4_d50dd820272ad24f7027_kernel.cu", "cake_dsv4_d50dd820272ad24f7027_binding.cu", From 5ee5417bfbea36f160fc2df24f0ee56e9d71818e Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Wed, 30 Sep 2026 00:31:20 -0700 Subject: [PATCH 3/3] perf(cake_sparse_mla): drop the round-6 program sources superseded by the round-7 export Co-Authored-By: Claude Fable 5.1 --- .../cake_dsv4_3b3789d6356f15f9cbe3_binding.cu | 454 --- .../cake_dsv4_3b3789d6356f15f9cbe3_kernel.cu | 1678 ---------- .../cake_dsv4_4ac65f00a4a032c10ecb_binding.cu | 385 --- .../cake_dsv4_4ac65f00a4a032c10ecb_kernel.cu | 1635 --------- .../cake_dsv4_57e7b93def70f2be356f_binding.cu | 389 --- .../cake_dsv4_57e7b93def70f2be356f_kernel.cu | 1643 --------- .../cake_dsv4_5e5d6648b7bd42bcd0d6_binding.cu | 441 --- .../cake_dsv4_5e5d6648b7bd42bcd0d6_kernel.cu | 2973 ----------------- .../cake_dsv4_75827bd38b443de66f67_binding.cu | 445 --- .../cake_dsv4_75827bd38b443de66f67_kernel.cu | 2392 ------------- .../cake_dsv4_7771bcc5c13a5ff83890_binding.cu | 461 --- .../cake_dsv4_7771bcc5c13a5ff83890_kernel.cu | 2803 ---------------- .../cake_dsv4_9526f80c89f2b5ef4c6c_binding.cu | 436 --- .../cake_dsv4_9526f80c89f2b5ef4c6c_kernel.cu | 2542 -------------- .../cake_dsv4_ba9ed6a6233761682cc3_binding.cu | 555 --- .../cake_dsv4_ba9ed6a6233761682cc3_kernel.cu | 2367 ------------- .../cake_dsv4_e1660a7f44d156286fea_binding.cu | 555 --- .../cake_dsv4_e1660a7f44d156286fea_kernel.cu | 2353 ------------- ...sv4_family_068c73b21d7cbcd958b8_binding.cu | 652 ---- ...sv4_family_a84dbf9944ef3cfb088b_binding.cu | 652 ---- .../cake_dsv4_fb6409904b35ed0f7a1a_binding.cu | 461 --- .../cake_dsv4_fb6409904b35ed0f7a1a_kernel.cu | 2812 ---------------- .../cake_dsv4_fbf138c989dfde9f764f_binding.cu | 446 --- .../cake_dsv4_fbf138c989dfde9f764f_kernel.cu | 2616 --------------- .../cake_dsv4_125332c2066d89bd7610_binding.cu | 555 --- .../cake_dsv4_125332c2066d89bd7610_kernel.cu | 2407 ------------- .../cake_dsv4_3e5a644d61420c9c7ea1_binding.cu | 516 --- .../cake_dsv4_3e5a644d61420c9c7ea1_kernel.cu | 2550 -------------- .../cake_dsv4_4c41fb8347f137d4d46f_binding.cu | 555 --- .../cake_dsv4_4c41fb8347f137d4d46f_kernel.cu | 2352 ------------- .../cake_dsv4_7a4b9556b3c84ebff6ce_binding.cu | 526 --- .../cake_dsv4_7a4b9556b3c84ebff6ce_kernel.cu | 2625 --------------- .../cake_dsv4_9cf30aacb43b71034660_binding.cu | 385 --- .../cake_dsv4_9cf30aacb43b71034660_kernel.cu | 1634 --------- .../cake_dsv4_9e1188b192c9c806d7c6_binding.cu | 461 --- .../cake_dsv4_9e1188b192c9c806d7c6_kernel.cu | 2802 ---------------- .../cake_dsv4_a2be51a1fd34ec0d9a3b_binding.cu | 525 --- .../cake_dsv4_a2be51a1fd34ec0d9a3b_kernel.cu | 2400 ------------- .../cake_dsv4_c20607c57cc11784015e_binding.cu | 461 --- .../cake_dsv4_c20607c57cc11784015e_kernel.cu | 2821 ---------------- .../cake_dsv4_c967e7efe220160e28d3_binding.cu | 534 --- .../cake_dsv4_c967e7efe220160e28d3_kernel.cu | 1686 ---------- .../cake_dsv4_d23652aae3370c30d978_binding.cu | 526 --- .../cake_dsv4_d23652aae3370c30d978_kernel.cu | 2704 --------------- .../cake_dsv4_f55f69afe4ac3ec96817_binding.cu | 441 --- .../cake_dsv4_f55f69afe4ac3ec96817_kernel.cu | 2972 ---------------- .../cake_dsv4_f89bc6561b38462e9538_binding.cu | 389 --- .../cake_dsv4_f89bc6561b38462e9538_kernel.cu | 1642 --------- ...sv4_family_956aa0f74f4a82042f2c_binding.cu | 652 ---- ...sv4_family_ec44ddd94e3be88cfc83_binding.cu | 652 ---- 50 files changed, 67919 deletions(-) delete mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_3b3789d6356f15f9cbe3_binding.cu delete mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_3b3789d6356f15f9cbe3_kernel.cu delete mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_4ac65f00a4a032c10ecb_binding.cu delete mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_4ac65f00a4a032c10ecb_kernel.cu delete mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_57e7b93def70f2be356f_binding.cu delete mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_57e7b93def70f2be356f_kernel.cu delete mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_5e5d6648b7bd42bcd0d6_binding.cu delete mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_5e5d6648b7bd42bcd0d6_kernel.cu delete mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_75827bd38b443de66f67_binding.cu delete mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_75827bd38b443de66f67_kernel.cu delete mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_7771bcc5c13a5ff83890_binding.cu delete mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_7771bcc5c13a5ff83890_kernel.cu delete mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_9526f80c89f2b5ef4c6c_binding.cu delete mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_9526f80c89f2b5ef4c6c_kernel.cu delete mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_ba9ed6a6233761682cc3_binding.cu delete mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_ba9ed6a6233761682cc3_kernel.cu delete mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_e1660a7f44d156286fea_binding.cu delete mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_e1660a7f44d156286fea_kernel.cu delete mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_family_068c73b21d7cbcd958b8_binding.cu delete mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_family_a84dbf9944ef3cfb088b_binding.cu delete mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_fb6409904b35ed0f7a1a_binding.cu delete mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_fb6409904b35ed0f7a1a_kernel.cu delete mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_fbf138c989dfde9f764f_binding.cu delete mode 100644 csrc/cake_dsv4/sm_100a/cake_dsv4_fbf138c989dfde9f764f_kernel.cu delete mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_125332c2066d89bd7610_binding.cu delete mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_125332c2066d89bd7610_kernel.cu delete mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_3e5a644d61420c9c7ea1_binding.cu delete mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_3e5a644d61420c9c7ea1_kernel.cu delete mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_4c41fb8347f137d4d46f_binding.cu delete mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_4c41fb8347f137d4d46f_kernel.cu delete mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_7a4b9556b3c84ebff6ce_binding.cu delete mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_7a4b9556b3c84ebff6ce_kernel.cu delete mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_9cf30aacb43b71034660_binding.cu delete mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_9cf30aacb43b71034660_kernel.cu delete mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_9e1188b192c9c806d7c6_binding.cu delete mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_9e1188b192c9c806d7c6_kernel.cu delete mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_a2be51a1fd34ec0d9a3b_binding.cu delete mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_a2be51a1fd34ec0d9a3b_kernel.cu delete mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_c20607c57cc11784015e_binding.cu delete mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_c20607c57cc11784015e_kernel.cu delete mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_c967e7efe220160e28d3_binding.cu delete mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_c967e7efe220160e28d3_kernel.cu delete mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_d23652aae3370c30d978_binding.cu delete mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_d23652aae3370c30d978_kernel.cu delete mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_f55f69afe4ac3ec96817_binding.cu delete mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_f55f69afe4ac3ec96817_kernel.cu delete mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_f89bc6561b38462e9538_binding.cu delete mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_f89bc6561b38462e9538_kernel.cu delete mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_family_956aa0f74f4a82042f2c_binding.cu delete mode 100644 csrc/cake_dsv4/sm_103a/cake_dsv4_family_ec44ddd94e3be88cfc83_binding.cu diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_3b3789d6356f15f9cbe3_binding.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_3b3789d6356f15f9cbe3_binding.cu deleted file mode 100644 index 08603a80ed1..00000000000 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_3b3789d6356f15f9cbe3_binding.cu +++ /dev/null @@ -1,454 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -// Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_3b3789d6356f15f9cbe3'. -#include -#include -#include -#include - -#include "tvm_ffi_utils.h" - -#include -#include -#include -#include -#include -#include -#include - -extern "C" __global__ void kernel_cake_dsv4_3b3789d6356f15f9cbe3(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int batch_size, int max_q_len, int ragged_query, int has_sinks); - - -namespace cake_host_shim_d2767e2d94919bb7 { - -using tvm::ffi::Optional; -using tvm::ffi::TensorView; - -class ScopedCudaDevice { - public: - explicit ScopedCudaDevice(int device_id) { - cudaError_t error = cudaGetDevice(&previous_device_); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed before host-shim launch: cudaError=" - << static_cast(error); - if (previous_device_ != device_id) { - error = cudaSetDevice(device_id); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaSetDevice failed before host-shim launch for cuda:" - << device_id << ": cudaError=" << static_cast(error); - restore_ = true; - } - } - - ScopedCudaDevice(const ScopedCudaDevice&) = delete; - ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; - - ~ScopedCudaDevice() noexcept { - if (restore_) { - (void)cudaSetDevice(previous_device_); - } - } - - private: - int previous_device_ = -1; - bool restore_ = false; -}; - -inline int64_t CakeDeviceMultiprocessorCount(int device_id) { - constexpr int kMaxCachedCudaDevices = 64; - TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) - << "physical-SM-count cache does not cover cuda:" << device_id; - static std::atomic count_by_device[kMaxCachedCudaDevices]{}; - int cached = count_by_device[device_id].load(std::memory_order_acquire); - if (cached > 0) return cached; - - int count = 0; - cudaError_t error = cudaDeviceGetAttribute( - &count, cudaDevAttrMultiProcessorCount, device_id); - TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) - << "querying multiProcessorCount failed for cuda:" << device_id - << ": cudaError=" << static_cast(error) << ", count=" << count; - // Concurrent first launches may repeat the immutable device query, but all - // publication is atomic and every later hot-path lookup is one acquire load. - count_by_device[device_id].store(count, std::memory_order_release); - return count; -} - -inline void CheckCudaTensor(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) - << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; -} - -inline void CheckSameCudaDevice( - const TensorView& t, - const TensorView& reference, - const char* name, - const char* reference_name) { - TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) - << name << " must be on the same CUDA device as " << reference_name - << ": got cuda:" << t.device().device_id - << " versus cuda:" << reference.device().device_id; -} - -inline void CheckCurrentCudaDevice( - const TensorView& reference, - const char* reference_name) { - int current_device = -1; - cudaError_t error = cudaGetDevice(¤t_device); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed while validating " << reference_name - << ": cudaError=" << static_cast(error); - TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) - << "current CUDA device must match " << reference_name - << ": current=cuda:" << current_device - << ", tensor=cuda:" << reference.device().device_id; -} - -inline void CheckContiguous(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; -} - -inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { - DLDataType d = t.dtype(); - TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) - << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits - << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits - << ", lanes=" << (int)d.lanes << ")"; -} - -// A logical axis.outer(trailing) folds every source dim above the trailing -// dimensions. Shape products are independent of physical strides, so verify -// the leading dimensions form one dense row-major chain instead of inventing -// a "folded stride". The descriptor reads its exact adjacent physical step -// separately through stride[-(trailing + 1)]. -inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { - TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) - << name << " cannot fold leading dimensions above " << trailing - << " trailing dims from ndim=" << t.ndim(); - int outer_last = t.ndim() - trailing - 1; - if (outer_last <= 0) { - return; - } - int64_t step = t.stride(outer_last); - TVM_FFI_CHECK(step > 0, ValueError) - << name << " physical strides must be positive"; - int64_t expected = step; - for (int axis = outer_last - 1; axis >= 0; --axis) { - expected *= t.size(axis + 1); - if (t.size(axis) > 1) { - TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) - << name << " leading dims are not physically foldable above " << trailing - << " trailing dims: stride(" << axis << ")=" << t.stride(axis) - << ", expected " << expected; - } - } -} - -// 4D TMA descriptor for buffer 'tmap_q' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 3, ValueError) - << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - int64_t d2 = t.size(t.ndim() - 2); - int64_t d3 = t.size(t.ndim() - 3); - TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) - << "TMA source 'tmap_q' trailing dims must be positive"; - TVM_FFI_CHECK(d1 % 64 == 0, ValueError) - << "TMA source 'tmap_q' extent " << d1 - << " must divide exactly by " << 64; - uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(d3)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 8u <= global_dim[2] && 1u <= global_dim[3], ValueError) - << "TMA box (64, 64, 8, 1) exceeds resolved global dims for 'tmap_q'"; - int64_t carrier_stride_0 = d1; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; - int64_t carrier_stride_1 = 64; - TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; - TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; - TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; - int64_t carrier_stride_2 = (d2 * d1); - TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; - TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; - TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; - uint64_t global_strides[3] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - (uint64_t)((carrier_stride_1 * 16) / 8), - (uint64_t)((carrier_stride_2 * 16) / 8), - }; - uint32_t box_dim[4] = {64u, 64u, 8u, 1u}; - uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_swa_kv' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_swa_kv' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_kv'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_compressed_kv' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_compressed_kv(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_compressed_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_compressed_kv' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_compressed_kv' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_compressed_kv"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_compressed_kv' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_kv'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_kv') failed: CUresult=" << (int)r; - return tm; -} - -void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_compressed_kv, TensorView arg_O, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_seq_lens, TensorView arg_cum_seq_lens_q, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_batch_size, int64_t arg_max_q_len, int64_t arg_ragged_query, int64_t arg_has_sinks, int64_t grid_x, int64_t grid_y, int64_t grid_z) { - DLDevice dev = arg_tmap_q.device(); - ScopedCudaDevice device_guard(dev.device_id); - CheckCudaTensor(arg_tmap_q, "tmap_q"); - CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); - CheckContiguous(arg_tmap_q, "tmap_q"); - CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); - CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 4, 16, 1); - CheckCudaTensor(arg_tmap_compressed_kv, "tmap_compressed_kv"); - CheckDtype(arg_tmap_compressed_kv, "tmap_compressed_kv", 4, 16, 1); - CheckCudaTensor(arg_O, "O"); - CheckDtype(arg_O, "O", 4, 16, 1); - CheckContiguous(arg_O, "O"); - CheckCudaTensor(arg_swa_indices, "swa_indices"); - CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); - CheckContiguous(arg_swa_indices, "swa_indices"); - CheckCudaTensor(arg_compressed_indices, "compressed_indices"); - CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); - CheckContiguous(arg_compressed_indices, "compressed_indices"); - CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); - CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckCudaTensor(arg_seq_lens, "seq_lens"); - CheckDtype(arg_seq_lens, "seq_lens", 0, 32, 1); - CheckContiguous(arg_seq_lens, "seq_lens"); - CheckCudaTensor(arg_cum_seq_lens_q, "cum_seq_lens_q"); - CheckDtype(arg_cum_seq_lens_q, "cum_seq_lens_q", 0, 32, 1); - CheckContiguous(arg_cum_seq_lens_q, "cum_seq_lens_q"); - CheckCudaTensor(arg_sinks, "sinks"); - CheckDtype(arg_sinks, "sinks", 2, 32, 1); - CheckContiguous(arg_sinks, "sinks"); - CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); - CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); - CheckContiguous(arg_bmm1_scale, "bmm1_scale"); - CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); - CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); - CheckContiguous(arg_bmm2_scale, "bmm2_scale"); - TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) - << "scalar 'num_heads' value " << arg_num_heads - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) - << "scalar 'swa_index_stride' value " << arg_swa_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) - << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) - << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) - << "scalar 'sparse_topk' value " << arg_sparse_topk - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_batch_size >= -2147483648LL && arg_batch_size <= 2147483647LL, ValueError) - << "scalar 'batch_size' value " << arg_batch_size - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_max_q_len >= -2147483648LL && arg_max_q_len <= 2147483647LL, ValueError) - << "scalar 'max_q_len' value " << arg_max_q_len - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_ragged_query >= -2147483648LL && arg_ragged_query <= 2147483647LL, ValueError) - << "scalar 'ragged_query' value " << arg_ragged_query - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) - << "scalar 'has_sinks' value " << arg_has_sinks - << " is outside i32 range [-2147483648, 2147483647]"; - CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); - CheckSameCudaDevice(arg_tmap_compressed_kv, arg_tmap_q, "tmap_compressed_kv", "tmap_q"); - CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); - CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); - CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); - CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); - CheckSameCudaDevice(arg_seq_lens, arg_tmap_q, "seq_lens", "tmap_q"); - CheckSameCudaDevice(arg_cum_seq_lens_q, arg_tmap_q, "cum_seq_lens_q", "tmap_q"); - CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); - CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); - CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); - TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) - << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y - << ", " << grid_z << ")"; - - cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); - CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); - CUtensorMap p_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); - CUtensorMap p_tmap_compressed_kv = EncodeTma_tmap_compressed_kv(arg_tmap_compressed_kv); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); - int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); - int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); - int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); - int* p_seq_lens = static_cast(arg_seq_lens.data_ptr()); - int* p_cum_seq_lens_q = static_cast(arg_cum_seq_lens_q.data_ptr()); - float* p_sinks = static_cast(arg_sinks.data_ptr()); - float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); - float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); - int32_t v_num_heads = (int32_t)arg_num_heads; - int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; - int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; - int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; - int32_t v_sparse_topk = (int32_t)arg_sparse_topk; - int32_t v_batch_size = (int32_t)arg_batch_size; - int32_t v_max_q_len = (int32_t)arg_max_q_len; - int32_t v_ragged_query = (int32_t)arg_ragged_query; - int32_t v_has_sinks = (int32_t)arg_has_sinks; - void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_tmap_compressed_kv, &p_O, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_seq_lens, &p_cum_seq_lens_q, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_batch_size, &v_max_q_len, &v_ragged_query, &v_has_sinks}; - - static std::mutex smem_mu; - static auto* smem_status_by_device = - new std::unordered_map(); - cudaError_t smem_status = cudaSuccess; - { - std::lock_guard lock(smem_mu); - auto it = smem_status_by_device->find(dev.device_id); - if (it == smem_status_by_device->end()) { - int current_device = -1; - smem_status = cudaGetDevice(¤t_device); - if (smem_status == cudaSuccess && current_device != dev.device_id) { - smem_status = cudaErrorInvalidDevice; - } - if (smem_status == cudaSuccess) { - smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_3b3789d6356f15f9cbe3), - cudaFuncAttributeMaxDynamicSharedMemorySize, - 198144); - } - smem_status_by_device->emplace(dev.device_id, smem_status); - } else { - smem_status = it->second; - } - } - TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_3b3789d6356f15f9cbe3 failed: " - << cudaGetErrorString(smem_status); - dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); - dim3 block(256u, 1u, 1u); - - cudaLaunchConfig_t config{}; - config.gridDim = grid; - config.blockDim = block; - config.dynamicSmemBytes = 198144u; - config.stream = stream; - cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_3b3789d6356f15f9cbe3), kargs); - TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_3b3789d6356f15f9cbe3 failed: " - << cudaGetErrorString(launch_status); -} - -} // namespace cake_host_shim_d2767e2d94919bb7 - -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_d2767e2d94919bb7::Run); diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_3b3789d6356f15f9cbe3_kernel.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_3b3789d6356f15f9cbe3_kernel.cu deleted file mode 100644 index 3a0cb8057d1..00000000000 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_3b3789d6356f15f9cbe3_kernel.cu +++ /dev/null @@ -1,1678 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -typedef signed char int8_t; -typedef unsigned char uint8_t; -typedef unsigned short uint16_t; -typedef unsigned int uint32_t; -#if defined(__CUDACC_RTC__) -typedef unsigned long long uint64_t; -#else -typedef unsigned long uint64_t; -#endif -static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); -typedef signed int int32_t; -typedef short int int16_t; -struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; -struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; -static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); -static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); - -#if defined(__CUDACC_RTC__) -typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; -#else -#include -#endif - -static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); -static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); -#include -#include - -#define CAKE_INF CUDART_INF_F -#define TMEM_NCOLS 512 -#define TMEM_TMEM_OFFSET 0 -#define NUM_KV_PIPE_STAGES 4 -#define SMEM_SMEM_Q_OFF 1024 -#define SMEM_SMEM_Q_STAGE_BYTES 8192 -#define SMEM_SMEM_Q_STRIDE 8192 -#define SMEM_SMEM_KV_OFF 66560 -#define SMEM_SMEM_KV_STAGE_BYTES 32768 -#define SMEM_SMEM_KV_STRIDE 32768 -#define SMEM_SMEM_V_OFF 66560 -#define SMEM_SMEM_V_STAGE_BYTES 32768 -#define SMEM_SMEM_V_STRIDE 32768 -#define SMEM_TOTAL 198144 -#define THREADS 256 - -#include - -__device__ __forceinline__ uint32_t elect_sync() { - uint32_t pred = 0; - asm volatile( - "{\n\t" - ".reg .pred %%px;\n\t" - "elect.sync _|%%px, %1;\n\t" - "@%%px mov.s32 %0, 1;\n\t" - "}\n" - : "+r"(pred) - : "r"(0xFFFFFFFF)); - return pred; -} - - -__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { - asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" - :: "r"(mbar_addr), "r"(count) : "memory"); -} - -__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { - asm volatile("mbarrier.init.b64 [%0], %1;" - :: "l"(mbar_addr), "r"(count) : "memory"); -} - - -__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - - -// CTA-local pipelines have short, resident producer/consumer edges. Omitting -// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the -// explicit loop still makes this helper blocking until acquire succeeds. -__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE;\n\t" - "bra.uni LAB_WAIT;\n\t" - "DONE:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Source-faithful relaxed CTA wait used only by a typed protocol that does -// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. -__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, 10000000;\n\t" - "@P1 bra.uni DONE_RELAXED;\n\t" - "bra.uni LAB_WAIT_RELAXED;\n\t" - "DONE_RELAXED:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Exact source ports may request the PTX suspendTimeHint operand explicitly. -// The hint is expressed in nanoseconds and is kept separate from the canonical -// no-hint CTA helper so unrelated schedules retain their existing retry path. -__device__ __forceinline__ void mbarrier_wait_suspend( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_SUSPEND:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_SUSPEND;\n\t" - "bra.uni LAB_WAIT_SUSPEND;\n\t" - "DONE_SUSPEND:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE_CLUSTER;\n\t" - "bra.uni LAB_WAIT_CLUSTER;\n\t" - "DONE_CLUSTER:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - ".reg .u32 WAIT_ADDR;\n\t" - "mov.u32 WAIT_ADDR, %0;\n\t" - "LAB_WAIT_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [WAIT_ADDR], %1, %2;\n\t" - "@P1 bra.uni DONE_HINT;\n\t" - "bra.uni LAB_WAIT_HINT;\n\t" - "DONE_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -// Exact unqualified CTA wait used by source schedules whose PTX intentionally -// omits the acquire qualifier while retaining a typed suspendTimeHint operand. -__device__ __forceinline__ void mbarrier_wait_relaxed_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED_HINT:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra DONE_RELAXED_HINT;\n\t" - "bra LAB_WAIT_RELAXED_HINT;\n\t" - "DONE_RELAXED_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); -} - -__device__ __forceinline__ void mbarrier_wait_cluster_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" - "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" - "DONE_CLUSTER_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_suspend( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait_cluster(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); - } -} - - -__device__ __forceinline__ void tcgen05_mma_f16( - int taddr, uint64_t a_desc, uint64_t b_desc, - uint32_t i_desc, int enable_input_d) { - asm volatile( - "{\n\t" - ".reg .pred p;\n\t" - "setp.ne.b32 p, %4, 0;\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" - "}\n" - :: "r"(taddr), "l"(a_desc), "l"(b_desc), - "r"(i_desc), "r"(enable_input_d) - : "memory"); -} - - -__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { - return (x & 0x3FFFFULL) >> 4ULL; -} - - -union MmaSmemDesc { - uint64_t u64; - uint32_t u32[2]; -}; - -__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { - MmaSmemDesc tmp; - tmp.u64 = smem_desc; - tmp.u32[0] += offset; - smem_desc = tmp.u64; -} - - -__device__ __forceinline__ void elect_commit(int mbar_addr) { - asm volatile( - "{\n\t" - ".reg .pred leader;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" - ".shared::cluster.b64 [%0];\n\t" - "}\n" - :: "r"(mbar_addr)); -} - - -__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { - asm volatile( - "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" - :: "r"(mbar_addr), "r"(bytes) : "memory"); -} - - -__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7," - " %8, %9, %10, %11, %12, %13, %14, %15," - " %16, %17, %18, %19, %20, %21, %22, %23," - " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), - "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), - "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), - "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), - "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), - "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), - "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ float approx_exp2(float x) { - float y; - asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float approx_rcp(float x) { - float y; - asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float max_noftz(float a, float b) { - float c; - asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); - return c; -} - - -__device__ __forceinline__ float warp_reduce_max(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); - return val; -} - - -__device__ __forceinline__ float warp_reduce_sum(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val += __shfl_xor_sync(0xFFFFFFFF, val, offset); - return val; -} - - -__device__ __forceinline__ float row_max_reduce(float2 acc) { - return max_noftz(acc.x, acc.y); -} - - -__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { - #pragma unroll - for (int j = 0; j < 16; j++) { - if (j % 2 == 0) - acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); - else - acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); - } -} - - -__device__ __forceinline__ float2 ex2_emulation_f32x2_value(float2 value) { - const float c0 = 1.0f, c1 = 0.695146143436431884765625f; - const float c2 = 0.227564394474029541015625f, c3 = 0.077119089663028717041015625f; - const float magic = 12582912.0f; - float x0 = max_noftz(value.x, -127.0f), x1 = max_noftz(value.y, -127.0f); - float2 xc2 = make_float2(x0, x1), magic2 = make_float2(magic, magic); - float2 xr2; - asm("add.rm.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xr2) - : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&magic2)); - float2 c3_2 = make_float2(c3, c3), c2_2 = make_float2(c2, c2); - float2 c1_2 = make_float2(c1, c1), c0_2 = make_float2(c0, c0); - float2 xrb2, xfrac2; - asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xrb2) - : "l"(*(unsigned long long*)&xr2), "l"(*(unsigned long long*)&magic2)); - asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xfrac2) - : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&xrb2)); - float2 poly2; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&c3_2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c2_2)); - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c1_2)); - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c0_2)); - int x0r_i, x1r_i, p0_i, p1_i; - asm("mov.b64 {%0, %1}, %2;" : "=r"(x0r_i), "=r"(x1r_i) : "l"(*(unsigned long long*)&xr2)); - asm("mov.b64 {%0, %1}, %2;" : "=r"(p0_i), "=r"(p1_i) : "l"(*(unsigned long long*)&poly2)); - float r0, r1; - asm("mov.b32 %0, %1;" : "=f"(r0) : "r"((x0r_i << 23) + p0_i)); - asm("mov.b32 %0, %1;" : "=f"(r1) : "r"((x1r_i << 23) + p1_i)); - return make_float2(r0, r1); -} - -__device__ __forceinline__ void ex2_emulation_f32x2(float* x0_ptr, float* x1_ptr) { - float2 result = ex2_emulation_f32x2_value(make_float2(*x0_ptr, *x1_ptr)); - *x0_ptr = result.x; *x1_ptr = result.y; -} - -__device__ __forceinline__ void softmax_frag_exp2_cast( - float* sv, uint32_t* pv, int use_emu) -{ - #pragma unroll - for (int j = 0; j < 16; j++) { - if (use_emu && j >= 12) - ex2_emulation_f32x2(&sv[j*2], &sv[j*2+1]); - else { - sv[j*2] = approx_exp2(sv[j*2]); - sv[j*2+1] = approx_exp2(sv[j*2+1]); - } - } - #pragma unroll - for (int j = 0; j < 16; j++) { - __nv_bfloat162 bf = __float22bfloat162_rn({sv[j*2], sv[j*2+1]}); - pv[j] = reinterpret_cast(bf); - } -} - - - -__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { - const float2* sv2 = reinterpret_cast(sv); - #pragma unroll - for (int j = 0; j < 16; j++) { - asm("add.f32x2 %0, %1, %2;" - : "+l"(reinterpret_cast(*acc)) - : "l"(reinterpret_cast(*acc)), - "l"(reinterpret_cast(sv2[j]))); - } -} - - -__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { - asm("mul.rn.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { - asm("mul.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { - asm("add.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { - asm("add.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { - asm("sub.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { - asm("sub.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("add.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { - float2 r; - asm("sub.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("sub.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ void fma_scale_x32( - float* sv, const float2* scale2, const float2* neg_max2) -{ - float2* sv_2 = reinterpret_cast(sv); - #pragma unroll - for (int j = 0; j < 16; j++) - fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); -} - -__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("mul.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - - -// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) - -__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("add.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("add.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("add.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("add.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("add.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("add.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("add.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - - -__device__ __forceinline__ void fence_async_shared() { - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); -} - - -__device__ __forceinline__ uint64_t make_smem_desc(int addr) { - const int SBO = 1024; - return desc_encode(addr) - | (desc_encode(SBO) << 32ULL) - | (1ULL << 46ULL) - | (2ULL << 61ULL); -} - - -__device__ __forceinline__ void tma_2d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3}], [%4];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_4d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.4d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5}], [%6];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr) { - // Canonical .shared::cta form for non-multicast gather4, matching - // trtllm-gen / cuda_ptx and the PTX ISA qualifier order - // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, - // .shared::cluster is reserved for the multicast variant (ctaMask). - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { - asm volatile( - "tcgen05.commit.cta_group::1.mbarrier::arrive::one" - ".shared::cluster.b64 [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { - uint32_t result; - asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" - : "=r"(result) : "r"(val)); - return result; -} - -extern "C" { - -__global__ __launch_bounds__(256, 1) void -kernel_cake_dsv4_3b3789d6356f15f9cbe3(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int batch_size, int max_q_len, int ragged_query, int has_sinks) -{ - const int tid = threadIdx.x; - const int warp = make_warp_uniform(tid / 32); - const int lane = tid % 32; - - extern __shared__ __align__(1024) char smem_raw[]; - int smem; - asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); - smem = make_warp_uniform(smem); - - const int mbar_base = smem; - #define q_full_addr (mbar_base + 0) - #define kv_full_addr (mbar_base + 8) - #define s_full_addr (mbar_base + 40) - #define p_full_addr (mbar_base + 56) - #define pv_done_addr (mbar_base + 72) - #define tmem_dealloc_addr (mbar_base + 80) - - const int bid = blockIdx.x; - const int num_bids = gridDim.x; - - const int cta_rank = 0; - - // Kernel setup ops - __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); - const int smem_q_addr = smem + 1024; - __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); - const int smem_kv_addr = smem + 66560; - __nv_bfloat16* smem_v = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); - const int smem_v_addr = smem + 66560; - - // Mbarrier init (6 pipeline groups, 0 ordered-sequence groups, 11 barriers) - // Mbarriers at smem_raw[0..88) - - if (warp == 0) { - uint32_t leader = elect_sync(); - if (leader) { - // q_full: 1 barriers, init_count=1 - mbarrier_init(smem + 0, 1); - // --- pipeline 'kv_pipe' --- - // kv_full: 4 barriers, init_count=3 - mbarrier_init(smem + 8, 3); - mbarrier_init(smem + 16, 3); - mbarrier_init(smem + 24, 3); - mbarrier_init(smem + 32, 3); - // s_full: 2 barriers, init_count=1 - mbarrier_init(smem + 40, 1); - mbarrier_init(smem + 48, 1); - // p_full: 2 barriers, init_count=4 - mbarrier_init(smem + 56, 4); - mbarrier_init(smem + 64, 4); - // pv_done: 1 barriers, init_count=1 - mbarrier_init(smem + 72, 1); - // tmem_dealloc: 1 barriers, init_count=7 - mbarrier_init(smem + 80, 7); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - - __syncwarp(); - - // TMEM alloc (512 columns, 512 used) - volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 88); - if (warp == 0) { - int _tmem_hold = smem + 88; - asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); - __syncwarp(); - asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); - } - - __syncthreads(); - asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); - - const int taddr = tmem_addr_storage[0]; - - // Kernel post-init ops - const int tmem_tmem = taddr; - - // ---- Ordered hardware-WG register redistribution ---- - // Dec phase frees registers before any WG attempts inc. - if (warp >= 4 && warp <= 7) { - asm volatile("setmaxnreg.dec.sync.aligned.u32 64;"); - } - - // ---- Role: softmax ---- - if (warp <= 3) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 176;"); - { // softmax_main - float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; - int query_idx = blockIdx.x; - int _max_2 = ((sparse_topk_lens[query_idx] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx] + sparse_topk_lens_offset) : (0)); - int _min_2 = ((_max_2) < (sparse_topk) ? (_max_2) : (sparse_topk)); - int active_topk = _min_2; - int tile_count = (active_topk + 128 - 1) / 128; - int query_batch = query_idx / max_q_len; - int query_offset = query_idx - query_batch * max_q_len; - int query_length = max_q_len; - if (ragged_query != 0) { - query_batch = 0; - #pragma unroll 1 - for (int batch = 0; batch < batch_size; batch++) { - if (query_idx >= cum_seq_lens_q[batch + 1]) { - query_batch = batch + 1; - } - } - int query_begin = cum_seq_lens_q[query_batch]; - query_length = cum_seq_lens_q[query_batch + 1] - query_begin; - query_offset = query_idx - query_begin; - } - int swa_valid_columns = seq_lens[query_batch] - query_length + query_offset + 1; - if (swa_valid_columns < 0) { - swa_valid_columns = 0; - } - if (swa_valid_columns > 128) { - swa_valid_columns = 128; - } - const int warp_in_compute = warp; - const int tmem_row_origin = warp_in_compute * 32; - const int logical_row_origin = warp_in_compute * 16; - const int my_row = logical_row_origin + lane % 16; - const int col_half = lane / 16; - float row_max = -CAKE_INF; - float row_sum = 0.0f; - if (has_sinks != 0 && my_row < num_heads) { - row_max = sinks[my_row] * 1.4426950408889634f / softmax_scale_log2; - row_sum = 1.0f; - } - #pragma unroll - for (int tile = 0; tile < 1; tile++) { - if (tile_count > tile) { - int phase = tile & 1; - int wait_phase = tile >> 1 & 1; - mbarrier_wait(s_full_addr + (phase) * 8, wait_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - int score_col = ((phase != 0) ? 128 : 0); - int score_addr = taddr + (unsigned int)score_col + (unsigned int)(tmem_row_origin << 16); - float _tmem_load_0[64]; - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_0[0])), "=r"(*reinterpret_cast(&_tmem_load_0[1])), "=r"(*reinterpret_cast(&_tmem_load_0[2])), "=r"(*reinterpret_cast(&_tmem_load_0[3])), "=r"(*reinterpret_cast(&_tmem_load_0[4])), "=r"(*reinterpret_cast(&_tmem_load_0[5])), "=r"(*reinterpret_cast(&_tmem_load_0[6])), "=r"(*reinterpret_cast(&_tmem_load_0[7])), "=r"(*reinterpret_cast(&_tmem_load_0[8])), "=r"(*reinterpret_cast(&_tmem_load_0[9])), "=r"(*reinterpret_cast(&_tmem_load_0[10])), "=r"(*reinterpret_cast(&_tmem_load_0[11])), "=r"(*reinterpret_cast(&_tmem_load_0[12])), "=r"(*reinterpret_cast(&_tmem_load_0[13])), "=r"(*reinterpret_cast(&_tmem_load_0[14])), "=r"(*reinterpret_cast(&_tmem_load_0[15])), "=r"(*reinterpret_cast(&_tmem_load_0[16])), "=r"(*reinterpret_cast(&_tmem_load_0[17])), "=r"(*reinterpret_cast(&_tmem_load_0[18])), "=r"(*reinterpret_cast(&_tmem_load_0[19])), "=r"(*reinterpret_cast(&_tmem_load_0[20])), "=r"(*reinterpret_cast(&_tmem_load_0[21])), "=r"(*reinterpret_cast(&_tmem_load_0[22])), "=r"(*reinterpret_cast(&_tmem_load_0[23])), "=r"(*reinterpret_cast(&_tmem_load_0[24])), "=r"(*reinterpret_cast(&_tmem_load_0[25])), "=r"(*reinterpret_cast(&_tmem_load_0[26])), "=r"(*reinterpret_cast(&_tmem_load_0[27])), "=r"(*reinterpret_cast(&_tmem_load_0[28])), "=r"(*reinterpret_cast(&_tmem_load_0[29])), "=r"(*reinterpret_cast(&_tmem_load_0[30])), "=r"(*reinterpret_cast(&_tmem_load_0[31])) - : "r"(score_addr)); - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_0[32])), "=r"(*reinterpret_cast(&_tmem_load_0[33])), "=r"(*reinterpret_cast(&_tmem_load_0[34])), "=r"(*reinterpret_cast(&_tmem_load_0[35])), "=r"(*reinterpret_cast(&_tmem_load_0[36])), "=r"(*reinterpret_cast(&_tmem_load_0[37])), "=r"(*reinterpret_cast(&_tmem_load_0[38])), "=r"(*reinterpret_cast(&_tmem_load_0[39])), "=r"(*reinterpret_cast(&_tmem_load_0[40])), "=r"(*reinterpret_cast(&_tmem_load_0[41])), "=r"(*reinterpret_cast(&_tmem_load_0[42])), "=r"(*reinterpret_cast(&_tmem_load_0[43])), "=r"(*reinterpret_cast(&_tmem_load_0[44])), "=r"(*reinterpret_cast(&_tmem_load_0[45])), "=r"(*reinterpret_cast(&_tmem_load_0[46])), "=r"(*reinterpret_cast(&_tmem_load_0[47])), "=r"(*reinterpret_cast(&_tmem_load_0[48])), "=r"(*reinterpret_cast(&_tmem_load_0[49])), "=r"(*reinterpret_cast(&_tmem_load_0[50])), "=r"(*reinterpret_cast(&_tmem_load_0[51])), "=r"(*reinterpret_cast(&_tmem_load_0[52])), "=r"(*reinterpret_cast(&_tmem_load_0[53])), "=r"(*reinterpret_cast(&_tmem_load_0[54])), "=r"(*reinterpret_cast(&_tmem_load_0[55])), "=r"(*reinterpret_cast(&_tmem_load_0[56])), "=r"(*reinterpret_cast(&_tmem_load_0[57])), "=r"(*reinterpret_cast(&_tmem_load_0[58])), "=r"(*reinterpret_cast(&_tmem_load_0[59])), "=r"(*reinterpret_cast(&_tmem_load_0[60])), "=r"(*reinterpret_cast(&_tmem_load_0[61])), "=r"(*reinterpret_cast(&_tmem_load_0[62])), "=r"(*reinterpret_cast(&_tmem_load_0[63])) - : "r"(score_addr + 32)); - int valid_cols = active_topk - tile * 128 - col_half * 64; - if (valid_cols < 0) { - valid_cols = 0; - } - if (valid_cols > 64) { - valid_cols = 64; - } - if (tile == 0) { - int valid_swa_cols = swa_valid_columns - col_half * 64; - if (valid_swa_cols < 0) { - valid_swa_cols = 0; - } - if (valid_swa_cols < valid_cols) { - valid_cols = valid_swa_cols; - } - } - uint32_t _slice_lo_mask_0; - { - int _lim_0 = valid_cols; - if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } - else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); - } - } - if (!(_slice_lo_mask_0 & (1u << 0))) _tmem_load_0[0] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 1))) _tmem_load_0[1] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 2))) _tmem_load_0[2] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 3))) _tmem_load_0[3] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 4))) _tmem_load_0[4] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 5))) _tmem_load_0[5] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 6))) _tmem_load_0[6] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 7))) _tmem_load_0[7] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 8))) _tmem_load_0[8] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 9))) _tmem_load_0[9] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 10))) _tmem_load_0[10] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 11))) _tmem_load_0[11] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 12))) _tmem_load_0[12] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 13))) _tmem_load_0[13] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 14))) _tmem_load_0[14] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 15))) _tmem_load_0[15] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 16))) _tmem_load_0[16] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 17))) _tmem_load_0[17] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 18))) _tmem_load_0[18] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 19))) _tmem_load_0[19] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 20))) _tmem_load_0[20] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 21))) _tmem_load_0[21] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 22))) _tmem_load_0[22] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 23))) _tmem_load_0[23] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 24))) _tmem_load_0[24] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 25))) _tmem_load_0[25] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 26))) _tmem_load_0[26] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 27))) _tmem_load_0[27] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 28))) _tmem_load_0[28] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 29))) _tmem_load_0[29] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 30))) _tmem_load_0[30] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 31))) _tmem_load_0[31] = -CAKE_INF; - uint32_t _slice_lo_mask_1; - { - int _lim_1 = valid_cols - 32; - if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } - else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); - } - } - if (!(_slice_lo_mask_1 & (1u << 0))) _tmem_load_0[32] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 1))) _tmem_load_0[33] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 2))) _tmem_load_0[34] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 3))) _tmem_load_0[35] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 4))) _tmem_load_0[36] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 5))) _tmem_load_0[37] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 6))) _tmem_load_0[38] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 7))) _tmem_load_0[39] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 8))) _tmem_load_0[40] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 9))) _tmem_load_0[41] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 10))) _tmem_load_0[42] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 11))) _tmem_load_0[43] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 12))) _tmem_load_0[44] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 13))) _tmem_load_0[45] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 14))) _tmem_load_0[46] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 15))) _tmem_load_0[47] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 16))) _tmem_load_0[48] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 17))) _tmem_load_0[49] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 18))) _tmem_load_0[50] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 19))) _tmem_load_0[51] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 20))) _tmem_load_0[52] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 21))) _tmem_load_0[53] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 22))) _tmem_load_0[54] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 23))) _tmem_load_0[55] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 24))) _tmem_load_0[56] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 25))) _tmem_load_0[57] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 26))) _tmem_load_0[58] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 27))) _tmem_load_0[59] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 28))) _tmem_load_0[60] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 29))) _tmem_load_0[61] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 30))) _tmem_load_0[62] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 31))) _tmem_load_0[63] = -CAKE_INF; - int* index_row = swa_indices + (query_idx * swa_index_stride + tile * 128 + col_half * 64); - #pragma unroll - for (int chunk = 0; chunk < 64; chunk += 4) { - if (valid_cols > chunk) { - int _vec_load_1[4]; - { - const int4* _ivptr_2 = reinterpret_cast(index_row + chunk + 0); - int4 _ivld_2; - _ivld_2 = *_ivptr_2; - _vec_load_1[0 + 0] = _ivld_2.x; - _vec_load_1[0 + 1] = _ivld_2.y; - _vec_load_1[0 + 2] = _ivld_2.z; - _vec_load_1[0 + 3] = _ivld_2.w; - } - #pragma unroll - for (int i = 0; i < 4; i++) { - if (_vec_load_1[i] < 0) { - _tmem_load_0[chunk + i] = -CAKE_INF; - } - } - } - } - float2 _reg_reduce_max2_3 = {-CAKE_INF, -CAKE_INF}; - row_max_x32_accum(&_tmem_load_0[0], _reg_reduce_max2_3); - row_max_x32_accum(&_tmem_load_0[32], _reg_reduce_max2_3); - float _tmem_load_0_max = row_max_reduce(_reg_reduce_max2_3); - float tile_max = _tmem_load_0_max; - float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, tile_max, 16); - float _max_3 = max_noftz(tile_max, _shfl_xor_0); - tile_max = _max_3; - float _max_4 = max_noftz(row_max, tile_max); - float new_max = _max_4; - float _fma_0 = __fmaf_rn(row_max, softmax_scale_log2, (-new_max) * softmax_scale_log2); - float delta = _fma_0; - float _exp2_0 = approx_exp2(delta); - float acc_scale = ((row_max > -CAKE_INF) ? _exp2_0 : 1.0f); - float safe_max = ((new_max == -CAKE_INF) ? 0.0f : new_max); - float max_scaled = safe_max * softmax_scale_log2; - const float2 _fma_b2_4 = {softmax_scale_log2, softmax_scale_log2}; - const float2 _fma_c2_5 = {-max_scaled, -max_scaled}; - #pragma unroll - for (int _lf = 0; _lf < 32; _lf++) - fma_f32x2_inplace(&reinterpret_cast(_tmem_load_0)[_lf], _fma_b2_4, _fma_c2_5); - #pragma unroll - for (int _le = 0; _le < 64; _le++) { - _tmem_load_0[_le] = approx_exp2(_tmem_load_0[_le]); - } - unsigned int packed_p[32]; - #pragma unroll - for (int _lp = 0; _lp < 32; _lp++) { - __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_0[_lp*2 + 0], _tmem_load_0[_lp*2+1 + 0])); - packed_p[_lp] = *(uint32_t*)&_bf2; - } - int p_addr = taddr + (unsigned int)score_col + 64 + (unsigned int)(tmem_row_origin << 16); - asm volatile( - "tcgen05.st.sync.aligned.16x32bx2.x16.b32" - " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" - :: "r"(p_addr), "r"(*reinterpret_cast(&packed_p[0])), "r"(*reinterpret_cast(&packed_p[1])), "r"(*reinterpret_cast(&packed_p[2])), "r"(*reinterpret_cast(&packed_p[3])), "r"(*reinterpret_cast(&packed_p[4])), "r"(*reinterpret_cast(&packed_p[5])), "r"(*reinterpret_cast(&packed_p[6])), "r"(*reinterpret_cast(&packed_p[7])), "r"(*reinterpret_cast(&packed_p[8])), "r"(*reinterpret_cast(&packed_p[9])), "r"(*reinterpret_cast(&packed_p[10])), "r"(*reinterpret_cast(&packed_p[11])), "r"(*reinterpret_cast(&packed_p[12])), "r"(*reinterpret_cast(&packed_p[13])), "r"(*reinterpret_cast(&packed_p[14])), "r"(*reinterpret_cast(&packed_p[15]))); - asm volatile( - "tcgen05.st.sync.aligned.16x32bx2.x16.b32" - " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" - :: "r"(p_addr + 16), "r"(*reinterpret_cast(&(packed_p + 16)[0])), "r"(*reinterpret_cast(&(packed_p + 16)[1])), "r"(*reinterpret_cast(&(packed_p + 16)[2])), "r"(*reinterpret_cast(&(packed_p + 16)[3])), "r"(*reinterpret_cast(&(packed_p + 16)[4])), "r"(*reinterpret_cast(&(packed_p + 16)[5])), "r"(*reinterpret_cast(&(packed_p + 16)[6])), "r"(*reinterpret_cast(&(packed_p + 16)[7])), "r"(*reinterpret_cast(&(packed_p + 16)[8])), "r"(*reinterpret_cast(&(packed_p + 16)[9])), "r"(*reinterpret_cast(&(packed_p + 16)[10])), "r"(*reinterpret_cast(&(packed_p + 16)[11])), "r"(*reinterpret_cast(&(packed_p + 16)[12])), "r"(*reinterpret_cast(&(packed_p + 16)[13])), "r"(*reinterpret_cast(&(packed_p + 16)[14])), "r"(*reinterpret_cast(&(packed_p + 16)[15]))); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - if (elect_sync()) { - mbarrier_arrive(p_full_addr + (phase) * 8); - } - float2 _reg_reduce_sum2_6 = make_float2(0.0f, 0.0f); - softmax_block_sum(&_tmem_load_0[0], &_reg_reduce_sum2_6); - softmax_block_sum(&_tmem_load_0[32], &_reg_reduce_sum2_6); - float _tmem_load_0_sum = _reg_reduce_sum2_6.x + _reg_reduce_sum2_6.y; - float block_sum = _tmem_load_0_sum; - float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, block_sum, 16); - block_sum = block_sum + _shfl_xor_1; - float _fma_1 = __fmaf_rn(row_sum, acc_scale, block_sum); - row_sum = _fma_1; - row_max = new_max; - } - } - if (sparse_topk == 128) { - mbarrier_wait(pv_done_addr, 0); - asm volatile("tcgen05.fence::after_thread_sync;"); - float output_scale = bmm2_scale[0]; - float _rcp_0 = approx_rcp(row_sum); - float inv_sum = ((row_sum > 0.0f) ? _rcp_0 : 0.0f); - int output_base = (query_idx * num_heads + my_row) * 512 + blockIdx.y * 256; - #pragma unroll - for (int v_stage = 0; v_stage < 2; v_stage++) { - int output_col = 256 + v_stage * 128; - int output_addr = taddr + (unsigned int)output_col + (unsigned int)(tmem_row_origin << 16); - float _tmem_load_1[64]; - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_1[0])), "=r"(*reinterpret_cast(&_tmem_load_1[1])), "=r"(*reinterpret_cast(&_tmem_load_1[2])), "=r"(*reinterpret_cast(&_tmem_load_1[3])), "=r"(*reinterpret_cast(&_tmem_load_1[4])), "=r"(*reinterpret_cast(&_tmem_load_1[5])), "=r"(*reinterpret_cast(&_tmem_load_1[6])), "=r"(*reinterpret_cast(&_tmem_load_1[7])), "=r"(*reinterpret_cast(&_tmem_load_1[8])), "=r"(*reinterpret_cast(&_tmem_load_1[9])), "=r"(*reinterpret_cast(&_tmem_load_1[10])), "=r"(*reinterpret_cast(&_tmem_load_1[11])), "=r"(*reinterpret_cast(&_tmem_load_1[12])), "=r"(*reinterpret_cast(&_tmem_load_1[13])), "=r"(*reinterpret_cast(&_tmem_load_1[14])), "=r"(*reinterpret_cast(&_tmem_load_1[15])), "=r"(*reinterpret_cast(&_tmem_load_1[16])), "=r"(*reinterpret_cast(&_tmem_load_1[17])), "=r"(*reinterpret_cast(&_tmem_load_1[18])), "=r"(*reinterpret_cast(&_tmem_load_1[19])), "=r"(*reinterpret_cast(&_tmem_load_1[20])), "=r"(*reinterpret_cast(&_tmem_load_1[21])), "=r"(*reinterpret_cast(&_tmem_load_1[22])), "=r"(*reinterpret_cast(&_tmem_load_1[23])), "=r"(*reinterpret_cast(&_tmem_load_1[24])), "=r"(*reinterpret_cast(&_tmem_load_1[25])), "=r"(*reinterpret_cast(&_tmem_load_1[26])), "=r"(*reinterpret_cast(&_tmem_load_1[27])), "=r"(*reinterpret_cast(&_tmem_load_1[28])), "=r"(*reinterpret_cast(&_tmem_load_1[29])), "=r"(*reinterpret_cast(&_tmem_load_1[30])), "=r"(*reinterpret_cast(&_tmem_load_1[31])) - : "r"(output_addr)); - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_1[32])), "=r"(*reinterpret_cast(&_tmem_load_1[33])), "=r"(*reinterpret_cast(&_tmem_load_1[34])), "=r"(*reinterpret_cast(&_tmem_load_1[35])), "=r"(*reinterpret_cast(&_tmem_load_1[36])), "=r"(*reinterpret_cast(&_tmem_load_1[37])), "=r"(*reinterpret_cast(&_tmem_load_1[38])), "=r"(*reinterpret_cast(&_tmem_load_1[39])), "=r"(*reinterpret_cast(&_tmem_load_1[40])), "=r"(*reinterpret_cast(&_tmem_load_1[41])), "=r"(*reinterpret_cast(&_tmem_load_1[42])), "=r"(*reinterpret_cast(&_tmem_load_1[43])), "=r"(*reinterpret_cast(&_tmem_load_1[44])), "=r"(*reinterpret_cast(&_tmem_load_1[45])), "=r"(*reinterpret_cast(&_tmem_load_1[46])), "=r"(*reinterpret_cast(&_tmem_load_1[47])), "=r"(*reinterpret_cast(&_tmem_load_1[48])), "=r"(*reinterpret_cast(&_tmem_load_1[49])), "=r"(*reinterpret_cast(&_tmem_load_1[50])), "=r"(*reinterpret_cast(&_tmem_load_1[51])), "=r"(*reinterpret_cast(&_tmem_load_1[52])), "=r"(*reinterpret_cast(&_tmem_load_1[53])), "=r"(*reinterpret_cast(&_tmem_load_1[54])), "=r"(*reinterpret_cast(&_tmem_load_1[55])), "=r"(*reinterpret_cast(&_tmem_load_1[56])), "=r"(*reinterpret_cast(&_tmem_load_1[57])), "=r"(*reinterpret_cast(&_tmem_load_1[58])), "=r"(*reinterpret_cast(&_tmem_load_1[59])), "=r"(*reinterpret_cast(&_tmem_load_1[60])), "=r"(*reinterpret_cast(&_tmem_load_1[61])), "=r"(*reinterpret_cast(&_tmem_load_1[62])), "=r"(*reinterpret_cast(&_tmem_load_1[63])) - : "r"(output_addr + 32)); - int col_base = v_stage * 128 + col_half * 64; - if (my_row < num_heads) { - #pragma unroll - for (int offset = 0; offset < 64; offset += 8) { - { - const float2 _prescale2_7 = {inv_sum * output_scale, inv_sum * output_scale}; - #if __CUDA_ARCH__ >= 1000 - #pragma unroll - for (int _ps = 0; _ps < 4; _ps++) - mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_1[offset])[_ps], _prescale2_7); - #else - #pragma unroll - for (int _ps = 0; _ps < 8; _ps++) - _tmem_load_1[offset + _ps] *= inv_sum * output_scale; - #endif - __nv_bfloat162 _pk[4]; - _pk[0] = __floats2bfloat162_rn(_tmem_load_1[offset + 0], _tmem_load_1[offset + 1]); - _pk[1] = __floats2bfloat162_rn(_tmem_load_1[offset + 2], _tmem_load_1[offset + 3]); - _pk[2] = __floats2bfloat162_rn(_tmem_load_1[offset + 4], _tmem_load_1[offset + 5]); - _pk[3] = __floats2bfloat162_rn(_tmem_load_1[offset + 6], _tmem_load_1[offset + 7]); - *reinterpret_cast(&((__nv_bfloat16*)(O + (output_base + col_base + offset)))[0]) = *reinterpret_cast(&_pk[0]); - } - } - } - } - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - if (elect_sync()) { - mbarrier_arrive(tmem_dealloc_addr); - } - } - } - // ---- Role: mma_warp ---- - if (warp == 4) { - { // mma_warp_main - int query_idx_1 = blockIdx.x; - int _max_1 = ((sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) : (0)); - int _min_1 = ((_max_1) < (sparse_topk) ? (_max_1) : (sparse_topk)); - int active_topk_1 = _min_1; - int tile_count_1 = (active_topk_1 + 128 - 1) / 128; - unsigned int _phase_q_full_0 = 0; - mbarrier_wait(q_full_addr, _phase_q_full_0); - _phase_q_full_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - #pragma unroll - for (int tile_1 = 0; tile_1 < 1; tile_1++) { - if (tile_count_1 > tile_1) { - int phase_1 = tile_1 & 1; - int score_col_1 = ((phase_1 != 0) ? 128 : 0); - #pragma unroll - for (int k_stage = 0; k_stage < 4; k_stage++) { - mbarrier_wait(kv_full_addr + (k_stage) * 8, phase_1); - asm volatile("tcgen05.fence::after_thread_sync;"); - int _mma_a_lo_0 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (k_stage * 2) * 512); - int _mma_b_lo_0 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (k_stage) * 2048); - { - uint64_t _mma_ss_a_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_0); - uint64_t _mma_ss_b_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_0); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, ((k_stage == 0) ? 0 : 1)); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); - } - } - int _mma_a_lo_1 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (k_stage * 2 + 1) * 512); - int _mma_b_lo_1 = make_warp_uniform((((smem_kv_addr + 16384) >> 4) & 0x3FFF) + (k_stage) * 2048); - { - uint64_t _mma_ss_a_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_1); - uint64_t _mma_ss_b_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_1); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); - } - } - if (k_stage == 3) { - elect_commit(s_full_addr + (phase_1) * 8); - } - } - int wait_phase_1 = tile_1 >> 1 & 1; - mbarrier_wait(p_full_addr + (phase_1) * 8, wait_phase_1); - asm volatile("tcgen05.fence::after_thread_sync;"); - #pragma unroll - for (int v_stage_1 = 0; v_stage_1 < 2; v_stage_1++) { - int _mma_b_lo_2 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (blockIdx.y * 2 + v_stage_1) * 2048); - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 dhi, blo, ta, id;\n\t" - ".reg .b64 db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "" - "mov.b32 dhi, 0x40004040;\n\t" - "mov.b32 id, 69272720;\n\t" - "mov.b32 ta, %2;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p0;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "}\n" - :: "r"((tmem_tmem + (256 + v_stage_1 * 128))), "r"(_mma_b_lo_2), "r"(tmem_tmem + (score_col_1 + 64)), "r"(0)); - } - elect_commit(pv_done_addr); - } - } - unsigned int _phase_tmem_dealloc_0 = 0; - mbarrier_wait(tmem_dealloc_addr, _phase_tmem_dealloc_0); - _phase_tmem_dealloc_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); - asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); - } - } - // ---- Role: load_warp ---- - if (warp >= 5 && warp <= 7) { - { // load_warp_main - const int load_warp_rank = warp - 5; - int query_idx_2 = blockIdx.x; - int _max_0 = ((sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) : (0)); - int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); - int active_topk_2 = _min_0; - int tile_count_2 = (active_topk_2 + 128 - 1) / 128; - unsigned int load_kv_stage = 0; - unsigned int load_kv_phase = 1; - if (load_warp_rank == 0) { - if (elect_sync()) { - mbarrier_arrive_expect_tx(q_full_addr, 65536); - #pragma unroll - for (int q_stage = 0; q_stage < 8; q_stage += 8) { - tma_4d_gmem2smem(smem_q_addr + (unsigned int)(q_stage * 8192), (&tmap_q), 0, 0, q_stage, query_idx_2, q_full_addr); - } - } - } - int group = load_warp_rank + lane * 3; - int row0 = 0; - int row1 = 0; - int row2 = 0; - int row3 = 0; - if (group < 32) { - int index_base = group * 4; - int* index_row_1 = swa_indices + (query_idx_2 * swa_index_stride); - int raw_rows[4]; - if (index_base + 3 < sparse_topk) { - int _vec_load_0[4]; - { - const int4* _ivptr_0 = reinterpret_cast(index_row_1 + index_base + 0); - int4 _ivld_0; - _ivld_0 = *_ivptr_0; - _vec_load_0[0 + 0] = _ivld_0.x; - _vec_load_0[0 + 1] = _ivld_0.y; - _vec_load_0[0 + 2] = _ivld_0.z; - _vec_load_0[0 + 3] = _ivld_0.w; - } - #pragma unroll - for (int i_1 = 0; i_1 < 4; i_1++) { - raw_rows[i_1] = _vec_load_0[i_1]; - } - } else { - #pragma unroll - for (int i_2 = 0; i_2 < 4; i_2++) { - raw_rows[i_2] = -1; - if (index_base + i_2 < sparse_topk) { - raw_rows[i_2] = index_row_1[index_base + i_2]; - } - } - } - row0 = ((raw_rows[0] >= 0) ? raw_rows[0] : 0); - row1 = ((raw_rows[1] >= 0) ? raw_rows[1] : 0); - row2 = ((raw_rows[2] >= 0) ? raw_rows[2] : 0); - row3 = ((raw_rows[3] >= 0) ? raw_rows[3] : 0); - } - #pragma unroll - for (int tile_2 = 0; tile_2 < 1; tile_2++) { - if (tile_count_2 > tile_2) { - #pragma unroll - for (int k_stage_1 = 0; k_stage_1 < 4; k_stage_1++) { - if (elect_sync()) { - int producer_groups = ((load_warp_rank < 2) ? 11 : 10); - mbarrier_arrive_expect_tx(kv_full_addr + (load_kv_stage) * 8, producer_groups * 512 * 2); - } - if (group < 32) { - int dst_k = smem_kv_addr + load_kv_stage * 32768; - if (tile_2 == 0) { - tma_gather4_gmem2smem(dst_k + group * 512, (&tmap_swa_kv), k_stage_1 * 128, row0, row1, row2, row3, kv_full_addr + (load_kv_stage) * 8); - tma_gather4_gmem2smem(dst_k + 16384 + group * 512, (&tmap_swa_kv), k_stage_1 * 128 + 64, row0, row1, row2, row3, kv_full_addr + (load_kv_stage) * 8); - } else { - tma_gather4_gmem2smem(dst_k + group * 512, (&tmap_compressed_kv), k_stage_1 * 128, row0, row1, row2, row3, kv_full_addr + (load_kv_stage) * 8); - tma_gather4_gmem2smem(dst_k + 16384 + group * 512, (&tmap_compressed_kv), k_stage_1 * 128 + 64, row0, row1, row2, row3, kv_full_addr + (load_kv_stage) * 8); - } - } - load_kv_stage += 1; - if (load_kv_stage == 4) { load_kv_stage = 0; load_kv_phase ^= 1; } - } - } - } - if (elect_sync()) { - mbarrier_arrive(tmem_dealloc_addr); - } - } - } - - // Cleanup -} - -} // extern "C" diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_4ac65f00a4a032c10ecb_binding.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_4ac65f00a4a032c10ecb_binding.cu deleted file mode 100644 index 245c914ddcf..00000000000 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_4ac65f00a4a032c10ecb_binding.cu +++ /dev/null @@ -1,385 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -// Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_4ac65f00a4a032c10ecb'. -#include -#include -#include -#include - -#include "tvm_ffi_utils.h" - -#include -#include -#include -#include -#include -#include -#include - -extern "C" __global__ void kernel_cake_dsv4_4ac65f00a4a032c10ecb(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks); - - -namespace cake_host_shim_71507c1f0a7f757c { - -using tvm::ffi::Optional; -using tvm::ffi::TensorView; - -class ScopedCudaDevice { - public: - explicit ScopedCudaDevice(int device_id) { - cudaError_t error = cudaGetDevice(&previous_device_); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed before host-shim launch: cudaError=" - << static_cast(error); - if (previous_device_ != device_id) { - error = cudaSetDevice(device_id); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaSetDevice failed before host-shim launch for cuda:" - << device_id << ": cudaError=" << static_cast(error); - restore_ = true; - } - } - - ScopedCudaDevice(const ScopedCudaDevice&) = delete; - ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; - - ~ScopedCudaDevice() noexcept { - if (restore_) { - (void)cudaSetDevice(previous_device_); - } - } - - private: - int previous_device_ = -1; - bool restore_ = false; -}; - -inline int64_t CakeDeviceMultiprocessorCount(int device_id) { - constexpr int kMaxCachedCudaDevices = 64; - TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) - << "physical-SM-count cache does not cover cuda:" << device_id; - static std::atomic count_by_device[kMaxCachedCudaDevices]{}; - int cached = count_by_device[device_id].load(std::memory_order_acquire); - if (cached > 0) return cached; - - int count = 0; - cudaError_t error = cudaDeviceGetAttribute( - &count, cudaDevAttrMultiProcessorCount, device_id); - TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) - << "querying multiProcessorCount failed for cuda:" << device_id - << ": cudaError=" << static_cast(error) << ", count=" << count; - // Concurrent first launches may repeat the immutable device query, but all - // publication is atomic and every later hot-path lookup is one acquire load. - count_by_device[device_id].store(count, std::memory_order_release); - return count; -} - -inline void CheckCudaTensor(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) - << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; -} - -inline void CheckSameCudaDevice( - const TensorView& t, - const TensorView& reference, - const char* name, - const char* reference_name) { - TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) - << name << " must be on the same CUDA device as " << reference_name - << ": got cuda:" << t.device().device_id - << " versus cuda:" << reference.device().device_id; -} - -inline void CheckCurrentCudaDevice( - const TensorView& reference, - const char* reference_name) { - int current_device = -1; - cudaError_t error = cudaGetDevice(¤t_device); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed while validating " << reference_name - << ": cudaError=" << static_cast(error); - TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) - << "current CUDA device must match " << reference_name - << ": current=cuda:" << current_device - << ", tensor=cuda:" << reference.device().device_id; -} - -inline void CheckContiguous(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; -} - -inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { - DLDataType d = t.dtype(); - TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) - << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits - << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits - << ", lanes=" << (int)d.lanes << ")"; -} - -// A logical axis.outer(trailing) folds every source dim above the trailing -// dimensions. Shape products are independent of physical strides, so verify -// the leading dimensions form one dense row-major chain instead of inventing -// a "folded stride". The descriptor reads its exact adjacent physical step -// separately through stride[-(trailing + 1)]. -inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { - TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) - << name << " cannot fold leading dimensions above " << trailing - << " trailing dims from ndim=" << t.ndim(); - int outer_last = t.ndim() - trailing - 1; - if (outer_last <= 0) { - return; - } - int64_t step = t.stride(outer_last); - TVM_FFI_CHECK(step > 0, ValueError) - << name << " physical strides must be positive"; - int64_t expected = step; - for (int axis = outer_last - 1; axis >= 0; --axis) { - expected *= t.size(axis + 1); - if (t.size(axis) > 1) { - TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) - << name << " leading dims are not physically foldable above " << trailing - << " trailing dims: stride(" << axis << ")=" << t.stride(axis) - << ", expected " << expected; - } - } -} - -// 4D TMA descriptor for buffer 'tmap_q' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 3, ValueError) - << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - int64_t d2 = t.size(t.ndim() - 2); - int64_t d3 = t.size(t.ndim() - 3); - TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) - << "TMA source 'tmap_q' trailing dims must be positive"; - TVM_FFI_CHECK(d1 % 64 == 0, ValueError) - << "TMA source 'tmap_q' extent " << d1 - << " must divide exactly by " << 64; - uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(d3)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) - << "TMA box (64, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; - int64_t carrier_stride_0 = d1; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; - int64_t carrier_stride_1 = 64; - TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; - TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; - TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; - int64_t carrier_stride_2 = (d2 * d1); - TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; - TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; - TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; - uint64_t global_strides[3] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - (uint64_t)((carrier_stride_1 * 16) / 8), - (uint64_t)((carrier_stride_2 * 16) / 8), - }; - uint32_t box_dim[4] = {64u, 64u, 1u, 1u}; - uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_swa_kv' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_swa_kv' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_kv'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; - return tm; -} - -void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_O, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t grid_x, int64_t grid_y, int64_t grid_z) { - DLDevice dev = arg_tmap_q.device(); - ScopedCudaDevice device_guard(dev.device_id); - CheckCudaTensor(arg_tmap_q, "tmap_q"); - CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); - CheckContiguous(arg_tmap_q, "tmap_q"); - CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); - CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 4, 16, 1); - CheckCudaTensor(arg_O, "O"); - CheckDtype(arg_O, "O", 4, 16, 1); - CheckContiguous(arg_O, "O"); - CheckCudaTensor(arg_swa_indices, "swa_indices"); - CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); - CheckContiguous(arg_swa_indices, "swa_indices"); - CheckCudaTensor(arg_compressed_indices, "compressed_indices"); - CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); - CheckContiguous(arg_compressed_indices, "compressed_indices"); - CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); - CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckCudaTensor(arg_sinks, "sinks"); - CheckDtype(arg_sinks, "sinks", 2, 32, 1); - CheckContiguous(arg_sinks, "sinks"); - CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); - CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); - CheckContiguous(arg_bmm1_scale, "bmm1_scale"); - CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); - CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); - CheckContiguous(arg_bmm2_scale, "bmm2_scale"); - TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) - << "scalar 'num_heads' value " << arg_num_heads - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) - << "scalar 'swa_index_stride' value " << arg_swa_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) - << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) - << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) - << "scalar 'sparse_topk' value " << arg_sparse_topk - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) - << "scalar 'has_sinks' value " << arg_has_sinks - << " is outside i32 range [-2147483648, 2147483647]"; - CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); - CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); - CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); - CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); - CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); - CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); - CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); - CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); - TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) - << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y - << ", " << grid_z << ")"; - - cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); - CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); - CUtensorMap p_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); - int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); - int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); - int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); - float* p_sinks = static_cast(arg_sinks.data_ptr()); - float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); - float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); - int32_t v_num_heads = (int32_t)arg_num_heads; - int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; - int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; - int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; - int32_t v_sparse_topk = (int32_t)arg_sparse_topk; - int32_t v_has_sinks = (int32_t)arg_has_sinks; - void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_O, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_has_sinks}; - - static std::mutex smem_mu; - static auto* smem_status_by_device = - new std::unordered_map(); - cudaError_t smem_status = cudaSuccess; - { - std::lock_guard lock(smem_mu); - auto it = smem_status_by_device->find(dev.device_id); - if (it == smem_status_by_device->end()) { - int current_device = -1; - smem_status = cudaGetDevice(¤t_device); - if (smem_status == cudaSuccess && current_device != dev.device_id) { - smem_status = cudaErrorInvalidDevice; - } - if (smem_status == cudaSuccess) { - smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_4ac65f00a4a032c10ecb), - cudaFuncAttributeMaxDynamicSharedMemorySize, - 198528); - } - smem_status_by_device->emplace(dev.device_id, smem_status); - } else { - smem_status = it->second; - } - } - TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_4ac65f00a4a032c10ecb failed: " - << cudaGetErrorString(smem_status); - dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); - dim3 block(512u, 1u, 1u); - - cudaLaunchConfig_t config{}; - config.gridDim = grid; - config.blockDim = block; - config.dynamicSmemBytes = 198528u; - config.stream = stream; - cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_4ac65f00a4a032c10ecb), kargs); - TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_4ac65f00a4a032c10ecb failed: " - << cudaGetErrorString(launch_status); -} - -} // namespace cake_host_shim_71507c1f0a7f757c - -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_71507c1f0a7f757c::Run); diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_4ac65f00a4a032c10ecb_kernel.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_4ac65f00a4a032c10ecb_kernel.cu deleted file mode 100644 index 1e58d437826..00000000000 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_4ac65f00a4a032c10ecb_kernel.cu +++ /dev/null @@ -1,1635 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -typedef signed char int8_t; -typedef unsigned char uint8_t; -typedef unsigned short uint16_t; -typedef unsigned int uint32_t; -#if defined(__CUDACC_RTC__) -typedef unsigned long long uint64_t; -#else -typedef unsigned long uint64_t; -#endif -static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); -typedef signed int int32_t; -typedef short int int16_t; -struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; -struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; -static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); -static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); - -#if defined(__CUDACC_RTC__) -typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; -#else -#include -#endif - -static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); -static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); -#include -#include - -#define CAKE_INF CUDART_INF_F -#define TMEM_NCOLS 384 -#define TMEM_TMEM_OFFSET 0 -#define NUM_KV_PIPE_STAGES 4 -#define SMEM_SMEM_Q_OFF 1024 -#define SMEM_SMEM_Q_STAGE_BYTES 8192 -#define SMEM_SMEM_Q_STRIDE 8192 -#define SMEM_SMEM_KV_OFF 66560 -#define SMEM_SMEM_KV_STAGE_BYTES 32768 -#define SMEM_SMEM_KV_STRIDE 32768 -#define SMEM_SMEM_V_OFF 66560 -#define SMEM_SMEM_V_STAGE_BYTES 32768 -#define SMEM_SMEM_V_STRIDE 32768 -#define SMEM_SMEM_SUM_OFF 197632 -#define SMEM_SMEM_SUM_STAGE_BYTES 256 -#define SMEM_SMEM_SUM_STRIDE 256 -#define SMEM_SMEM_INDEX_FLAGS_OFF 198400 -#define SMEM_SMEM_INDEX_FLAGS_STAGE_BYTES 16 -#define SMEM_SMEM_INDEX_FLAGS_STRIDE 16 -#define SMEM_SMEM_INDICES_OFF 197888 -#define SMEM_SMEM_INDICES_STAGE_BYTES 512 -#define SMEM_SMEM_INDICES_STRIDE 512 -#define SMEM_TOTAL 198528 -#define THREADS 512 - -#include - -__device__ __forceinline__ uint32_t elect_sync() { - uint32_t pred = 0; - asm volatile( - "{\n\t" - ".reg .pred %%px;\n\t" - "elect.sync _|%%px, %1;\n\t" - "@%%px mov.s32 %0, 1;\n\t" - "}\n" - : "+r"(pred) - : "r"(0xFFFFFFFF)); - return pred; -} - - -__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { - asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" - :: "r"(mbar_addr), "r"(count) : "memory"); -} - -__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { - asm volatile("mbarrier.init.b64 [%0], %1;" - :: "l"(mbar_addr), "r"(count) : "memory"); -} - - -__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - - -// CTA-local pipelines have short, resident producer/consumer edges. Omitting -// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the -// explicit loop still makes this helper blocking until acquire succeeds. -__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE;\n\t" - "bra.uni LAB_WAIT;\n\t" - "DONE:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Source-faithful relaxed CTA wait used only by a typed protocol that does -// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. -__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, 10000000;\n\t" - "@P1 bra.uni DONE_RELAXED;\n\t" - "bra.uni LAB_WAIT_RELAXED;\n\t" - "DONE_RELAXED:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Exact source ports may request the PTX suspendTimeHint operand explicitly. -// The hint is expressed in nanoseconds and is kept separate from the canonical -// no-hint CTA helper so unrelated schedules retain their existing retry path. -__device__ __forceinline__ void mbarrier_wait_suspend( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_SUSPEND:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_SUSPEND;\n\t" - "bra.uni LAB_WAIT_SUSPEND;\n\t" - "DONE_SUSPEND:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE_CLUSTER;\n\t" - "bra.uni LAB_WAIT_CLUSTER;\n\t" - "DONE_CLUSTER:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - ".reg .u32 WAIT_ADDR;\n\t" - "mov.u32 WAIT_ADDR, %0;\n\t" - "LAB_WAIT_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [WAIT_ADDR], %1, %2;\n\t" - "@P1 bra.uni DONE_HINT;\n\t" - "bra.uni LAB_WAIT_HINT;\n\t" - "DONE_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -// Exact unqualified CTA wait used by source schedules whose PTX intentionally -// omits the acquire qualifier while retaining a typed suspendTimeHint operand. -__device__ __forceinline__ void mbarrier_wait_relaxed_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED_HINT:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra DONE_RELAXED_HINT;\n\t" - "bra LAB_WAIT_RELAXED_HINT;\n\t" - "DONE_RELAXED_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); -} - -__device__ __forceinline__ void mbarrier_wait_cluster_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" - "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" - "DONE_CLUSTER_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_suspend( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait_cluster(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); - } -} - - -__device__ __forceinline__ void tcgen05_mma_f16( - int taddr, uint64_t a_desc, uint64_t b_desc, - uint32_t i_desc, int enable_input_d) { - asm volatile( - "{\n\t" - ".reg .pred p;\n\t" - "setp.ne.b32 p, %4, 0;\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" - "}\n" - :: "r"(taddr), "l"(a_desc), "l"(b_desc), - "r"(i_desc), "r"(enable_input_d) - : "memory"); -} - - -__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { - return (x & 0x3FFFFULL) >> 4ULL; -} - - -union MmaSmemDesc { - uint64_t u64; - uint32_t u32[2]; -}; - -__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { - MmaSmemDesc tmp; - tmp.u64 = smem_desc; - tmp.u32[0] += offset; - smem_desc = tmp.u64; -} - - -__device__ __forceinline__ void elect_commit(int mbar_addr) { - asm volatile( - "{\n\t" - ".reg .pred leader;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" - ".shared::cluster.b64 [%0];\n\t" - "}\n" - :: "r"(mbar_addr)); -} - - -__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { - asm volatile( - "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" - :: "r"(mbar_addr), "r"(bytes) : "memory"); -} - - -__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7," - " %8, %9, %10, %11, %12, %13, %14, %15," - " %16, %17, %18, %19, %20, %21, %22, %23," - " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), - "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), - "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), - "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), - "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), - "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), - "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ float approx_exp2(float x) { - float y; - asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float approx_rcp(float x) { - float y; - asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float max_noftz(float a, float b) { - float c; - asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); - return c; -} - - -__device__ __forceinline__ float warp_reduce_max(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); - return val; -} - - -__device__ __forceinline__ float warp_reduce_sum(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val += __shfl_xor_sync(0xFFFFFFFF, val, offset); - return val; -} - - -__device__ __forceinline__ float row_max_reduce(float2 acc) { - return max_noftz(acc.x, acc.y); -} - - -__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { - #pragma unroll - for (int j = 0; j < 16; j++) { - if (j % 2 == 0) - acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); - else - acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); - } -} - - -__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { - asm("mul.rn.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { - asm("mul.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { - asm("add.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { - asm("add.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { - asm("sub.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { - asm("sub.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("add.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { - float2 r; - asm("sub.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("sub.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ void fma_scale_x32( - float* sv, const float2* scale2, const float2* neg_max2) -{ - float2* sv_2 = reinterpret_cast(sv); - #pragma unroll - for (int j = 0; j < 16; j++) - fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); -} - -__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("mul.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - - -// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) - -__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("add.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("add.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("add.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("add.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("add.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("add.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("add.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - - -__device__ __forceinline__ void fence_async_shared() { - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); -} - - -__device__ __forceinline__ uint64_t make_smem_desc(int addr) { - const int SBO = 1024; - return desc_encode(addr) - | (desc_encode(SBO) << 32ULL) - | (1ULL << 46ULL) - | (2ULL << 61ULL); -} - - -__device__ __forceinline__ void tma_2d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3}], [%4];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_4d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.4d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5}], [%6];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr) { - // Canonical .shared::cta form for non-multicast gather4, matching - // trtllm-gen / cuda_ptx and the PTX ISA qualifier order - // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, - // .shared::cluster is reserved for the multicast variant (ctaMask). - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { - asm volatile( - "tcgen05.commit.cta_group::1.mbarrier::arrive::one" - ".shared::cluster.b64 [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { - uint32_t result; - asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" - : "=r"(result) : "r"(val)); - return result; -} - -extern "C" { - -__global__ __launch_bounds__(512, 1) void -kernel_cake_dsv4_4ac65f00a4a032c10ecb(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks) -{ - const int tid = threadIdx.x; - const int warp = make_warp_uniform(tid / 32); - const int lane = tid % 32; - - extern __shared__ __align__(1024) char smem_raw[]; - int smem; - asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); - smem = make_warp_uniform(smem); - - const int mbar_base = smem; - #define q_full_addr (mbar_base + 0) - #define kv_full_addr (mbar_base + 8) - #define kv_empty_addr (mbar_base + 40) - #define s_full_addr (mbar_base + 72) - #define p_full_addr (mbar_base + 80) - #define sum_ready_addr (mbar_base + 88) - #define o_done_addr (mbar_base + 96) - #define tmem_dealloc_addr (mbar_base + 104) - - const int bid = blockIdx.x; - const int num_bids = gridDim.x; - - const int cta_rank = 0; - - // Kernel setup ops - __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); - const int smem_q_addr = smem + 1024; - __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); - const int smem_kv_addr = smem + 66560; - __nv_bfloat16* smem_v = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); - const int smem_v_addr = smem + 66560; - float* smem_sum = reinterpret_cast(smem_raw + 197632); - const int smem_sum_addr = smem + 197632; - int* smem_index_flags = reinterpret_cast(smem_raw + 198400); - const int smem_index_flags_addr = smem + 198400; - int* smem_indices = reinterpret_cast(smem_raw + 197888); - const int smem_indices_addr = smem + 197888; - - // Mbarrier init (8 pipeline groups, 0 ordered-sequence groups, 14 barriers) - // Mbarriers at smem_raw[0..112) - - if (warp == 0) { - uint32_t leader = elect_sync(); - if (leader) { - // q_full: 1 barriers, init_count=1 - mbarrier_init(smem + 0, 1); - // --- pipeline 'kv_pipe' --- - // kv_full: 4 barriers, init_count=4 - mbarrier_init(smem + 8, 4); - mbarrier_init(smem + 16, 4); - mbarrier_init(smem + 24, 4); - mbarrier_init(smem + 32, 4); - // kv_empty: 4 barriers, init_count=1 - mbarrier_init(smem + 40, 1); - mbarrier_init(smem + 48, 1); - mbarrier_init(smem + 56, 1); - mbarrier_init(smem + 64, 1); - // s_full: 1 barriers, init_count=1 - mbarrier_init(smem + 72, 1); - // p_full: 1 barriers, init_count=128 - mbarrier_init(smem + 80, 128); - // sum_ready: 1 barriers, init_count=128 - mbarrier_init(smem + 88, 128); - // o_done: 1 barriers, init_count=1 - mbarrier_init(smem + 96, 1); - // tmem_dealloc: 1 barriers, init_count=256 - mbarrier_init(smem + 104, 256); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - - __syncwarp(); - - // TMEM alloc (512 columns, 384 used) - volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 112); - if (warp == 0) { - int _tmem_hold = smem + 112; - asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); - __syncwarp(); - asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); - } - - __syncthreads(); - asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); - - const int taddr = tmem_addr_storage[0]; - - // Kernel post-init ops - const int tmem_tmem = taddr; - - // ---- Ordered hardware-WG register redistribution ---- - // Dec phase frees registers before any WG attempts inc. - if (warp >= 8 && warp <= 11) { - asm volatile("setmaxnreg.dec.sync.aligned.u32 64;"); - } - - // ---- Role: softmax ---- - if (warp <= 3) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); - { // softmax_main - float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; - int query_idx = blockIdx.x >> 2; - int _max_0 = ((sparse_topk_lens[query_idx] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx] + sparse_topk_lens_offset) : (0)); - int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); - int active_topk = _min_0; - const int warp_in_compute = warp; - const int tmem_row_origin = warp_in_compute * 32; - const int logical_row_origin = warp_in_compute * 16; - const int my_row = logical_row_origin + lane % 16; - const int col_half = lane / 16; - unsigned int _phase_s_full_0 = 0; - mbarrier_wait(s_full_addr, _phase_s_full_0); - _phase_s_full_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - int score_addr = taddr + (unsigned int)(tmem_row_origin << 16); - float _tmem_load_0[64]; - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_0[0])), "=r"(*reinterpret_cast(&_tmem_load_0[1])), "=r"(*reinterpret_cast(&_tmem_load_0[2])), "=r"(*reinterpret_cast(&_tmem_load_0[3])), "=r"(*reinterpret_cast(&_tmem_load_0[4])), "=r"(*reinterpret_cast(&_tmem_load_0[5])), "=r"(*reinterpret_cast(&_tmem_load_0[6])), "=r"(*reinterpret_cast(&_tmem_load_0[7])), "=r"(*reinterpret_cast(&_tmem_load_0[8])), "=r"(*reinterpret_cast(&_tmem_load_0[9])), "=r"(*reinterpret_cast(&_tmem_load_0[10])), "=r"(*reinterpret_cast(&_tmem_load_0[11])), "=r"(*reinterpret_cast(&_tmem_load_0[12])), "=r"(*reinterpret_cast(&_tmem_load_0[13])), "=r"(*reinterpret_cast(&_tmem_load_0[14])), "=r"(*reinterpret_cast(&_tmem_load_0[15])), "=r"(*reinterpret_cast(&_tmem_load_0[16])), "=r"(*reinterpret_cast(&_tmem_load_0[17])), "=r"(*reinterpret_cast(&_tmem_load_0[18])), "=r"(*reinterpret_cast(&_tmem_load_0[19])), "=r"(*reinterpret_cast(&_tmem_load_0[20])), "=r"(*reinterpret_cast(&_tmem_load_0[21])), "=r"(*reinterpret_cast(&_tmem_load_0[22])), "=r"(*reinterpret_cast(&_tmem_load_0[23])), "=r"(*reinterpret_cast(&_tmem_load_0[24])), "=r"(*reinterpret_cast(&_tmem_load_0[25])), "=r"(*reinterpret_cast(&_tmem_load_0[26])), "=r"(*reinterpret_cast(&_tmem_load_0[27])), "=r"(*reinterpret_cast(&_tmem_load_0[28])), "=r"(*reinterpret_cast(&_tmem_load_0[29])), "=r"(*reinterpret_cast(&_tmem_load_0[30])), "=r"(*reinterpret_cast(&_tmem_load_0[31])) - : "r"(score_addr)); - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_0[32])), "=r"(*reinterpret_cast(&_tmem_load_0[33])), "=r"(*reinterpret_cast(&_tmem_load_0[34])), "=r"(*reinterpret_cast(&_tmem_load_0[35])), "=r"(*reinterpret_cast(&_tmem_load_0[36])), "=r"(*reinterpret_cast(&_tmem_load_0[37])), "=r"(*reinterpret_cast(&_tmem_load_0[38])), "=r"(*reinterpret_cast(&_tmem_load_0[39])), "=r"(*reinterpret_cast(&_tmem_load_0[40])), "=r"(*reinterpret_cast(&_tmem_load_0[41])), "=r"(*reinterpret_cast(&_tmem_load_0[42])), "=r"(*reinterpret_cast(&_tmem_load_0[43])), "=r"(*reinterpret_cast(&_tmem_load_0[44])), "=r"(*reinterpret_cast(&_tmem_load_0[45])), "=r"(*reinterpret_cast(&_tmem_load_0[46])), "=r"(*reinterpret_cast(&_tmem_load_0[47])), "=r"(*reinterpret_cast(&_tmem_load_0[48])), "=r"(*reinterpret_cast(&_tmem_load_0[49])), "=r"(*reinterpret_cast(&_tmem_load_0[50])), "=r"(*reinterpret_cast(&_tmem_load_0[51])), "=r"(*reinterpret_cast(&_tmem_load_0[52])), "=r"(*reinterpret_cast(&_tmem_load_0[53])), "=r"(*reinterpret_cast(&_tmem_load_0[54])), "=r"(*reinterpret_cast(&_tmem_load_0[55])), "=r"(*reinterpret_cast(&_tmem_load_0[56])), "=r"(*reinterpret_cast(&_tmem_load_0[57])), "=r"(*reinterpret_cast(&_tmem_load_0[58])), "=r"(*reinterpret_cast(&_tmem_load_0[59])), "=r"(*reinterpret_cast(&_tmem_load_0[60])), "=r"(*reinterpret_cast(&_tmem_load_0[61])), "=r"(*reinterpret_cast(&_tmem_load_0[62])), "=r"(*reinterpret_cast(&_tmem_load_0[63])) - : "r"(score_addr + 32)); - int valid_cols = active_topk - col_half * 64; - int causal_cols = smem_index_flags[col_half * 2 + 1]; - valid_cols = ((valid_cols < causal_cols) ? valid_cols : causal_cols); - uint32_t _slice_lo_mask_0; - { - int _lim_0 = valid_cols; - if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } - else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); - } - } - if (!(_slice_lo_mask_0 & (1u << 0))) _tmem_load_0[0] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 1))) _tmem_load_0[1] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 2))) _tmem_load_0[2] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 3))) _tmem_load_0[3] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 4))) _tmem_load_0[4] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 5))) _tmem_load_0[5] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 6))) _tmem_load_0[6] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 7))) _tmem_load_0[7] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 8))) _tmem_load_0[8] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 9))) _tmem_load_0[9] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 10))) _tmem_load_0[10] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 11))) _tmem_load_0[11] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 12))) _tmem_load_0[12] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 13))) _tmem_load_0[13] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 14))) _tmem_load_0[14] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 15))) _tmem_load_0[15] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 16))) _tmem_load_0[16] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 17))) _tmem_load_0[17] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 18))) _tmem_load_0[18] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 19))) _tmem_load_0[19] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 20))) _tmem_load_0[20] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 21))) _tmem_load_0[21] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 22))) _tmem_load_0[22] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 23))) _tmem_load_0[23] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 24))) _tmem_load_0[24] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 25))) _tmem_load_0[25] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 26))) _tmem_load_0[26] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 27))) _tmem_load_0[27] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 28))) _tmem_load_0[28] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 29))) _tmem_load_0[29] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 30))) _tmem_load_0[30] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 31))) _tmem_load_0[31] = -CAKE_INF; - uint32_t _slice_lo_mask_1; - { - int _lim_1 = valid_cols - 32; - if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } - else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); - } - } - if (!(_slice_lo_mask_1 & (1u << 0))) _tmem_load_0[32] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 1))) _tmem_load_0[33] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 2))) _tmem_load_0[34] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 3))) _tmem_load_0[35] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 4))) _tmem_load_0[36] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 5))) _tmem_load_0[37] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 6))) _tmem_load_0[38] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 7))) _tmem_load_0[39] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 8))) _tmem_load_0[40] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 9))) _tmem_load_0[41] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 10))) _tmem_load_0[42] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 11))) _tmem_load_0[43] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 12))) _tmem_load_0[44] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 13))) _tmem_load_0[45] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 14))) _tmem_load_0[46] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 15))) _tmem_load_0[47] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 16))) _tmem_load_0[48] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 17))) _tmem_load_0[49] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 18))) _tmem_load_0[50] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 19))) _tmem_load_0[51] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 20))) _tmem_load_0[52] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 21))) _tmem_load_0[53] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 22))) _tmem_load_0[54] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 23))) _tmem_load_0[55] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 24))) _tmem_load_0[56] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 25))) _tmem_load_0[57] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 26))) _tmem_load_0[58] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 27))) _tmem_load_0[59] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 28))) _tmem_load_0[60] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 29))) _tmem_load_0[61] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 30))) _tmem_load_0[62] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 31))) _tmem_load_0[63] = -CAKE_INF; - float2 _reg_reduce_max2_2 = {-CAKE_INF, -CAKE_INF}; - row_max_x32_accum(&_tmem_load_0[0], _reg_reduce_max2_2); - row_max_x32_accum(&_tmem_load_0[32], _reg_reduce_max2_2); - float _tmem_load_0_max = row_max_reduce(_reg_reduce_max2_2); - float row_max = _tmem_load_0_max; - float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, row_max, 16); - float _max_1 = max_noftz(row_max, _shfl_xor_0); - row_max = _max_1; - if (has_sinks != 0 && my_row < num_heads) { - float sink_unscaled = sinks[my_row] * 1.4426950408889634f / softmax_scale_log2; - float _max_2 = max_noftz(row_max, sink_unscaled); - row_max = _max_2; - } - bool empty_row = row_max == -CAKE_INF; - float safe_max = ((empty_row) ? 0.0f : row_max); - float max_scaled = safe_max * softmax_scale_log2; - float2 _f2_0 = make_float2(softmax_scale_log2, softmax_scale_log2); - float2 _f2_1 = make_float2(-max_scaled, -max_scaled); - float2 _f2_2 = make_float2(0.0f, 0.0f); - float2 sum0 = _f2_2; - float2 _f2_3 = make_float2(0.0f, 0.0f); - float2 sum1 = _f2_3; - float2 _f2_4 = make_float2(0.0f, 0.0f); - float2 sum2 = _f2_4; - float2 _f2_5 = make_float2(0.0f, 0.0f); - float2 sum3 = _f2_5; - #pragma unroll - for (int i = 0; i < 64; i += 8) { - float2 _f2_6 = make_float2(_tmem_load_0[i], _tmem_load_0[i + 1]); - float2 _f2_7 = make_float2(_tmem_load_0[i + 2], _tmem_load_0[i + 3]); - float2 _f2_8 = make_float2(_tmem_load_0[i + 4], _tmem_load_0[i + 5]); - float2 _f2_9 = make_float2(_tmem_load_0[i + 6], _tmem_load_0[i + 7]); - float2 affine0 = fma_f32x2_rn_ftz(_f2_6, _f2_0, _f2_1); - float2 affine1 = fma_f32x2_rn_ftz(_f2_7, _f2_0, _f2_1); - float2 affine2 = fma_f32x2_rn_ftz(_f2_8, _f2_0, _f2_1); - float2 affine3 = fma_f32x2_rn_ftz(_f2_9, _f2_0, _f2_1); - float _exp2_0 = approx_exp2(affine0.x); - float exp0 = _exp2_0; - float _exp2_1 = approx_exp2(affine0.y); - float exp1 = _exp2_1; - float _exp2_2 = approx_exp2(affine1.x); - float exp2 = _exp2_2; - float _exp2_3 = approx_exp2(affine1.y); - float exp3 = _exp2_3; - float _exp2_4 = approx_exp2(affine2.x); - float exp4 = _exp2_4; - float _exp2_5 = approx_exp2(affine2.y); - float exp5 = _exp2_5; - float _exp2_6 = approx_exp2(affine3.x); - float exp6 = _exp2_6; - float _exp2_7 = approx_exp2(affine3.y); - float exp7 = _exp2_7; - _tmem_load_0[i] = exp0; - _tmem_load_0[i + 1] = exp1; - _tmem_load_0[i + 2] = exp2; - _tmem_load_0[i + 3] = exp3; - _tmem_load_0[i + 4] = exp4; - _tmem_load_0[i + 5] = exp5; - _tmem_load_0[i + 6] = exp6; - _tmem_load_0[i + 7] = exp7; - float2 _f2_10 = make_float2(exp0, exp1); - sum0 = add_f32x2(sum0, _f2_10); - float2 _f2_11 = make_float2(exp2, exp3); - sum1 = add_f32x2(sum1, _f2_11); - float2 _f2_12 = make_float2(exp4, exp5); - sum2 = add_f32x2(sum2, _f2_12); - float2 _f2_13 = make_float2(exp6, exp7); - sum3 = add_f32x2(sum3, _f2_13); - } - float2 sum01 = add_f32x2(sum0, sum1); - float2 sum23 = add_f32x2(sum2, sum3); - float2 row_sum_pair = add_f32x2(sum01, sum23); - float row_sum = row_sum_pair.x + row_sum_pair.y; - float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, row_sum, 16); - row_sum = row_sum + _shfl_xor_1; - if (has_sinks != 0 && my_row < num_heads) { - float _exp2_8 = approx_exp2(sinks[my_row] * 1.4426950408889634f - max_scaled); - row_sum = row_sum + _exp2_8; - } - unsigned int packed_p[32]; - #pragma unroll - for (int _lp = 0; _lp < 32; _lp++) { - __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_0[_lp*2 + 0], _tmem_load_0[_lp*2+1 + 0])); - packed_p[_lp] = *(uint32_t*)&_bf2; - } - int p_addr = taddr + 128 + (unsigned int)(tmem_row_origin << 16); - asm volatile( - "tcgen05.st.sync.aligned.16x32bx2.x16.b32" - " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" - :: "r"(p_addr), "r"(*reinterpret_cast(&packed_p[0])), "r"(*reinterpret_cast(&packed_p[1])), "r"(*reinterpret_cast(&packed_p[2])), "r"(*reinterpret_cast(&packed_p[3])), "r"(*reinterpret_cast(&packed_p[4])), "r"(*reinterpret_cast(&packed_p[5])), "r"(*reinterpret_cast(&packed_p[6])), "r"(*reinterpret_cast(&packed_p[7])), "r"(*reinterpret_cast(&packed_p[8])), "r"(*reinterpret_cast(&packed_p[9])), "r"(*reinterpret_cast(&packed_p[10])), "r"(*reinterpret_cast(&packed_p[11])), "r"(*reinterpret_cast(&packed_p[12])), "r"(*reinterpret_cast(&packed_p[13])), "r"(*reinterpret_cast(&packed_p[14])), "r"(*reinterpret_cast(&packed_p[15]))); - asm volatile( - "tcgen05.st.sync.aligned.16x32bx2.x16.b32" - " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" - :: "r"(p_addr + 16), "r"(*reinterpret_cast(&(packed_p + 16)[0])), "r"(*reinterpret_cast(&(packed_p + 16)[1])), "r"(*reinterpret_cast(&(packed_p + 16)[2])), "r"(*reinterpret_cast(&(packed_p + 16)[3])), "r"(*reinterpret_cast(&(packed_p + 16)[4])), "r"(*reinterpret_cast(&(packed_p + 16)[5])), "r"(*reinterpret_cast(&(packed_p + 16)[6])), "r"(*reinterpret_cast(&(packed_p + 16)[7])), "r"(*reinterpret_cast(&(packed_p + 16)[8])), "r"(*reinterpret_cast(&(packed_p + 16)[9])), "r"(*reinterpret_cast(&(packed_p + 16)[10])), "r"(*reinterpret_cast(&(packed_p + 16)[11])), "r"(*reinterpret_cast(&(packed_p + 16)[12])), "r"(*reinterpret_cast(&(packed_p + 16)[13])), "r"(*reinterpret_cast(&(packed_p + 16)[14])), "r"(*reinterpret_cast(&(packed_p + 16)[15]))); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - mbarrier_arrive(p_full_addr); - if (col_half == 0) { - float _max_3 = max_noftz(row_sum, 1.1754943508222875e-38f); - float published_sum = _max_3; - smem_sum[my_row] = published_sum; - } - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(sum_ready_addr); - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(tmem_dealloc_addr); - } - } - // ---- Role: epilogue ---- - if (warp >= 4 && warp <= 7) { - asm volatile("setmaxnreg.dec.sync.aligned.u32 128;"); - { // epilogue_main - float output_scale = bmm2_scale[0]; - int work_idx = blockIdx.x; - int query_idx_1 = work_idx >> 2; - int v_chunk = work_idx & 3; - const int warp_in_role = warp - 4; - const int tmem_row_origin_1 = warp_in_role * 32; - const int logical_row_origin_1 = warp_in_role * 16; - const int my_row_1 = logical_row_origin_1 + lane % 16; - const int col_half_1 = lane / 16; - const int row_addr = tmem_row_origin_1 << 16; - unsigned int _phase_o_done_0 = 0; - mbarrier_wait(o_done_addr, _phase_o_done_0); - _phase_o_done_0 ^= 1; - unsigned int _phase_sum_ready_0 = 0; - mbarrier_wait(sum_ready_addr, _phase_sum_ready_0); - _phase_sum_ready_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - float row_sum_1 = smem_sum[my_row_1]; - float _rcp_0 = approx_rcp(row_sum_1); - float inv_sum = _rcp_0; - int output_base = (query_idx_1 * num_heads + my_row_1) * 512 + v_chunk * 128; - float _tmem_load_1[64]; - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_1[0])), "=r"(*reinterpret_cast(&_tmem_load_1[1])), "=r"(*reinterpret_cast(&_tmem_load_1[2])), "=r"(*reinterpret_cast(&_tmem_load_1[3])), "=r"(*reinterpret_cast(&_tmem_load_1[4])), "=r"(*reinterpret_cast(&_tmem_load_1[5])), "=r"(*reinterpret_cast(&_tmem_load_1[6])), "=r"(*reinterpret_cast(&_tmem_load_1[7])), "=r"(*reinterpret_cast(&_tmem_load_1[8])), "=r"(*reinterpret_cast(&_tmem_load_1[9])), "=r"(*reinterpret_cast(&_tmem_load_1[10])), "=r"(*reinterpret_cast(&_tmem_load_1[11])), "=r"(*reinterpret_cast(&_tmem_load_1[12])), "=r"(*reinterpret_cast(&_tmem_load_1[13])), "=r"(*reinterpret_cast(&_tmem_load_1[14])), "=r"(*reinterpret_cast(&_tmem_load_1[15])), "=r"(*reinterpret_cast(&_tmem_load_1[16])), "=r"(*reinterpret_cast(&_tmem_load_1[17])), "=r"(*reinterpret_cast(&_tmem_load_1[18])), "=r"(*reinterpret_cast(&_tmem_load_1[19])), "=r"(*reinterpret_cast(&_tmem_load_1[20])), "=r"(*reinterpret_cast(&_tmem_load_1[21])), "=r"(*reinterpret_cast(&_tmem_load_1[22])), "=r"(*reinterpret_cast(&_tmem_load_1[23])), "=r"(*reinterpret_cast(&_tmem_load_1[24])), "=r"(*reinterpret_cast(&_tmem_load_1[25])), "=r"(*reinterpret_cast(&_tmem_load_1[26])), "=r"(*reinterpret_cast(&_tmem_load_1[27])), "=r"(*reinterpret_cast(&_tmem_load_1[28])), "=r"(*reinterpret_cast(&_tmem_load_1[29])), "=r"(*reinterpret_cast(&_tmem_load_1[30])), "=r"(*reinterpret_cast(&_tmem_load_1[31])) - : "r"(taddr + 256 + (unsigned int)row_addr)); - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_1[32])), "=r"(*reinterpret_cast(&_tmem_load_1[33])), "=r"(*reinterpret_cast(&_tmem_load_1[34])), "=r"(*reinterpret_cast(&_tmem_load_1[35])), "=r"(*reinterpret_cast(&_tmem_load_1[36])), "=r"(*reinterpret_cast(&_tmem_load_1[37])), "=r"(*reinterpret_cast(&_tmem_load_1[38])), "=r"(*reinterpret_cast(&_tmem_load_1[39])), "=r"(*reinterpret_cast(&_tmem_load_1[40])), "=r"(*reinterpret_cast(&_tmem_load_1[41])), "=r"(*reinterpret_cast(&_tmem_load_1[42])), "=r"(*reinterpret_cast(&_tmem_load_1[43])), "=r"(*reinterpret_cast(&_tmem_load_1[44])), "=r"(*reinterpret_cast(&_tmem_load_1[45])), "=r"(*reinterpret_cast(&_tmem_load_1[46])), "=r"(*reinterpret_cast(&_tmem_load_1[47])), "=r"(*reinterpret_cast(&_tmem_load_1[48])), "=r"(*reinterpret_cast(&_tmem_load_1[49])), "=r"(*reinterpret_cast(&_tmem_load_1[50])), "=r"(*reinterpret_cast(&_tmem_load_1[51])), "=r"(*reinterpret_cast(&_tmem_load_1[52])), "=r"(*reinterpret_cast(&_tmem_load_1[53])), "=r"(*reinterpret_cast(&_tmem_load_1[54])), "=r"(*reinterpret_cast(&_tmem_load_1[55])), "=r"(*reinterpret_cast(&_tmem_load_1[56])), "=r"(*reinterpret_cast(&_tmem_load_1[57])), "=r"(*reinterpret_cast(&_tmem_load_1[58])), "=r"(*reinterpret_cast(&_tmem_load_1[59])), "=r"(*reinterpret_cast(&_tmem_load_1[60])), "=r"(*reinterpret_cast(&_tmem_load_1[61])), "=r"(*reinterpret_cast(&_tmem_load_1[62])), "=r"(*reinterpret_cast(&_tmem_load_1[63])) - : "r"(taddr + 256 + (unsigned int)row_addr + 32)); - int col_base = col_half_1 * 64; - if (my_row_1 < num_heads) { - #pragma unroll - for (int offset = 0; offset < 64; offset += 8) { - { - const float2 _prescale2_0 = {inv_sum * output_scale, inv_sum * output_scale}; - #if __CUDA_ARCH__ >= 1000 - #pragma unroll - for (int _ps = 0; _ps < 4; _ps++) - mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_1[offset])[_ps], _prescale2_0); - #else - #pragma unroll - for (int _ps = 0; _ps < 8; _ps++) - _tmem_load_1[offset + _ps] *= inv_sum * output_scale; - #endif - __nv_bfloat162 _pk[4]; - _pk[0] = __floats2bfloat162_rn(_tmem_load_1[offset + 0], _tmem_load_1[offset + 1]); - _pk[1] = __floats2bfloat162_rn(_tmem_load_1[offset + 2], _tmem_load_1[offset + 3]); - _pk[2] = __floats2bfloat162_rn(_tmem_load_1[offset + 4], _tmem_load_1[offset + 5]); - _pk[3] = __floats2bfloat162_rn(_tmem_load_1[offset + 6], _tmem_load_1[offset + 7]); - *reinterpret_cast(&((__nv_bfloat16*)(O + (output_base + col_base + offset)))[0]) = *reinterpret_cast(&_pk[0]); - } - } - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(tmem_dealloc_addr); - } - } - // ---- Role: mma_warp ---- - if (warp == 8) { - { // mma_warp_main - unsigned int _phase_q_full_0 = 0; - mbarrier_wait(q_full_addr, _phase_q_full_0); - _phase_q_full_0 ^= 1; - unsigned int _phase_kv_full = 0; - #pragma unroll - for (int k_stage = 0; k_stage < 4; k_stage++) { - mbarrier_wait(kv_full_addr + (k_stage) * 8, _phase_kv_full); - asm volatile("tcgen05.fence::after_thread_sync;"); - int _mma_a_lo_0 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (k_stage * 2) * 512); - int _mma_b_lo_0 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (k_stage) * 2048); - { - uint64_t _mma_ss_a_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_0); - uint64_t _mma_ss_b_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_0); - if (elect_sync()) { - tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, ((k_stage == 0) ? 0 : 1)); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); - } - } - int _mma_a_lo_1 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (k_stage * 2 + 1) * 512); - int _mma_b_lo_1 = make_warp_uniform((((smem_kv_addr + 16384) >> 4) & 0x3FFF) + (k_stage) * 2048); - { - uint64_t _mma_ss_a_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_1); - uint64_t _mma_ss_b_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_1); - if (elect_sync()) { - tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); - } - } - if (k_stage == 3) { - elect_commit(s_full_addr); - } - elect_commit(kv_empty_addr + (k_stage) * 8); - } - int _mma_prepared_b_lo_2 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (0) * 2048); - unsigned int _phase_p_full_0 = 0; - mbarrier_wait(p_full_addr, _phase_p_full_0); - _phase_p_full_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - mbarrier_wait(kv_full_addr, 1); - asm volatile("tcgen05.fence::after_thread_sync;"); - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 dhi, blo, ta, id;\n\t" - ".reg .b64 db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "" - "mov.b32 dhi, 0x40004040;\n\t" - "mov.b32 id, 69272720;\n\t" - "mov.b32 ta, %2;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p0;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "}\n" - :: "r"((tmem_tmem + (256))), "r"(_mma_prepared_b_lo_2), "r"(tmem_tmem + 128), "r"(0)); - elect_commit(o_done_addr); - elect_commit(kv_empty_addr); - unsigned int _phase_tmem_dealloc_0 = 0; - mbarrier_wait(tmem_dealloc_addr, _phase_tmem_dealloc_0); - _phase_tmem_dealloc_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); - asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); - } - } - // ---- Role: empty ---- - if (warp >= 9 && warp <= 11) { - // idle — no tasks assigned - } - // ---- Role: load_warp ---- - if (warp >= 12 && warp <= 15) { - asm volatile("setmaxnreg.dec.sync.aligned.u32 64;"); - { // load_warp_main - const int load_warp_rank = warp - 12; - int work_idx_1 = blockIdx.x; - int query_idx_2 = work_idx_1 >> 2; - int v_chunk_1 = work_idx_1 & 3; - int* row_ptr = swa_indices + (query_idx_2 * swa_index_stride); - int index_offset = load_warp_rank * 32 + lane; - int sparse_row = -1; - if (index_offset < sparse_topk) { - sparse_row = row_ptr[index_offset]; - } - if (load_warp_rank == 0) { - if (elect_sync()) { - mbarrier_arrive_expect_tx(q_full_addr, 65536); - #pragma unroll - for (int q_stage = 0; q_stage < 8; q_stage++) { - tma_4d_gmem2smem(smem_q_addr + (unsigned int)(q_stage * 8192), (&tmap_q), 0, 0, q_stage, query_idx_2, q_full_addr); - } - } - } - smem_indices[index_offset] = sparse_row; - unsigned int _vote_0 = __ballot_sync(0xFFFFFFFF, sparse_row < 0); - unsigned int invalid_indices = _vote_0; - int valid_prefix = 32; - if (invalid_indices != 0) { - int _ffs_0 = __ffs(invalid_indices); - valid_prefix = _ffs_0 - 1; - } - if (elect_sync()) { - smem_index_flags[load_warp_rank] = valid_prefix; - } - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - asm volatile("barrier.sync 8, 128;" ::: "memory"); - if ((load_warp_rank & 1) != 0) { - if (elect_sync()) { - int lower_prefix = smem_index_flags[load_warp_rank - 1]; - int upper_prefix = smem_index_flags[load_warp_rank]; - smem_index_flags[load_warp_rank] = ((lower_prefix < 32) ? lower_prefix : 32 + upper_prefix); - } - } - unsigned int _phase_kv_empty = 1; - #pragma unroll - for (int k_stage_1 = 0; k_stage_1 < 4; k_stage_1++) { - mbarrier_wait(kv_empty_addr + (k_stage_1) * 8, _phase_kv_empty); - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (k_stage_1) * 8, 8192); - } - #pragma unroll - for (int local_group = 0; local_group < 8; local_group++) { - int dst_k = smem_kv_addr + (unsigned int)(k_stage_1 * 32768); - int group = load_warp_rank * 8 + local_group; - int group_offset = group * 4; - int raw_rows[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows[0])), "=r"(*reinterpret_cast(&raw_rows[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows[(0) + 3])) - : "r"(smem_indices_addr + (unsigned int)(group_offset * 4))); - int raw0 = raw_rows[0]; - int raw1 = raw_rows[1]; - int raw2 = raw_rows[2]; - int raw3 = raw_rows[3]; - int row0 = ((raw0 >= 0) ? raw0 : 0); - int row1 = ((raw1 >= 0) ? raw1 : 0); - int row2 = ((raw2 >= 0) ? raw2 : 0); - int row3 = ((raw3 >= 0) ? raw3 : 0); - if (elect_sync()) { - tma_gather4_gmem2smem(dst_k + group * 512, (&tmap_swa_kv), k_stage_1 * 128, row0, row1, row2, row3, kv_full_addr + (k_stage_1) * 8); - tma_gather4_gmem2smem(dst_k + 16384 + group * 512, (&tmap_swa_kv), k_stage_1 * 128 + 64, row0, row1, row2, row3, kv_full_addr + (k_stage_1) * 8); - } - } - } - #pragma unroll - for (int v_stage = 0; v_stage < 1; v_stage++) { - mbarrier_wait(kv_empty_addr + (v_stage) * 8, 0); - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (v_stage) * 8, 8192); - } - #pragma unroll - for (int local_group_1 = 0; local_group_1 < 8; local_group_1++) { - int dst_v = smem_v_addr + (unsigned int)(v_stage * 32768); - int v_col = v_chunk_1 * 128 + v_stage * 64; - int group_1 = load_warp_rank * 8 + local_group_1; - int group_offset_1 = group_1 * 4; - int raw_rows_1[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_1[0])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 3])) - : "r"(smem_indices_addr + (unsigned int)(group_offset_1 * 4))); - int raw0_1 = raw_rows_1[0]; - int raw1_1 = raw_rows_1[1]; - int raw2_1 = raw_rows_1[2]; - int raw3_1 = raw_rows_1[3]; - int row0_1 = ((raw0_1 >= 0) ? raw0_1 : 0); - int row1_1 = ((raw1_1 >= 0) ? raw1_1 : 0); - int row2_1 = ((raw2_1 >= 0) ? raw2_1 : 0); - int row3_1 = ((raw3_1 >= 0) ? raw3_1 : 0); - if (elect_sync()) { - tma_gather4_gmem2smem(dst_v + group_1 * 512, (&tmap_swa_kv), v_col, row0_1, row1_1, row2_1, row3_1, kv_full_addr + (v_stage) * 8); - tma_gather4_gmem2smem(dst_v + 16384 + group_1 * 512, (&tmap_swa_kv), v_col + 64, row0_1, row1_1, row2_1, row3_1, kv_full_addr + (v_stage) * 8); - } - } - } - } - } - - // Cleanup -} - -} // extern "C" diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_57e7b93def70f2be356f_binding.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_57e7b93def70f2be356f_binding.cu deleted file mode 100644 index da00a1e33ac..00000000000 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_57e7b93def70f2be356f_binding.cu +++ /dev/null @@ -1,389 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -// Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_57e7b93def70f2be356f'. -#include -#include -#include -#include - -#include "tvm_ffi_utils.h" - -#include -#include -#include -#include -#include -#include -#include - -extern "C" __global__ void kernel_cake_dsv4_57e7b93def70f2be356f(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int num_head_tiles, int has_sinks); - - -namespace cake_host_shim_16cf3a7d0d01e13d { - -using tvm::ffi::Optional; -using tvm::ffi::TensorView; - -class ScopedCudaDevice { - public: - explicit ScopedCudaDevice(int device_id) { - cudaError_t error = cudaGetDevice(&previous_device_); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed before host-shim launch: cudaError=" - << static_cast(error); - if (previous_device_ != device_id) { - error = cudaSetDevice(device_id); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaSetDevice failed before host-shim launch for cuda:" - << device_id << ": cudaError=" << static_cast(error); - restore_ = true; - } - } - - ScopedCudaDevice(const ScopedCudaDevice&) = delete; - ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; - - ~ScopedCudaDevice() noexcept { - if (restore_) { - (void)cudaSetDevice(previous_device_); - } - } - - private: - int previous_device_ = -1; - bool restore_ = false; -}; - -inline int64_t CakeDeviceMultiprocessorCount(int device_id) { - constexpr int kMaxCachedCudaDevices = 64; - TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) - << "physical-SM-count cache does not cover cuda:" << device_id; - static std::atomic count_by_device[kMaxCachedCudaDevices]{}; - int cached = count_by_device[device_id].load(std::memory_order_acquire); - if (cached > 0) return cached; - - int count = 0; - cudaError_t error = cudaDeviceGetAttribute( - &count, cudaDevAttrMultiProcessorCount, device_id); - TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) - << "querying multiProcessorCount failed for cuda:" << device_id - << ": cudaError=" << static_cast(error) << ", count=" << count; - // Concurrent first launches may repeat the immutable device query, but all - // publication is atomic and every later hot-path lookup is one acquire load. - count_by_device[device_id].store(count, std::memory_order_release); - return count; -} - -inline void CheckCudaTensor(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) - << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; -} - -inline void CheckSameCudaDevice( - const TensorView& t, - const TensorView& reference, - const char* name, - const char* reference_name) { - TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) - << name << " must be on the same CUDA device as " << reference_name - << ": got cuda:" << t.device().device_id - << " versus cuda:" << reference.device().device_id; -} - -inline void CheckCurrentCudaDevice( - const TensorView& reference, - const char* reference_name) { - int current_device = -1; - cudaError_t error = cudaGetDevice(¤t_device); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed while validating " << reference_name - << ": cudaError=" << static_cast(error); - TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) - << "current CUDA device must match " << reference_name - << ": current=cuda:" << current_device - << ", tensor=cuda:" << reference.device().device_id; -} - -inline void CheckContiguous(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; -} - -inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { - DLDataType d = t.dtype(); - TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) - << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits - << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits - << ", lanes=" << (int)d.lanes << ")"; -} - -// A logical axis.outer(trailing) folds every source dim above the trailing -// dimensions. Shape products are independent of physical strides, so verify -// the leading dimensions form one dense row-major chain instead of inventing -// a "folded stride". The descriptor reads its exact adjacent physical step -// separately through stride[-(trailing + 1)]. -inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { - TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) - << name << " cannot fold leading dimensions above " << trailing - << " trailing dims from ndim=" << t.ndim(); - int outer_last = t.ndim() - trailing - 1; - if (outer_last <= 0) { - return; - } - int64_t step = t.stride(outer_last); - TVM_FFI_CHECK(step > 0, ValueError) - << name << " physical strides must be positive"; - int64_t expected = step; - for (int axis = outer_last - 1; axis >= 0; --axis) { - expected *= t.size(axis + 1); - if (t.size(axis) > 1) { - TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) - << name << " leading dims are not physically foldable above " << trailing - << " trailing dims: stride(" << axis << ")=" << t.stride(axis) - << ", expected " << expected; - } - } -} - -// 4D TMA descriptor for buffer 'tmap_q' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 3, ValueError) - << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - int64_t d2 = t.size(t.ndim() - 2); - int64_t d3 = t.size(t.ndim() - 3); - TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) - << "TMA source 'tmap_q' trailing dims must be positive"; - TVM_FFI_CHECK(d1 % 64 == 0, ValueError) - << "TMA source 'tmap_q' extent " << d1 - << " must divide exactly by " << 64; - uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(d3)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) - << "TMA box (64, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; - int64_t carrier_stride_0 = d1; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; - int64_t carrier_stride_1 = 64; - TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; - TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; - TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; - int64_t carrier_stride_2 = (d2 * d1); - TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; - TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; - TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; - uint64_t global_strides[3] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - (uint64_t)((carrier_stride_1 * 16) / 8), - (uint64_t)((carrier_stride_2 * 16) / 8), - }; - uint32_t box_dim[4] = {64u, 64u, 1u, 1u}; - uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_swa_kv' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_swa_kv' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_kv'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; - return tm; -} - -void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_O, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_num_head_tiles, int64_t arg_has_sinks, int64_t grid_x, int64_t grid_y, int64_t grid_z) { - DLDevice dev = arg_tmap_q.device(); - ScopedCudaDevice device_guard(dev.device_id); - CheckCudaTensor(arg_tmap_q, "tmap_q"); - CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); - CheckContiguous(arg_tmap_q, "tmap_q"); - CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); - CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 4, 16, 1); - CheckCudaTensor(arg_O, "O"); - CheckDtype(arg_O, "O", 4, 16, 1); - CheckContiguous(arg_O, "O"); - CheckCudaTensor(arg_swa_indices, "swa_indices"); - CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); - CheckContiguous(arg_swa_indices, "swa_indices"); - CheckCudaTensor(arg_compressed_indices, "compressed_indices"); - CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); - CheckContiguous(arg_compressed_indices, "compressed_indices"); - CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); - CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckCudaTensor(arg_sinks, "sinks"); - CheckDtype(arg_sinks, "sinks", 2, 32, 1); - CheckContiguous(arg_sinks, "sinks"); - CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); - CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); - CheckContiguous(arg_bmm1_scale, "bmm1_scale"); - CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); - CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); - CheckContiguous(arg_bmm2_scale, "bmm2_scale"); - TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) - << "scalar 'num_heads' value " << arg_num_heads - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) - << "scalar 'swa_index_stride' value " << arg_swa_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) - << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) - << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) - << "scalar 'sparse_topk' value " << arg_sparse_topk - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_num_head_tiles >= -2147483648LL && arg_num_head_tiles <= 2147483647LL, ValueError) - << "scalar 'num_head_tiles' value " << arg_num_head_tiles - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) - << "scalar 'has_sinks' value " << arg_has_sinks - << " is outside i32 range [-2147483648, 2147483647]"; - CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); - CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); - CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); - CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); - CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); - CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); - CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); - CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); - TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) - << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y - << ", " << grid_z << ")"; - - cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); - CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); - CUtensorMap p_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); - int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); - int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); - int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); - float* p_sinks = static_cast(arg_sinks.data_ptr()); - float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); - float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); - int32_t v_num_heads = (int32_t)arg_num_heads; - int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; - int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; - int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; - int32_t v_sparse_topk = (int32_t)arg_sparse_topk; - int32_t v_num_head_tiles = (int32_t)arg_num_head_tiles; - int32_t v_has_sinks = (int32_t)arg_has_sinks; - void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_O, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_num_head_tiles, &v_has_sinks}; - - static std::mutex smem_mu; - static auto* smem_status_by_device = - new std::unordered_map(); - cudaError_t smem_status = cudaSuccess; - { - std::lock_guard lock(smem_mu); - auto it = smem_status_by_device->find(dev.device_id); - if (it == smem_status_by_device->end()) { - int current_device = -1; - smem_status = cudaGetDevice(¤t_device); - if (smem_status == cudaSuccess && current_device != dev.device_id) { - smem_status = cudaErrorInvalidDevice; - } - if (smem_status == cudaSuccess) { - smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_57e7b93def70f2be356f), - cudaFuncAttributeMaxDynamicSharedMemorySize, - 198528); - } - smem_status_by_device->emplace(dev.device_id, smem_status); - } else { - smem_status = it->second; - } - } - TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_57e7b93def70f2be356f failed: " - << cudaGetErrorString(smem_status); - dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); - dim3 block(512u, 1u, 1u); - - cudaLaunchConfig_t config{}; - config.gridDim = grid; - config.blockDim = block; - config.dynamicSmemBytes = 198528u; - config.stream = stream; - cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_57e7b93def70f2be356f), kargs); - TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_57e7b93def70f2be356f failed: " - << cudaGetErrorString(launch_status); -} - -} // namespace cake_host_shim_16cf3a7d0d01e13d - -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_16cf3a7d0d01e13d::Run); diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_57e7b93def70f2be356f_kernel.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_57e7b93def70f2be356f_kernel.cu deleted file mode 100644 index bb28b1bf0d1..00000000000 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_57e7b93def70f2be356f_kernel.cu +++ /dev/null @@ -1,1643 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -typedef signed char int8_t; -typedef unsigned char uint8_t; -typedef unsigned short uint16_t; -typedef unsigned int uint32_t; -#if defined(__CUDACC_RTC__) -typedef unsigned long long uint64_t; -#else -typedef unsigned long uint64_t; -#endif -static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); -typedef signed int int32_t; -typedef short int int16_t; -struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; -struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; -static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); -static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); - -#if defined(__CUDACC_RTC__) -typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; -#else -#include -#endif - -static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); -static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); -#include -#include - -#define CAKE_INF CUDART_INF_F -#define TMEM_NCOLS 384 -#define TMEM_TMEM_OFFSET 0 -#define NUM_KV_PIPE_STAGES 4 -#define SMEM_SMEM_Q_OFF 1024 -#define SMEM_SMEM_Q_STAGE_BYTES 8192 -#define SMEM_SMEM_Q_STRIDE 8192 -#define SMEM_SMEM_KV_OFF 66560 -#define SMEM_SMEM_KV_STAGE_BYTES 32768 -#define SMEM_SMEM_KV_STRIDE 32768 -#define SMEM_SMEM_V_OFF 66560 -#define SMEM_SMEM_V_STAGE_BYTES 32768 -#define SMEM_SMEM_V_STRIDE 32768 -#define SMEM_SMEM_SUM_OFF 197632 -#define SMEM_SMEM_SUM_STAGE_BYTES 256 -#define SMEM_SMEM_SUM_STRIDE 256 -#define SMEM_SMEM_INDEX_FLAGS_OFF 198400 -#define SMEM_SMEM_INDEX_FLAGS_STAGE_BYTES 16 -#define SMEM_SMEM_INDEX_FLAGS_STRIDE 16 -#define SMEM_SMEM_INDICES_OFF 197888 -#define SMEM_SMEM_INDICES_STAGE_BYTES 512 -#define SMEM_SMEM_INDICES_STRIDE 512 -#define SMEM_TOTAL 198528 -#define THREADS 512 - -#include - -__device__ __forceinline__ uint32_t elect_sync() { - uint32_t pred = 0; - asm volatile( - "{\n\t" - ".reg .pred %%px;\n\t" - "elect.sync _|%%px, %1;\n\t" - "@%%px mov.s32 %0, 1;\n\t" - "}\n" - : "+r"(pred) - : "r"(0xFFFFFFFF)); - return pred; -} - - -__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { - asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" - :: "r"(mbar_addr), "r"(count) : "memory"); -} - -__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { - asm volatile("mbarrier.init.b64 [%0], %1;" - :: "l"(mbar_addr), "r"(count) : "memory"); -} - - -__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - - -// CTA-local pipelines have short, resident producer/consumer edges. Omitting -// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the -// explicit loop still makes this helper blocking until acquire succeeds. -__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE;\n\t" - "bra.uni LAB_WAIT;\n\t" - "DONE:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Source-faithful relaxed CTA wait used only by a typed protocol that does -// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. -__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, 10000000;\n\t" - "@P1 bra.uni DONE_RELAXED;\n\t" - "bra.uni LAB_WAIT_RELAXED;\n\t" - "DONE_RELAXED:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Exact source ports may request the PTX suspendTimeHint operand explicitly. -// The hint is expressed in nanoseconds and is kept separate from the canonical -// no-hint CTA helper so unrelated schedules retain their existing retry path. -__device__ __forceinline__ void mbarrier_wait_suspend( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_SUSPEND:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_SUSPEND;\n\t" - "bra.uni LAB_WAIT_SUSPEND;\n\t" - "DONE_SUSPEND:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE_CLUSTER;\n\t" - "bra.uni LAB_WAIT_CLUSTER;\n\t" - "DONE_CLUSTER:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - ".reg .u32 WAIT_ADDR;\n\t" - "mov.u32 WAIT_ADDR, %0;\n\t" - "LAB_WAIT_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [WAIT_ADDR], %1, %2;\n\t" - "@P1 bra.uni DONE_HINT;\n\t" - "bra.uni LAB_WAIT_HINT;\n\t" - "DONE_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -// Exact unqualified CTA wait used by source schedules whose PTX intentionally -// omits the acquire qualifier while retaining a typed suspendTimeHint operand. -__device__ __forceinline__ void mbarrier_wait_relaxed_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED_HINT:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra DONE_RELAXED_HINT;\n\t" - "bra LAB_WAIT_RELAXED_HINT;\n\t" - "DONE_RELAXED_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); -} - -__device__ __forceinline__ void mbarrier_wait_cluster_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" - "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" - "DONE_CLUSTER_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_suspend( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait_cluster(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); - } -} - - -__device__ __forceinline__ void tcgen05_mma_f16( - int taddr, uint64_t a_desc, uint64_t b_desc, - uint32_t i_desc, int enable_input_d) { - asm volatile( - "{\n\t" - ".reg .pred p;\n\t" - "setp.ne.b32 p, %4, 0;\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" - "}\n" - :: "r"(taddr), "l"(a_desc), "l"(b_desc), - "r"(i_desc), "r"(enable_input_d) - : "memory"); -} - - -__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { - return (x & 0x3FFFFULL) >> 4ULL; -} - - -union MmaSmemDesc { - uint64_t u64; - uint32_t u32[2]; -}; - -__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { - MmaSmemDesc tmp; - tmp.u64 = smem_desc; - tmp.u32[0] += offset; - smem_desc = tmp.u64; -} - - -__device__ __forceinline__ void elect_commit(int mbar_addr) { - asm volatile( - "{\n\t" - ".reg .pred leader;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" - ".shared::cluster.b64 [%0];\n\t" - "}\n" - :: "r"(mbar_addr)); -} - - -__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { - asm volatile( - "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" - :: "r"(mbar_addr), "r"(bytes) : "memory"); -} - - -__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7," - " %8, %9, %10, %11, %12, %13, %14, %15," - " %16, %17, %18, %19, %20, %21, %22, %23," - " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), - "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), - "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), - "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), - "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), - "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), - "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ float approx_exp2(float x) { - float y; - asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float approx_rcp(float x) { - float y; - asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float max_noftz(float a, float b) { - float c; - asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); - return c; -} - - -__device__ __forceinline__ float warp_reduce_max(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); - return val; -} - - -__device__ __forceinline__ float warp_reduce_sum(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val += __shfl_xor_sync(0xFFFFFFFF, val, offset); - return val; -} - - -__device__ __forceinline__ float row_max_reduce(float2 acc) { - return max_noftz(acc.x, acc.y); -} - - -__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { - #pragma unroll - for (int j = 0; j < 16; j++) { - if (j % 2 == 0) - acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); - else - acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); - } -} - - -__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { - asm("mul.rn.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { - asm("mul.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { - asm("add.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { - asm("add.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { - asm("sub.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { - asm("sub.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("add.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { - float2 r; - asm("sub.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("sub.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ void fma_scale_x32( - float* sv, const float2* scale2, const float2* neg_max2) -{ - float2* sv_2 = reinterpret_cast(sv); - #pragma unroll - for (int j = 0; j < 16; j++) - fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); -} - -__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("mul.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - - -// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) - -__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("add.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("add.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("add.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("add.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("add.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("add.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("add.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - - -__device__ __forceinline__ void fence_async_shared() { - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); -} - - -__device__ __forceinline__ uint64_t make_smem_desc(int addr) { - const int SBO = 1024; - return desc_encode(addr) - | (desc_encode(SBO) << 32ULL) - | (1ULL << 46ULL) - | (2ULL << 61ULL); -} - - -__device__ __forceinline__ void tma_2d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3}], [%4];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_4d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.4d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5}], [%6];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr) { - // Canonical .shared::cta form for non-multicast gather4, matching - // trtllm-gen / cuda_ptx and the PTX ISA qualifier order - // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, - // .shared::cluster is reserved for the multicast variant (ctaMask). - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { - asm volatile( - "tcgen05.commit.cta_group::1.mbarrier::arrive::one" - ".shared::cluster.b64 [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { - uint32_t result; - asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" - : "=r"(result) : "r"(val)); - return result; -} - -extern "C" { - -__global__ __launch_bounds__(512, 1) void -kernel_cake_dsv4_57e7b93def70f2be356f(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int num_head_tiles, int has_sinks) -{ - const int tid = threadIdx.x; - const int warp = make_warp_uniform(tid / 32); - const int lane = tid % 32; - - extern __shared__ __align__(1024) char smem_raw[]; - int smem; - asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); - smem = make_warp_uniform(smem); - - const int mbar_base = smem; - #define q_full_addr (mbar_base + 0) - #define kv_full_addr (mbar_base + 8) - #define kv_empty_addr (mbar_base + 40) - #define s_full_addr (mbar_base + 72) - #define p_full_addr (mbar_base + 80) - #define sum_ready_addr (mbar_base + 88) - #define o_done_addr (mbar_base + 96) - #define tmem_dealloc_addr (mbar_base + 104) - - const int bid = blockIdx.x; - const int num_bids = gridDim.x; - - const int cta_rank = 0; - - // Kernel setup ops - __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); - const int smem_q_addr = smem + 1024; - __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); - const int smem_kv_addr = smem + 66560; - __nv_bfloat16* smem_v = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); - const int smem_v_addr = smem + 66560; - float* smem_sum = reinterpret_cast(smem_raw + 197632); - const int smem_sum_addr = smem + 197632; - int* smem_index_flags = reinterpret_cast(smem_raw + 198400); - const int smem_index_flags_addr = smem + 198400; - int* smem_indices = reinterpret_cast(smem_raw + 197888); - const int smem_indices_addr = smem + 197888; - - // Mbarrier init (8 pipeline groups, 0 ordered-sequence groups, 14 barriers) - // Mbarriers at smem_raw[0..112) - - if (warp == 0) { - uint32_t leader = elect_sync(); - if (leader) { - // q_full: 1 barriers, init_count=1 - mbarrier_init(smem + 0, 1); - // --- pipeline 'kv_pipe' --- - // kv_full: 4 barriers, init_count=4 - mbarrier_init(smem + 8, 4); - mbarrier_init(smem + 16, 4); - mbarrier_init(smem + 24, 4); - mbarrier_init(smem + 32, 4); - // kv_empty: 4 barriers, init_count=1 - mbarrier_init(smem + 40, 1); - mbarrier_init(smem + 48, 1); - mbarrier_init(smem + 56, 1); - mbarrier_init(smem + 64, 1); - // s_full: 1 barriers, init_count=1 - mbarrier_init(smem + 72, 1); - // p_full: 1 barriers, init_count=128 - mbarrier_init(smem + 80, 128); - // sum_ready: 1 barriers, init_count=128 - mbarrier_init(smem + 88, 128); - // o_done: 1 barriers, init_count=1 - mbarrier_init(smem + 96, 1); - // tmem_dealloc: 1 barriers, init_count=256 - mbarrier_init(smem + 104, 256); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - - __syncwarp(); - - // TMEM alloc (512 columns, 384 used) - volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 112); - if (warp == 0) { - int _tmem_hold = smem + 112; - asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); - __syncwarp(); - asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); - } - - __syncthreads(); - asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); - - const int taddr = tmem_addr_storage[0]; - - // Kernel post-init ops - const int tmem_tmem = taddr; - - // ---- Ordered hardware-WG register redistribution ---- - // Dec phase frees registers before any WG attempts inc. - if (warp >= 8 && warp <= 11) { - asm volatile("setmaxnreg.dec.sync.aligned.u32 64;"); - } - - // ---- Role: softmax ---- - if (warp <= 3) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); - { // softmax_main - float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; - int query_head_work = blockIdx.x >> 2; - int query_idx = query_head_work / num_head_tiles; - int head_tile = query_head_work % num_head_tiles; - int _max_0 = ((sparse_topk_lens[query_idx] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx] + sparse_topk_lens_offset) : (0)); - int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); - int active_topk = _min_0; - const int warp_in_compute = warp; - const int tmem_row_origin = warp_in_compute * 32; - const int logical_row_origin = warp_in_compute * 16; - const int my_row = logical_row_origin + lane % 16; - int head_idx = head_tile * 64 + my_row; - const int col_half = lane / 16; - unsigned int _phase_s_full_0 = 0; - mbarrier_wait(s_full_addr, _phase_s_full_0); - _phase_s_full_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - int score_addr = taddr + (unsigned int)(tmem_row_origin << 16); - float _tmem_load_0[64]; - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_0[0])), "=r"(*reinterpret_cast(&_tmem_load_0[1])), "=r"(*reinterpret_cast(&_tmem_load_0[2])), "=r"(*reinterpret_cast(&_tmem_load_0[3])), "=r"(*reinterpret_cast(&_tmem_load_0[4])), "=r"(*reinterpret_cast(&_tmem_load_0[5])), "=r"(*reinterpret_cast(&_tmem_load_0[6])), "=r"(*reinterpret_cast(&_tmem_load_0[7])), "=r"(*reinterpret_cast(&_tmem_load_0[8])), "=r"(*reinterpret_cast(&_tmem_load_0[9])), "=r"(*reinterpret_cast(&_tmem_load_0[10])), "=r"(*reinterpret_cast(&_tmem_load_0[11])), "=r"(*reinterpret_cast(&_tmem_load_0[12])), "=r"(*reinterpret_cast(&_tmem_load_0[13])), "=r"(*reinterpret_cast(&_tmem_load_0[14])), "=r"(*reinterpret_cast(&_tmem_load_0[15])), "=r"(*reinterpret_cast(&_tmem_load_0[16])), "=r"(*reinterpret_cast(&_tmem_load_0[17])), "=r"(*reinterpret_cast(&_tmem_load_0[18])), "=r"(*reinterpret_cast(&_tmem_load_0[19])), "=r"(*reinterpret_cast(&_tmem_load_0[20])), "=r"(*reinterpret_cast(&_tmem_load_0[21])), "=r"(*reinterpret_cast(&_tmem_load_0[22])), "=r"(*reinterpret_cast(&_tmem_load_0[23])), "=r"(*reinterpret_cast(&_tmem_load_0[24])), "=r"(*reinterpret_cast(&_tmem_load_0[25])), "=r"(*reinterpret_cast(&_tmem_load_0[26])), "=r"(*reinterpret_cast(&_tmem_load_0[27])), "=r"(*reinterpret_cast(&_tmem_load_0[28])), "=r"(*reinterpret_cast(&_tmem_load_0[29])), "=r"(*reinterpret_cast(&_tmem_load_0[30])), "=r"(*reinterpret_cast(&_tmem_load_0[31])) - : "r"(score_addr)); - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_0[32])), "=r"(*reinterpret_cast(&_tmem_load_0[33])), "=r"(*reinterpret_cast(&_tmem_load_0[34])), "=r"(*reinterpret_cast(&_tmem_load_0[35])), "=r"(*reinterpret_cast(&_tmem_load_0[36])), "=r"(*reinterpret_cast(&_tmem_load_0[37])), "=r"(*reinterpret_cast(&_tmem_load_0[38])), "=r"(*reinterpret_cast(&_tmem_load_0[39])), "=r"(*reinterpret_cast(&_tmem_load_0[40])), "=r"(*reinterpret_cast(&_tmem_load_0[41])), "=r"(*reinterpret_cast(&_tmem_load_0[42])), "=r"(*reinterpret_cast(&_tmem_load_0[43])), "=r"(*reinterpret_cast(&_tmem_load_0[44])), "=r"(*reinterpret_cast(&_tmem_load_0[45])), "=r"(*reinterpret_cast(&_tmem_load_0[46])), "=r"(*reinterpret_cast(&_tmem_load_0[47])), "=r"(*reinterpret_cast(&_tmem_load_0[48])), "=r"(*reinterpret_cast(&_tmem_load_0[49])), "=r"(*reinterpret_cast(&_tmem_load_0[50])), "=r"(*reinterpret_cast(&_tmem_load_0[51])), "=r"(*reinterpret_cast(&_tmem_load_0[52])), "=r"(*reinterpret_cast(&_tmem_load_0[53])), "=r"(*reinterpret_cast(&_tmem_load_0[54])), "=r"(*reinterpret_cast(&_tmem_load_0[55])), "=r"(*reinterpret_cast(&_tmem_load_0[56])), "=r"(*reinterpret_cast(&_tmem_load_0[57])), "=r"(*reinterpret_cast(&_tmem_load_0[58])), "=r"(*reinterpret_cast(&_tmem_load_0[59])), "=r"(*reinterpret_cast(&_tmem_load_0[60])), "=r"(*reinterpret_cast(&_tmem_load_0[61])), "=r"(*reinterpret_cast(&_tmem_load_0[62])), "=r"(*reinterpret_cast(&_tmem_load_0[63])) - : "r"(score_addr + 32)); - int valid_cols = active_topk - col_half * 64; - int causal_cols = smem_index_flags[col_half * 2 + 1]; - valid_cols = ((valid_cols < causal_cols) ? valid_cols : causal_cols); - uint32_t _slice_lo_mask_0; - { - int _lim_0 = valid_cols; - if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } - else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); - } - } - if (!(_slice_lo_mask_0 & (1u << 0))) _tmem_load_0[0] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 1))) _tmem_load_0[1] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 2))) _tmem_load_0[2] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 3))) _tmem_load_0[3] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 4))) _tmem_load_0[4] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 5))) _tmem_load_0[5] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 6))) _tmem_load_0[6] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 7))) _tmem_load_0[7] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 8))) _tmem_load_0[8] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 9))) _tmem_load_0[9] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 10))) _tmem_load_0[10] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 11))) _tmem_load_0[11] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 12))) _tmem_load_0[12] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 13))) _tmem_load_0[13] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 14))) _tmem_load_0[14] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 15))) _tmem_load_0[15] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 16))) _tmem_load_0[16] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 17))) _tmem_load_0[17] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 18))) _tmem_load_0[18] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 19))) _tmem_load_0[19] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 20))) _tmem_load_0[20] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 21))) _tmem_load_0[21] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 22))) _tmem_load_0[22] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 23))) _tmem_load_0[23] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 24))) _tmem_load_0[24] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 25))) _tmem_load_0[25] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 26))) _tmem_load_0[26] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 27))) _tmem_load_0[27] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 28))) _tmem_load_0[28] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 29))) _tmem_load_0[29] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 30))) _tmem_load_0[30] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 31))) _tmem_load_0[31] = -CAKE_INF; - uint32_t _slice_lo_mask_1; - { - int _lim_1 = valid_cols - 32; - if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } - else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); - } - } - if (!(_slice_lo_mask_1 & (1u << 0))) _tmem_load_0[32] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 1))) _tmem_load_0[33] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 2))) _tmem_load_0[34] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 3))) _tmem_load_0[35] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 4))) _tmem_load_0[36] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 5))) _tmem_load_0[37] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 6))) _tmem_load_0[38] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 7))) _tmem_load_0[39] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 8))) _tmem_load_0[40] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 9))) _tmem_load_0[41] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 10))) _tmem_load_0[42] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 11))) _tmem_load_0[43] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 12))) _tmem_load_0[44] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 13))) _tmem_load_0[45] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 14))) _tmem_load_0[46] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 15))) _tmem_load_0[47] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 16))) _tmem_load_0[48] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 17))) _tmem_load_0[49] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 18))) _tmem_load_0[50] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 19))) _tmem_load_0[51] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 20))) _tmem_load_0[52] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 21))) _tmem_load_0[53] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 22))) _tmem_load_0[54] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 23))) _tmem_load_0[55] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 24))) _tmem_load_0[56] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 25))) _tmem_load_0[57] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 26))) _tmem_load_0[58] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 27))) _tmem_load_0[59] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 28))) _tmem_load_0[60] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 29))) _tmem_load_0[61] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 30))) _tmem_load_0[62] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 31))) _tmem_load_0[63] = -CAKE_INF; - float2 _reg_reduce_max2_2 = {-CAKE_INF, -CAKE_INF}; - row_max_x32_accum(&_tmem_load_0[0], _reg_reduce_max2_2); - row_max_x32_accum(&_tmem_load_0[32], _reg_reduce_max2_2); - float _tmem_load_0_max = row_max_reduce(_reg_reduce_max2_2); - float row_max = _tmem_load_0_max; - float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, row_max, 16); - float _max_1 = max_noftz(row_max, _shfl_xor_0); - row_max = _max_1; - if (has_sinks != 0 && head_idx < num_heads) { - float sink_unscaled = sinks[head_idx] * 1.4426950408889634f / softmax_scale_log2; - float _max_2 = max_noftz(row_max, sink_unscaled); - row_max = _max_2; - } - bool empty_row = row_max == -CAKE_INF; - float safe_max = ((empty_row) ? 0.0f : row_max); - float max_scaled = safe_max * softmax_scale_log2; - float2 _f2_0 = make_float2(softmax_scale_log2, softmax_scale_log2); - float2 _f2_1 = make_float2(-max_scaled, -max_scaled); - float2 _f2_2 = make_float2(0.0f, 0.0f); - float2 sum0 = _f2_2; - float2 _f2_3 = make_float2(0.0f, 0.0f); - float2 sum1 = _f2_3; - float2 _f2_4 = make_float2(0.0f, 0.0f); - float2 sum2 = _f2_4; - float2 _f2_5 = make_float2(0.0f, 0.0f); - float2 sum3 = _f2_5; - #pragma unroll - for (int i = 0; i < 64; i += 8) { - float2 _f2_6 = make_float2(_tmem_load_0[i], _tmem_load_0[i + 1]); - float2 _f2_7 = make_float2(_tmem_load_0[i + 2], _tmem_load_0[i + 3]); - float2 _f2_8 = make_float2(_tmem_load_0[i + 4], _tmem_load_0[i + 5]); - float2 _f2_9 = make_float2(_tmem_load_0[i + 6], _tmem_load_0[i + 7]); - float2 affine0 = fma_f32x2_rn_ftz(_f2_6, _f2_0, _f2_1); - float2 affine1 = fma_f32x2_rn_ftz(_f2_7, _f2_0, _f2_1); - float2 affine2 = fma_f32x2_rn_ftz(_f2_8, _f2_0, _f2_1); - float2 affine3 = fma_f32x2_rn_ftz(_f2_9, _f2_0, _f2_1); - float _exp2_0 = approx_exp2(affine0.x); - float exp0 = _exp2_0; - float _exp2_1 = approx_exp2(affine0.y); - float exp1 = _exp2_1; - float _exp2_2 = approx_exp2(affine1.x); - float exp2 = _exp2_2; - float _exp2_3 = approx_exp2(affine1.y); - float exp3 = _exp2_3; - float _exp2_4 = approx_exp2(affine2.x); - float exp4 = _exp2_4; - float _exp2_5 = approx_exp2(affine2.y); - float exp5 = _exp2_5; - float _exp2_6 = approx_exp2(affine3.x); - float exp6 = _exp2_6; - float _exp2_7 = approx_exp2(affine3.y); - float exp7 = _exp2_7; - _tmem_load_0[i] = exp0; - _tmem_load_0[i + 1] = exp1; - _tmem_load_0[i + 2] = exp2; - _tmem_load_0[i + 3] = exp3; - _tmem_load_0[i + 4] = exp4; - _tmem_load_0[i + 5] = exp5; - _tmem_load_0[i + 6] = exp6; - _tmem_load_0[i + 7] = exp7; - float2 _f2_10 = make_float2(exp0, exp1); - sum0 = add_f32x2(sum0, _f2_10); - float2 _f2_11 = make_float2(exp2, exp3); - sum1 = add_f32x2(sum1, _f2_11); - float2 _f2_12 = make_float2(exp4, exp5); - sum2 = add_f32x2(sum2, _f2_12); - float2 _f2_13 = make_float2(exp6, exp7); - sum3 = add_f32x2(sum3, _f2_13); - } - float2 sum01 = add_f32x2(sum0, sum1); - float2 sum23 = add_f32x2(sum2, sum3); - float2 row_sum_pair = add_f32x2(sum01, sum23); - float row_sum = row_sum_pair.x + row_sum_pair.y; - float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, row_sum, 16); - row_sum = row_sum + _shfl_xor_1; - if (has_sinks != 0 && head_idx < num_heads) { - float _exp2_8 = approx_exp2(sinks[head_idx] * 1.4426950408889634f - max_scaled); - row_sum = row_sum + _exp2_8; - } - unsigned int packed_p[32]; - #pragma unroll - for (int _lp = 0; _lp < 32; _lp++) { - __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_0[_lp*2 + 0], _tmem_load_0[_lp*2+1 + 0])); - packed_p[_lp] = *(uint32_t*)&_bf2; - } - int p_addr = taddr + 128 + (unsigned int)(tmem_row_origin << 16); - asm volatile( - "tcgen05.st.sync.aligned.16x32bx2.x16.b32" - " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" - :: "r"(p_addr), "r"(*reinterpret_cast(&packed_p[0])), "r"(*reinterpret_cast(&packed_p[1])), "r"(*reinterpret_cast(&packed_p[2])), "r"(*reinterpret_cast(&packed_p[3])), "r"(*reinterpret_cast(&packed_p[4])), "r"(*reinterpret_cast(&packed_p[5])), "r"(*reinterpret_cast(&packed_p[6])), "r"(*reinterpret_cast(&packed_p[7])), "r"(*reinterpret_cast(&packed_p[8])), "r"(*reinterpret_cast(&packed_p[9])), "r"(*reinterpret_cast(&packed_p[10])), "r"(*reinterpret_cast(&packed_p[11])), "r"(*reinterpret_cast(&packed_p[12])), "r"(*reinterpret_cast(&packed_p[13])), "r"(*reinterpret_cast(&packed_p[14])), "r"(*reinterpret_cast(&packed_p[15]))); - asm volatile( - "tcgen05.st.sync.aligned.16x32bx2.x16.b32" - " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" - :: "r"(p_addr + 16), "r"(*reinterpret_cast(&(packed_p + 16)[0])), "r"(*reinterpret_cast(&(packed_p + 16)[1])), "r"(*reinterpret_cast(&(packed_p + 16)[2])), "r"(*reinterpret_cast(&(packed_p + 16)[3])), "r"(*reinterpret_cast(&(packed_p + 16)[4])), "r"(*reinterpret_cast(&(packed_p + 16)[5])), "r"(*reinterpret_cast(&(packed_p + 16)[6])), "r"(*reinterpret_cast(&(packed_p + 16)[7])), "r"(*reinterpret_cast(&(packed_p + 16)[8])), "r"(*reinterpret_cast(&(packed_p + 16)[9])), "r"(*reinterpret_cast(&(packed_p + 16)[10])), "r"(*reinterpret_cast(&(packed_p + 16)[11])), "r"(*reinterpret_cast(&(packed_p + 16)[12])), "r"(*reinterpret_cast(&(packed_p + 16)[13])), "r"(*reinterpret_cast(&(packed_p + 16)[14])), "r"(*reinterpret_cast(&(packed_p + 16)[15]))); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - if (col_half == 0) { - float _max_3 = max_noftz(row_sum, 1.1754943508222875e-38f); - float published_sum = _max_3; - smem_sum[my_row] = published_sum; - } - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(p_full_addr); - mbarrier_arrive(sum_ready_addr); - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(tmem_dealloc_addr); - } - } - // ---- Role: epilogue ---- - if (warp >= 4 && warp <= 7) { - asm volatile("setmaxnreg.dec.sync.aligned.u32 128;"); - { // epilogue_main - float output_scale = bmm2_scale[0]; - int work_idx = blockIdx.x; - int query_head_work_1 = work_idx >> 2; - int query_idx_1 = query_head_work_1 / num_head_tiles; - int head_tile_1 = query_head_work_1 % num_head_tiles; - int v_chunk = work_idx & 3; - const int warp_in_role = warp - 4; - const int tmem_row_origin_1 = warp_in_role * 32; - const int logical_row_origin_1 = warp_in_role * 16; - const int my_row_1 = logical_row_origin_1 + lane % 16; - int head_idx_1 = head_tile_1 * 64 + my_row_1; - const int col_half_1 = lane / 16; - const int row_addr = tmem_row_origin_1 << 16; - unsigned int _phase_o_done_0 = 0; - mbarrier_wait(o_done_addr, _phase_o_done_0); - _phase_o_done_0 ^= 1; - unsigned int _phase_sum_ready_0 = 0; - mbarrier_wait(sum_ready_addr, _phase_sum_ready_0); - _phase_sum_ready_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - float row_sum_1 = smem_sum[my_row_1]; - float _rcp_0 = approx_rcp(row_sum_1); - float inv_sum = _rcp_0; - int output_base = (query_idx_1 * num_heads + head_idx_1) * 512 + v_chunk * 128; - float _tmem_load_1[64]; - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_1[0])), "=r"(*reinterpret_cast(&_tmem_load_1[1])), "=r"(*reinterpret_cast(&_tmem_load_1[2])), "=r"(*reinterpret_cast(&_tmem_load_1[3])), "=r"(*reinterpret_cast(&_tmem_load_1[4])), "=r"(*reinterpret_cast(&_tmem_load_1[5])), "=r"(*reinterpret_cast(&_tmem_load_1[6])), "=r"(*reinterpret_cast(&_tmem_load_1[7])), "=r"(*reinterpret_cast(&_tmem_load_1[8])), "=r"(*reinterpret_cast(&_tmem_load_1[9])), "=r"(*reinterpret_cast(&_tmem_load_1[10])), "=r"(*reinterpret_cast(&_tmem_load_1[11])), "=r"(*reinterpret_cast(&_tmem_load_1[12])), "=r"(*reinterpret_cast(&_tmem_load_1[13])), "=r"(*reinterpret_cast(&_tmem_load_1[14])), "=r"(*reinterpret_cast(&_tmem_load_1[15])), "=r"(*reinterpret_cast(&_tmem_load_1[16])), "=r"(*reinterpret_cast(&_tmem_load_1[17])), "=r"(*reinterpret_cast(&_tmem_load_1[18])), "=r"(*reinterpret_cast(&_tmem_load_1[19])), "=r"(*reinterpret_cast(&_tmem_load_1[20])), "=r"(*reinterpret_cast(&_tmem_load_1[21])), "=r"(*reinterpret_cast(&_tmem_load_1[22])), "=r"(*reinterpret_cast(&_tmem_load_1[23])), "=r"(*reinterpret_cast(&_tmem_load_1[24])), "=r"(*reinterpret_cast(&_tmem_load_1[25])), "=r"(*reinterpret_cast(&_tmem_load_1[26])), "=r"(*reinterpret_cast(&_tmem_load_1[27])), "=r"(*reinterpret_cast(&_tmem_load_1[28])), "=r"(*reinterpret_cast(&_tmem_load_1[29])), "=r"(*reinterpret_cast(&_tmem_load_1[30])), "=r"(*reinterpret_cast(&_tmem_load_1[31])) - : "r"(taddr + 256 + (unsigned int)row_addr)); - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_1[32])), "=r"(*reinterpret_cast(&_tmem_load_1[33])), "=r"(*reinterpret_cast(&_tmem_load_1[34])), "=r"(*reinterpret_cast(&_tmem_load_1[35])), "=r"(*reinterpret_cast(&_tmem_load_1[36])), "=r"(*reinterpret_cast(&_tmem_load_1[37])), "=r"(*reinterpret_cast(&_tmem_load_1[38])), "=r"(*reinterpret_cast(&_tmem_load_1[39])), "=r"(*reinterpret_cast(&_tmem_load_1[40])), "=r"(*reinterpret_cast(&_tmem_load_1[41])), "=r"(*reinterpret_cast(&_tmem_load_1[42])), "=r"(*reinterpret_cast(&_tmem_load_1[43])), "=r"(*reinterpret_cast(&_tmem_load_1[44])), "=r"(*reinterpret_cast(&_tmem_load_1[45])), "=r"(*reinterpret_cast(&_tmem_load_1[46])), "=r"(*reinterpret_cast(&_tmem_load_1[47])), "=r"(*reinterpret_cast(&_tmem_load_1[48])), "=r"(*reinterpret_cast(&_tmem_load_1[49])), "=r"(*reinterpret_cast(&_tmem_load_1[50])), "=r"(*reinterpret_cast(&_tmem_load_1[51])), "=r"(*reinterpret_cast(&_tmem_load_1[52])), "=r"(*reinterpret_cast(&_tmem_load_1[53])), "=r"(*reinterpret_cast(&_tmem_load_1[54])), "=r"(*reinterpret_cast(&_tmem_load_1[55])), "=r"(*reinterpret_cast(&_tmem_load_1[56])), "=r"(*reinterpret_cast(&_tmem_load_1[57])), "=r"(*reinterpret_cast(&_tmem_load_1[58])), "=r"(*reinterpret_cast(&_tmem_load_1[59])), "=r"(*reinterpret_cast(&_tmem_load_1[60])), "=r"(*reinterpret_cast(&_tmem_load_1[61])), "=r"(*reinterpret_cast(&_tmem_load_1[62])), "=r"(*reinterpret_cast(&_tmem_load_1[63])) - : "r"(taddr + 256 + (unsigned int)row_addr + 32)); - int col_base = col_half_1 * 64; - if (head_idx_1 < num_heads) { - #pragma unroll - for (int offset = 0; offset < 64; offset += 8) { - { - const float2 _prescale2_0 = {inv_sum * output_scale, inv_sum * output_scale}; - #if __CUDA_ARCH__ >= 1000 - #pragma unroll - for (int _ps = 0; _ps < 4; _ps++) - mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_1[offset])[_ps], _prescale2_0); - #else - #pragma unroll - for (int _ps = 0; _ps < 8; _ps++) - _tmem_load_1[offset + _ps] *= inv_sum * output_scale; - #endif - __nv_bfloat162 _pk[4]; - _pk[0] = __floats2bfloat162_rn(_tmem_load_1[offset + 0], _tmem_load_1[offset + 1]); - _pk[1] = __floats2bfloat162_rn(_tmem_load_1[offset + 2], _tmem_load_1[offset + 3]); - _pk[2] = __floats2bfloat162_rn(_tmem_load_1[offset + 4], _tmem_load_1[offset + 5]); - _pk[3] = __floats2bfloat162_rn(_tmem_load_1[offset + 6], _tmem_load_1[offset + 7]); - *reinterpret_cast(&((__nv_bfloat16*)(O + (output_base + col_base + offset)))[0]) = *reinterpret_cast(&_pk[0]); - } - } - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(tmem_dealloc_addr); - } - } - // ---- Role: mma_warp ---- - if (warp == 8) { - { // mma_warp_main - unsigned int _phase_q_full_0 = 0; - mbarrier_wait(q_full_addr, _phase_q_full_0); - _phase_q_full_0 ^= 1; - unsigned int _phase_kv_full = 0; - #pragma unroll - for (int k_stage = 0; k_stage < 4; k_stage++) { - mbarrier_wait(kv_full_addr + (k_stage) * 8, _phase_kv_full); - asm volatile("tcgen05.fence::after_thread_sync;"); - int _mma_a_lo_0 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (k_stage * 2) * 512); - int _mma_b_lo_0 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (k_stage) * 2048); - { - uint64_t _mma_ss_a_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_0); - uint64_t _mma_ss_b_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_0); - if (elect_sync()) { - tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, ((k_stage == 0) ? 0 : 1)); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); - } - } - int _mma_a_lo_1 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (k_stage * 2 + 1) * 512); - int _mma_b_lo_1 = make_warp_uniform((((smem_kv_addr + 16384) >> 4) & 0x3FFF) + (k_stage) * 2048); - { - uint64_t _mma_ss_a_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_1); - uint64_t _mma_ss_b_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_1); - if (elect_sync()) { - tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); - } - } - if (k_stage == 3) { - elect_commit(s_full_addr); - } - elect_commit(kv_empty_addr + (k_stage) * 8); - } - int _mma_prepared_b_lo_2 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (0) * 2048); - unsigned int _phase_p_full_0 = 0; - mbarrier_wait(p_full_addr, _phase_p_full_0); - _phase_p_full_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - mbarrier_wait(kv_full_addr, 1); - asm volatile("tcgen05.fence::after_thread_sync;"); - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 dhi, blo, ta, id;\n\t" - ".reg .b64 db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "" - "mov.b32 dhi, 0x40004040;\n\t" - "mov.b32 id, 69272720;\n\t" - "mov.b32 ta, %2;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p0;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "}\n" - :: "r"((tmem_tmem + (256))), "r"(_mma_prepared_b_lo_2), "r"(tmem_tmem + 128), "r"(0)); - elect_commit(o_done_addr); - elect_commit(kv_empty_addr); - unsigned int _phase_tmem_dealloc_0 = 0; - mbarrier_wait(tmem_dealloc_addr, _phase_tmem_dealloc_0); - _phase_tmem_dealloc_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); - asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); - } - } - // ---- Role: empty ---- - if (warp >= 9 && warp <= 11) { - // idle — no tasks assigned - } - // ---- Role: load_warp ---- - if (warp >= 12 && warp <= 15) { - asm volatile("setmaxnreg.dec.sync.aligned.u32 64;"); - { // load_warp_main - const int load_warp_rank = warp - 12; - int work_idx_1 = blockIdx.x; - int query_head_work_2 = work_idx_1 >> 2; - int query_idx_2 = query_head_work_2 / num_head_tiles; - int head_tile_2 = query_head_work_2 % num_head_tiles; - int v_chunk_1 = work_idx_1 & 3; - int* row_ptr = swa_indices + (query_idx_2 * swa_index_stride); - int index_offset = load_warp_rank * 32 + lane; - int sparse_row = -1; - if (index_offset < sparse_topk) { - sparse_row = row_ptr[index_offset]; - } - if (load_warp_rank == 0) { - if (elect_sync()) { - mbarrier_arrive_expect_tx(q_full_addr, 65536); - #pragma unroll - for (int q_stage = 0; q_stage < 8; q_stage++) { - tma_4d_gmem2smem(smem_q_addr + (unsigned int)(q_stage * 8192), (&tmap_q), 0, head_tile_2 * 64, q_stage, query_idx_2, q_full_addr); - } - } - } - smem_indices[index_offset] = sparse_row; - unsigned int _vote_0 = __ballot_sync(0xFFFFFFFF, sparse_row < 0); - unsigned int invalid_indices = _vote_0; - int valid_prefix = 32; - if (invalid_indices != 0) { - int _ffs_0 = __ffs(invalid_indices); - valid_prefix = _ffs_0 - 1; - } - if (elect_sync()) { - smem_index_flags[load_warp_rank] = valid_prefix; - } - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - asm volatile("barrier.sync 8, 128;" ::: "memory"); - if ((load_warp_rank & 1) != 0) { - if (elect_sync()) { - int lower_prefix = smem_index_flags[load_warp_rank - 1]; - int upper_prefix = smem_index_flags[load_warp_rank]; - smem_index_flags[load_warp_rank] = ((lower_prefix < 32) ? lower_prefix : 32 + upper_prefix); - } - } - unsigned int _phase_kv_empty = 1; - #pragma unroll - for (int k_stage_1 = 0; k_stage_1 < 4; k_stage_1++) { - mbarrier_wait(kv_empty_addr + (k_stage_1) * 8, _phase_kv_empty); - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (k_stage_1) * 8, 8192); - } - #pragma unroll - for (int local_group = 0; local_group < 8; local_group++) { - int dst_k = smem_kv_addr + (unsigned int)(k_stage_1 * 32768); - int group = load_warp_rank * 8 + local_group; - int group_offset = group * 4; - int raw_rows[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows[0])), "=r"(*reinterpret_cast(&raw_rows[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows[(0) + 3])) - : "r"(smem_indices_addr + (unsigned int)(group_offset * 4))); - int raw0 = raw_rows[0]; - int raw1 = raw_rows[1]; - int raw2 = raw_rows[2]; - int raw3 = raw_rows[3]; - int row0 = ((raw0 >= 0) ? raw0 : 0); - int row1 = ((raw1 >= 0) ? raw1 : 0); - int row2 = ((raw2 >= 0) ? raw2 : 0); - int row3 = ((raw3 >= 0) ? raw3 : 0); - if (elect_sync()) { - tma_gather4_gmem2smem(dst_k + group * 512, (&tmap_swa_kv), k_stage_1 * 128, row0, row1, row2, row3, kv_full_addr + (k_stage_1) * 8); - tma_gather4_gmem2smem(dst_k + 16384 + group * 512, (&tmap_swa_kv), k_stage_1 * 128 + 64, row0, row1, row2, row3, kv_full_addr + (k_stage_1) * 8); - } - } - } - #pragma unroll - for (int v_stage = 0; v_stage < 1; v_stage++) { - mbarrier_wait(kv_empty_addr + (v_stage) * 8, 0); - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (v_stage) * 8, 8192); - } - #pragma unroll - for (int local_group_1 = 0; local_group_1 < 8; local_group_1++) { - int dst_v = smem_v_addr + (unsigned int)(v_stage * 32768); - int v_col = v_chunk_1 * 128 + v_stage * 64; - int group_1 = load_warp_rank * 8 + local_group_1; - int group_offset_1 = group_1 * 4; - int raw_rows_1[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_1[0])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 3])) - : "r"(smem_indices_addr + (unsigned int)(group_offset_1 * 4))); - int raw0_1 = raw_rows_1[0]; - int raw1_1 = raw_rows_1[1]; - int raw2_1 = raw_rows_1[2]; - int raw3_1 = raw_rows_1[3]; - int row0_1 = ((raw0_1 >= 0) ? raw0_1 : 0); - int row1_1 = ((raw1_1 >= 0) ? raw1_1 : 0); - int row2_1 = ((raw2_1 >= 0) ? raw2_1 : 0); - int row3_1 = ((raw3_1 >= 0) ? raw3_1 : 0); - if (elect_sync()) { - tma_gather4_gmem2smem(dst_v + group_1 * 512, (&tmap_swa_kv), v_col, row0_1, row1_1, row2_1, row3_1, kv_full_addr + (v_stage) * 8); - tma_gather4_gmem2smem(dst_v + 16384 + group_1 * 512, (&tmap_swa_kv), v_col + 64, row0_1, row1_1, row2_1, row3_1, kv_full_addr + (v_stage) * 8); - } - } - } - } - } - - // Cleanup -} - -} // extern "C" diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_5e5d6648b7bd42bcd0d6_binding.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_5e5d6648b7bd42bcd0d6_binding.cu deleted file mode 100644 index b4398bf2025..00000000000 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_5e5d6648b7bd42bcd0d6_binding.cu +++ /dev/null @@ -1,441 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -// Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_5e5d6648b7bd42bcd0d6'. -#include -#include -#include -#include - -#include "tvm_ffi_utils.h" - -#include -#include -#include -#include -#include -#include -#include - -extern "C" __global__ void kernel_cake_dsv4_5e5d6648b7bd42bcd0d6(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ cum_seq_lens_q, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); - - -namespace cake_host_shim_1d4218665228ec14 { - -using tvm::ffi::Optional; -using tvm::ffi::TensorView; - -class ScopedCudaDevice { - public: - explicit ScopedCudaDevice(int device_id) { - cudaError_t error = cudaGetDevice(&previous_device_); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed before host-shim launch: cudaError=" - << static_cast(error); - if (previous_device_ != device_id) { - error = cudaSetDevice(device_id); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaSetDevice failed before host-shim launch for cuda:" - << device_id << ": cudaError=" << static_cast(error); - restore_ = true; - } - } - - ScopedCudaDevice(const ScopedCudaDevice&) = delete; - ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; - - ~ScopedCudaDevice() noexcept { - if (restore_) { - (void)cudaSetDevice(previous_device_); - } - } - - private: - int previous_device_ = -1; - bool restore_ = false; -}; - -inline int64_t CakeDeviceMultiprocessorCount(int device_id) { - constexpr int kMaxCachedCudaDevices = 64; - TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) - << "physical-SM-count cache does not cover cuda:" << device_id; - static std::atomic count_by_device[kMaxCachedCudaDevices]{}; - int cached = count_by_device[device_id].load(std::memory_order_acquire); - if (cached > 0) return cached; - - int count = 0; - cudaError_t error = cudaDeviceGetAttribute( - &count, cudaDevAttrMultiProcessorCount, device_id); - TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) - << "querying multiProcessorCount failed for cuda:" << device_id - << ": cudaError=" << static_cast(error) << ", count=" << count; - // Concurrent first launches may repeat the immutable device query, but all - // publication is atomic and every later hot-path lookup is one acquire load. - count_by_device[device_id].store(count, std::memory_order_release); - return count; -} - -inline void CheckCudaTensor(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) - << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; -} - -inline void CheckSameCudaDevice( - const TensorView& t, - const TensorView& reference, - const char* name, - const char* reference_name) { - TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) - << name << " must be on the same CUDA device as " << reference_name - << ": got cuda:" << t.device().device_id - << " versus cuda:" << reference.device().device_id; -} - -inline void CheckCurrentCudaDevice( - const TensorView& reference, - const char* reference_name) { - int current_device = -1; - cudaError_t error = cudaGetDevice(¤t_device); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed while validating " << reference_name - << ": cudaError=" << static_cast(error); - TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) - << "current CUDA device must match " << reference_name - << ": current=cuda:" << current_device - << ", tensor=cuda:" << reference.device().device_id; -} - -inline void CheckContiguous(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; -} - -inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { - DLDataType d = t.dtype(); - TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) - << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits - << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits - << ", lanes=" << (int)d.lanes << ")"; -} - -// A logical axis.outer(trailing) folds every source dim above the trailing -// dimensions. Shape products are independent of physical strides, so verify -// the leading dimensions form one dense row-major chain instead of inventing -// a "folded stride". The descriptor reads its exact adjacent physical step -// separately through stride[-(trailing + 1)]. -inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { - TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) - << name << " cannot fold leading dimensions above " << trailing - << " trailing dims from ndim=" << t.ndim(); - int outer_last = t.ndim() - trailing - 1; - if (outer_last <= 0) { - return; - } - int64_t step = t.stride(outer_last); - TVM_FFI_CHECK(step > 0, ValueError) - << name << " physical strides must be positive"; - int64_t expected = step; - for (int axis = outer_last - 1; axis >= 0; --axis) { - expected *= t.size(axis + 1); - if (t.size(axis) > 1) { - TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) - << name << " leading dims are not physically foldable above " << trailing - << " trailing dims: stride(" << axis << ")=" << t.stride(axis) - << ", expected " << expected; - } - } -} - -// 4D TMA descriptor for buffer 'tmap_q' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 3, ValueError) - << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - int64_t d2 = t.size(t.ndim() - 2); - int64_t d3 = t.size(t.ndim() - 3); - TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) - << "TMA source 'tmap_q' trailing dims must be positive"; - TVM_FFI_CHECK(d2 % 64 == 0, ValueError) - << "TMA source 'tmap_q' extent " << d2 - << " must divide exactly by " << 64; - uint64_t global_dim[4] = {(uint64_t)(d1), (uint64_t)(64), (uint64_t)((d2 / 64)), (uint64_t)(d3)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; - TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) - << "TMA box (128, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; - int64_t carrier_stride_0 = d1; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; - int64_t carrier_stride_1 = (64 * d1); - TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; - TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; - TVM_FFI_CHECK((carrier_stride_1 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; - int64_t carrier_stride_2 = (d2 * d1); - TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; - TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; - TVM_FFI_CHECK((carrier_stride_2 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; - uint64_t global_strides[3] = { - (uint64_t)((carrier_stride_0 * 8) / 8), - (uint64_t)((carrier_stride_1 * 8) / 8), - (uint64_t)((carrier_stride_2 * 8) / 8), - }; - uint32_t box_dim[4] = {128u, 64u, 1u, 1u}; - uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_L2_128B, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_swa_kv' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_swa_kv' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; - TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (128, 1) exceeds resolved global dims for 'tmap_swa_kv'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 8) / 8), - }; - uint32_t box_dim[2] = {128u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_compressed_kv' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_compressed_kv(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_compressed_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_compressed_kv' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_compressed_kv' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_compressed_kv"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_compressed_kv' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved a non-positive global dim"; - TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (128, 1) exceeds resolved global dims for 'tmap_compressed_kv'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 8) / 8), - }; - uint32_t box_dim[2] = {128u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_kv') failed: CUresult=" << (int)r; - return tm; -} - -void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_compressed_kv, TensorView arg_O, TensorView arg_cum_seq_lens_q, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, int64_t grid_x, int64_t grid_y, int64_t grid_z) { - DLDevice dev = arg_tmap_q.device(); - ScopedCudaDevice device_guard(dev.device_id); - CheckCudaTensor(arg_tmap_q, "tmap_q"); - CheckDtype(arg_tmap_q, "tmap_q", 1, 8, 1); - CheckContiguous(arg_tmap_q, "tmap_q"); - CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); - CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 1, 8, 1); - CheckCudaTensor(arg_tmap_compressed_kv, "tmap_compressed_kv"); - CheckDtype(arg_tmap_compressed_kv, "tmap_compressed_kv", 1, 8, 1); - CheckCudaTensor(arg_O, "O"); - CheckDtype(arg_O, "O", 4, 16, 1); - CheckContiguous(arg_O, "O"); - CheckCudaTensor(arg_cum_seq_lens_q, "cum_seq_lens_q"); - CheckDtype(arg_cum_seq_lens_q, "cum_seq_lens_q", 0, 32, 1); - CheckContiguous(arg_cum_seq_lens_q, "cum_seq_lens_q"); - CheckCudaTensor(arg_swa_indices, "swa_indices"); - CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); - CheckContiguous(arg_swa_indices, "swa_indices"); - CheckCudaTensor(arg_compressed_indices, "compressed_indices"); - CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); - CheckContiguous(arg_compressed_indices, "compressed_indices"); - CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); - CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckCudaTensor(arg_sinks, "sinks"); - CheckDtype(arg_sinks, "sinks", 2, 32, 1); - CheckContiguous(arg_sinks, "sinks"); - CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); - CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); - CheckContiguous(arg_bmm1_scale, "bmm1_scale"); - CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); - CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); - CheckContiguous(arg_bmm2_scale, "bmm2_scale"); - TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) - << "scalar 'num_heads' value " << arg_num_heads - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) - << "scalar 'swa_index_stride' value " << arg_swa_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) - << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) - << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) - << "scalar 'sparse_topk' value " << arg_sparse_topk - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) - << "scalar 'has_sinks' value " << arg_has_sinks - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) - << "scalar 'total_work_items' value " << arg_total_work_items - << " is outside i32 range [-2147483648, 2147483647]"; - CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); - CheckSameCudaDevice(arg_tmap_compressed_kv, arg_tmap_q, "tmap_compressed_kv", "tmap_q"); - CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); - CheckSameCudaDevice(arg_cum_seq_lens_q, arg_tmap_q, "cum_seq_lens_q", "tmap_q"); - CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); - CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); - CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); - CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); - CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); - CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); - TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) - << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y - << ", " << grid_z << ")"; - - cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); - CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); - CUtensorMap p_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); - CUtensorMap p_tmap_compressed_kv = EncodeTma_tmap_compressed_kv(arg_tmap_compressed_kv); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); - int* p_cum_seq_lens_q = static_cast(arg_cum_seq_lens_q.data_ptr()); - int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); - int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); - int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); - float* p_sinks = static_cast(arg_sinks.data_ptr()); - float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); - float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); - int32_t v_num_heads = (int32_t)arg_num_heads; - int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; - int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; - int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; - int32_t v_sparse_topk = (int32_t)arg_sparse_topk; - int32_t v_has_sinks = (int32_t)arg_has_sinks; - int32_t v_total_work_items = (int32_t)arg_total_work_items; - void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_tmap_compressed_kv, &p_O, &p_cum_seq_lens_q, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_has_sinks, &v_total_work_items}; - - static std::mutex smem_mu; - static auto* smem_status_by_device = - new std::unordered_map(); - cudaError_t smem_status = cudaSuccess; - { - std::lock_guard lock(smem_mu); - auto it = smem_status_by_device->find(dev.device_id); - if (it == smem_status_by_device->end()) { - int current_device = -1; - smem_status = cudaGetDevice(¤t_device); - if (smem_status == cudaSuccess && current_device != dev.device_id) { - smem_status = cudaErrorInvalidDevice; - } - if (smem_status == cudaSuccess) { - smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_5e5d6648b7bd42bcd0d6), - cudaFuncAttributeMaxDynamicSharedMemorySize, - 203136); - } - smem_status_by_device->emplace(dev.device_id, smem_status); - } else { - smem_status = it->second; - } - } - TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_5e5d6648b7bd42bcd0d6 failed: " - << cudaGetErrorString(smem_status); - dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); - dim3 block(512u, 1u, 1u); - - cudaLaunchConfig_t config{}; - config.gridDim = grid; - config.blockDim = block; - config.dynamicSmemBytes = 203136u; - config.stream = stream; - cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_5e5d6648b7bd42bcd0d6), kargs); - TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_5e5d6648b7bd42bcd0d6 failed: " - << cudaGetErrorString(launch_status); -} - -} // namespace cake_host_shim_1d4218665228ec14 - -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_1d4218665228ec14::Run); diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_5e5d6648b7bd42bcd0d6_kernel.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_5e5d6648b7bd42bcd0d6_kernel.cu deleted file mode 100644 index 2a64d3f6d3e..00000000000 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_5e5d6648b7bd42bcd0d6_kernel.cu +++ /dev/null @@ -1,2973 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -typedef signed char int8_t; -typedef unsigned char uint8_t; -typedef unsigned short uint16_t; -typedef unsigned int uint32_t; -#if defined(__CUDACC_RTC__) -typedef unsigned long long uint64_t; -#else -typedef unsigned long uint64_t; -#endif -static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); -typedef signed int int32_t; -typedef short int int16_t; -struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; -struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; -static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); -static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); - -#if defined(__CUDACC_RTC__) -typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; -#else -#include -#endif - -static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); -static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); -#include -#include - -#define CAKE_INF CUDART_INF_F -#define TMEM_NCOLS 512 -#define TMEM_TMEM_OFFSET 0 -#define NUM_Q_PIPE_STAGES 2 -#define NUM_KV_PIPE_STAGES 4 -#define NUM_PAGE_PIPE_STAGES 6 -#define NUM_SCORE_PIPE_STAGES 2 -#define NUM_STATS_PIPE_STAGES 2 -#define NUM_WORK_PIPE_STAGES 2 -#define NUM_THROTTLE_PIPE_STAGES 2 -#define SMEM_SMEM_Q_OFF 0 -#define SMEM_SMEM_Q_STAGE_BYTES 32768 -#define SMEM_SMEM_Q_STRIDE 32768 -#define SMEM_SMEM_K_OFF 65536 -#define SMEM_SMEM_K_STAGE_BYTES 32768 -#define SMEM_SMEM_K_STRIDE 32768 -#define SMEM_SMEM_V_OFF 65536 -#define SMEM_SMEM_V_STAGE_BYTES 32768 -#define SMEM_SMEM_V_STRIDE 32768 -#define SMEM_SMEM_PAGE_OFFSETS_OFF 196608 -#define SMEM_SMEM_PAGE_OFFSETS_STAGE_BYTES 1024 -#define SMEM_SMEM_PAGE_OFFSETS_STRIDE 1024 -#define SMEM_WORK_RESPONSE_OFF 202752 -#define SMEM_WORK_RESPONSE_STAGE_BYTES 16 -#define SMEM_WORK_RESPONSE_STRIDE 16 -#define SMEM_TOTAL 203136 -#define THREADS 512 - -#include - -__device__ __forceinline__ uint32_t elect_sync() { - uint32_t pred = 0; - asm volatile( - "{\n\t" - ".reg .pred %%px;\n\t" - "elect.sync _|%%px, %1;\n\t" - "@%%px mov.s32 %0, 1;\n\t" - "}\n" - : "+r"(pred) - : "r"(0xFFFFFFFF)); - return pred; -} - - -__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { - asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" - :: "r"(mbar_addr), "r"(count) : "memory"); -} - -__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { - asm volatile("mbarrier.init.b64 [%0], %1;" - :: "l"(mbar_addr), "r"(count) : "memory"); -} - - -__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - - -// CTA-local pipelines have short, resident producer/consumer edges. Omitting -// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the -// explicit loop still makes this helper blocking until acquire succeeds. -__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE;\n\t" - "bra.uni LAB_WAIT;\n\t" - "DONE:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Source-faithful relaxed CTA wait used only by a typed protocol that does -// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. -__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, 10000000;\n\t" - "@P1 bra.uni DONE_RELAXED;\n\t" - "bra.uni LAB_WAIT_RELAXED;\n\t" - "DONE_RELAXED:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Exact source ports may request the PTX suspendTimeHint operand explicitly. -// The hint is expressed in nanoseconds and is kept separate from the canonical -// no-hint CTA helper so unrelated schedules retain their existing retry path. -__device__ __forceinline__ void mbarrier_wait_suspend( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_SUSPEND:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_SUSPEND;\n\t" - "bra.uni LAB_WAIT_SUSPEND;\n\t" - "DONE_SUSPEND:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE_CLUSTER;\n\t" - "bra.uni LAB_WAIT_CLUSTER;\n\t" - "DONE_CLUSTER:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - ".reg .u32 WAIT_ADDR;\n\t" - "mov.u32 WAIT_ADDR, %0;\n\t" - "LAB_WAIT_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [WAIT_ADDR], %1, %2;\n\t" - "@P1 bra.uni DONE_HINT;\n\t" - "bra.uni LAB_WAIT_HINT;\n\t" - "DONE_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -// Exact unqualified CTA wait used by source schedules whose PTX intentionally -// omits the acquire qualifier while retaining a typed suspendTimeHint operand. -__device__ __forceinline__ void mbarrier_wait_relaxed_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED_HINT:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra DONE_RELAXED_HINT;\n\t" - "bra LAB_WAIT_RELAXED_HINT;\n\t" - "DONE_RELAXED_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); -} - -__device__ __forceinline__ void mbarrier_wait_cluster_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" - "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" - "DONE_CLUSTER_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_suspend( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait_cluster(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); - } -} - - -__device__ __forceinline__ void tcgen05_mma_f8f6f4( - int taddr, uint64_t a_desc, uint64_t b_desc, - uint32_t i_desc, int enable_input_d) { - asm volatile( - "{\n\t" - ".reg .pred p;\n\t" - "setp.ne.b32 p, %4, 0;\n\t" - "tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], %1, %2, %3, p;\n\t" - "}\n" - :: "r"(taddr), "l"(a_desc), "l"(b_desc), - "r"(i_desc), "r"(enable_input_d) - : "memory"); -} - - -__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { - return (x & 0x3FFFFULL) >> 4ULL; -} - - -union MmaSmemDesc { - uint64_t u64; - uint32_t u32[2]; -}; - -__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { - MmaSmemDesc tmp; - tmp.u64 = smem_desc; - tmp.u32[0] += offset; - smem_desc = tmp.u64; -} - - -__device__ __forceinline__ void elect_commit(int mbar_addr) { - asm volatile( - "{\n\t" - ".reg .pred leader;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" - ".shared::cluster.b64 [%0];\n\t" - "}\n" - :: "r"(mbar_addr)); -} - - -__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { - asm volatile( - "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" - :: "r"(mbar_addr), "r"(bytes) : "memory"); -} - - -__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7," - " %8, %9, %10, %11, %12, %13, %14, %15," - " %16, %17, %18, %19, %20, %21, %22, %23," - " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), - "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), - "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), - "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), - "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), - "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), - "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ float approx_exp2(float x) { - float y; - asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float approx_rcp(float x) { - float y; - asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float max_noftz(float a, float b) { - float c; - asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); - return c; -} - - -__device__ __forceinline__ float warp_reduce_max(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); - return val; -} - - -__device__ __forceinline__ float warp_reduce_sum(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val += __shfl_xor_sync(0xFFFFFFFF, val, offset); - return val; -} - - -__device__ __forceinline__ float row_max_reduce(float2 acc) { - return max_noftz(acc.x, acc.y); -} - - -__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { - #pragma unroll - for (int j = 0; j < 16; j++) { - if (j % 2 == 0) - acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); - else - acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); - } -} - - -__device__ __forceinline__ float2 ex2_emulation_f32x2_value(float2 value) { - const float c0 = 1.0f, c1 = 0.695146143436431884765625f; - const float c2 = 0.227564394474029541015625f, c3 = 0.077119089663028717041015625f; - const float magic = 12582912.0f; - float x0 = max_noftz(value.x, -127.0f), x1 = max_noftz(value.y, -127.0f); - float2 xc2 = make_float2(x0, x1), magic2 = make_float2(magic, magic); - float2 xr2; - asm("add.rm.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xr2) - : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&magic2)); - float2 c3_2 = make_float2(c3, c3), c2_2 = make_float2(c2, c2); - float2 c1_2 = make_float2(c1, c1), c0_2 = make_float2(c0, c0); - float2 xrb2, xfrac2; - asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xrb2) - : "l"(*(unsigned long long*)&xr2), "l"(*(unsigned long long*)&magic2)); - asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xfrac2) - : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&xrb2)); - float2 poly2; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&c3_2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c2_2)); - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c1_2)); - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c0_2)); - int x0r_i, x1r_i, p0_i, p1_i; - asm("mov.b64 {%0, %1}, %2;" : "=r"(x0r_i), "=r"(x1r_i) : "l"(*(unsigned long long*)&xr2)); - asm("mov.b64 {%0, %1}, %2;" : "=r"(p0_i), "=r"(p1_i) : "l"(*(unsigned long long*)&poly2)); - float r0, r1; - asm("mov.b32 %0, %1;" : "=f"(r0) : "r"((x0r_i << 23) + p0_i)); - asm("mov.b32 %0, %1;" : "=f"(r1) : "r"((x1r_i << 23) + p1_i)); - return make_float2(r0, r1); -} - -__device__ __forceinline__ void ex2_emulation_f32x2(float* x0_ptr, float* x1_ptr) { - float2 result = ex2_emulation_f32x2_value(make_float2(*x0_ptr, *x1_ptr)); - *x0_ptr = result.x; *x1_ptr = result.y; -} - -__device__ __forceinline__ void softmax_frag_exp2_cast( - float* sv, uint32_t* pv, int use_emu) -{ - #pragma unroll - for (int j = 0; j < 16; j++) { - if (use_emu && j >= 12) - ex2_emulation_f32x2(&sv[j*2], &sv[j*2+1]); - else { - sv[j*2] = approx_exp2(sv[j*2]); - sv[j*2+1] = approx_exp2(sv[j*2+1]); - } - } - #pragma unroll - for (int j = 0; j < 16; j++) { - __nv_bfloat162 bf = __float22bfloat162_rn({sv[j*2], sv[j*2+1]}); - pv[j] = reinterpret_cast(bf); - } -} - - - -__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { - const float2* sv2 = reinterpret_cast(sv); - #pragma unroll - for (int j = 0; j < 16; j++) { - asm("add.f32x2 %0, %1, %2;" - : "+l"(reinterpret_cast(*acc)) - : "l"(reinterpret_cast(*acc)), - "l"(reinterpret_cast(sv2[j]))); - } -} - - -__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { - asm("mul.rn.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { - asm("mul.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { - asm("add.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { - asm("add.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { - asm("sub.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { - asm("sub.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("add.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { - float2 r; - asm("sub.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("sub.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ void fma_scale_x32( - float* sv, const float2* scale2, const float2* neg_max2) -{ - float2* sv_2 = reinterpret_cast(sv); - #pragma unroll - for (int j = 0; j < 16; j++) - fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); -} - -__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("mul.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - - -// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) - -__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("add.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("add.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("add.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("add.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("add.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("add.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("add.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - - -__device__ __forceinline__ void fence_async_shared() { - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); -} - - -__device__ __forceinline__ uint64_t make_smem_desc(int addr) { - const int SBO = 1024; - return desc_encode(addr) - | (desc_encode(SBO) << 32ULL) - | (1ULL << 46ULL) - | (2ULL << 61ULL); -} - - -__device__ __forceinline__ void tma_2d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3}], [%4];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_4d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.4d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5}], [%6];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr) { - // Canonical .shared::cta form for non-multicast gather4, matching - // trtllm-gen / cuda_ptx and the PTX ISA qualifier order - // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, - // .shared::cluster is reserved for the multicast variant (ctaMask). - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { - asm volatile( - "tcgen05.commit.cta_group::1.mbarrier::arrive::one" - ".shared::cluster.b64 [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x16.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7," - " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), - "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), - "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { - tmem_ld_x16(dst, addr); - asm volatile("tcgen05.wait::ld.sync.aligned;"); -} - - -__device__ __forceinline__ void tmem_ld_x8(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x8.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { - uint32_t result; - asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" - : "=r"(result) : "r"(val)); - return result; -} - - -template -struct CakePipelineState { - uint32_t index_; - uint32_t phase_; - uint32_t count_; - - __device__ __forceinline__ uint32_t index() const { return index_; } - __device__ __forceinline__ uint32_t phase() const { return phase_; } - __device__ __forceinline__ uint32_t count() const { return count_; } - __device__ __forceinline__ CakePipelineState& operator++() { - ++count_; - ++index_; - if (index_ == Stages) { - index_ = 0; - phase_ ^= 1; - } - return *this; - } -}; - -template -__device__ __forceinline__ CakePipelineState -cake_pipeline_state_at(CakePipelineState state) { - #pragma unroll - for (int step = 0; step < Offset; ++step) { - ++state; - } - return state; -} - -extern "C" { - -__global__ __launch_bounds__(512, 1) void -kernel_cake_dsv4_5e5d6648b7bd42bcd0d6(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ cum_seq_lens_q, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items) -{ - const int tid = threadIdx.x; - const int warp = make_warp_uniform(tid / 32); - const int lane = tid % 32; - - extern __shared__ __align__(1024) char smem_raw[]; - int smem; - asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); - smem = make_warp_uniform(smem); - - const int mbar_base = smem + 202784; - #define q_full_addr (mbar_base + 0) - #define q_empty_addr (mbar_base + 16) - #define kv_full_addr (mbar_base + 32) - #define kv_empty_addr (mbar_base + 64) - #define page_full_addr (mbar_base + 96) - #define page_empty_addr (mbar_base + 144) - #define score_full_addr (mbar_base + 192) - #define score_empty_addr (mbar_base + 208) - #define stats_full_addr (mbar_base + 224) - #define stats_empty_addr (mbar_base + 240) - #define o_full_addr (mbar_base + 256) - #define o_empty_addr (mbar_base + 264) - #define work_full_addr (mbar_base + 272) - #define work_empty_addr (mbar_base + 288) - #define throttle_full_addr (mbar_base + 304) - #define throttle_empty_addr (mbar_base + 320) - - const int bid = blockIdx.x; - const int num_bids = gridDim.x; - - const int cta_rank = 0; - - // Kernel setup ops - uint8_t* smem_q = reinterpret_cast(smem_raw + 0); - const int smem_q_addr = smem + 0; - uint8_t* smem_k = reinterpret_cast(smem_raw + 65536); - const int smem_k_addr = smem + 65536; - uint8_t* smem_v = reinterpret_cast(smem_raw + 65536); - const int smem_v_addr = smem + 65536; - int* smem_page_offsets = reinterpret_cast(smem_raw + 196608); - const int smem_page_offsets_addr = smem + 196608; - unsigned int* work_response = reinterpret_cast(smem_raw + 202752); - const int work_response_addr = smem + 202752; - asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_q))) : "memory"); - asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_swa_kv))) : "memory"); - asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_compressed_kv))) : "memory"); - - // Mbarrier init (16 pipeline groups, 0 ordered-sequence groups, 42 barriers) - // Mbarriers at smem_raw[202784..203120) - - if (warp == 12) { - uint32_t leader = elect_sync(); - if (leader) { - // --- pipeline 'q_pipe' --- - // q_full: 2 barriers, init_count=1 - mbarrier_init(smem + 202784, 1); - mbarrier_init(smem + 202792, 1); - // q_empty: 2 barriers, init_count=1 - mbarrier_init(smem + 202800, 1); - mbarrier_init(smem + 202808, 1); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - if (warp == 0) { - uint32_t leader = elect_sync(); - if (leader) { - // --- pipeline 'kv_pipe' --- - // kv_full: 4 barriers, init_count=4 - mbarrier_init(smem + 202816, 4); - mbarrier_init(smem + 202824, 4); - mbarrier_init(smem + 202832, 4); - mbarrier_init(smem + 202840, 4); - // kv_empty: 4 barriers, init_count=1 - mbarrier_init(smem + 202848, 1); - mbarrier_init(smem + 202856, 1); - mbarrier_init(smem + 202864, 1); - mbarrier_init(smem + 202872, 1); - // --- pipeline 'page_pipe' --- - // page_full: 6 barriers, init_count=32 - mbarrier_init(smem + 202880, 32); - mbarrier_init(smem + 202888, 32); - mbarrier_init(smem + 202896, 32); - mbarrier_init(smem + 202904, 32); - mbarrier_init(smem + 202912, 32); - mbarrier_init(smem + 202920, 32); - // page_empty: 6 barriers, init_count=128 - mbarrier_init(smem + 202928, 128); - mbarrier_init(smem + 202936, 128); - mbarrier_init(smem + 202944, 128); - mbarrier_init(smem + 202952, 128); - mbarrier_init(smem + 202960, 128); - mbarrier_init(smem + 202968, 128); - // --- pipeline 'score_pipe' --- - // score_full: 2 barriers, init_count=1 - mbarrier_init(smem + 202976, 1); - mbarrier_init(smem + 202984, 1); - // score_empty: 2 barriers, init_count=128 - mbarrier_init(smem + 202992, 128); - mbarrier_init(smem + 203000, 128); - // --- pipeline 'stats_pipe' --- - // stats_full: 2 barriers, init_count=128 - mbarrier_init(smem + 203008, 128); - mbarrier_init(smem + 203016, 128); - // stats_empty: 2 barriers, init_count=128 - mbarrier_init(smem + 203024, 128); - mbarrier_init(smem + 203032, 128); - // o_full: 1 barriers, init_count=1 - mbarrier_init(smem + 203040, 1); - // o_empty: 1 barriers, init_count=128 - mbarrier_init(smem + 203048, 128); - // --- pipeline 'work_pipe' --- - // work_full: 2 barriers, init_count=1 - mbarrier_init(smem + 203056, 1); - mbarrier_init(smem + 203064, 1); - // work_empty: 2 barriers, init_count=512 - mbarrier_init(smem + 203072, 512); - mbarrier_init(smem + 203080, 512); - // --- pipeline 'throttle_pipe' --- - // throttle_full: 2 barriers, init_count=128 - mbarrier_init(smem + 203088, 128); - mbarrier_init(smem + 203096, 128); - // throttle_empty: 2 barriers, init_count=32 - mbarrier_init(smem + 203104, 32); - mbarrier_init(smem + 203112, 32); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - - __syncwarp(); - - // TMEM alloc (512 columns, 512 used) - volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 203120); - if (warp == 0) { - int _tmem_hold = smem + 203120; - asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); - __syncwarp(); - asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); - } - - __syncthreads(); - asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); - - const int taddr = tmem_addr_storage[0]; - - // Kernel post-init ops - const int tmem_tmem = taddr; - - // ---- Ordered hardware-WG register redistribution ---- - // Inc phase consumes the registers released above. - if (warp >= 8 && warp <= 11) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 136;"); - } - - // ---- Role: softmax ---- - if (warp <= 3) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 152;"); - { // softmax_main - const int warp_rank = warp; - const int logical_row_origin = warp_rank * 16; - const int my_row = logical_row_origin + lane % 16; - const int col_half = lane / 16; - int score_cursor = 0; - int stats_cursor = 0; - unsigned int softmax_work_stage = 0; - unsigned int softmax_q = blockIdx.x; - unsigned int softmax_head = blockIdx.y; - unsigned int softmax_batch = blockIdx.z; - unsigned int _phase_work_full = 0; - #pragma unroll 1 - for (unsigned int _work = 0; _work < total_work_items; _work++) { - int q_begin = cum_seq_lens_q[(int)softmax_batch]; - int q_end = cum_seq_lens_q[(int)(softmax_batch + 1)]; - int q_local = (int)softmax_q; - int work_valid = q_local < q_end - q_begin; - asm volatile("griddepcontrol.wait;" ::: "memory"); - float softmax_rescale_log2 = bmm1_scale[0] * 1.4426950408889634f; - asm volatile("griddepcontrol.wait;" ::: "memory"); - float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; - if (work_valid != 0) { - int query_idx = q_begin + q_local; - int _max_4 = ((sparse_topk_lens[query_idx] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx] + sparse_topk_lens_offset) : (0)); - int _min_3 = ((_max_4) < (sparse_topk) ? (_max_4) : (sparse_topk)); - int active_topk = _min_3; - int num_tiles = (active_topk + 128 - 1) / 128; - float row_max = -CAKE_INF; - float row_sum = 0.0f; - #pragma unroll 1 - for (int tile = 0; tile < num_tiles; tile++) { - int pipeline_tile = score_cursor + tile; - int score_stage = pipeline_tile & 1; - int score_phase = pipeline_tile >> 1 & 1; - mbarrier_wait(score_full_addr + (score_stage) * 8, score_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - int score_col = ((score_stage != 0) ? 128 : 0); - float _tmem_load_0[64]; - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_0[0])), "=r"(*reinterpret_cast(&_tmem_load_0[1])), "=r"(*reinterpret_cast(&_tmem_load_0[2])), "=r"(*reinterpret_cast(&_tmem_load_0[3])), "=r"(*reinterpret_cast(&_tmem_load_0[4])), "=r"(*reinterpret_cast(&_tmem_load_0[5])), "=r"(*reinterpret_cast(&_tmem_load_0[6])), "=r"(*reinterpret_cast(&_tmem_load_0[7])), "=r"(*reinterpret_cast(&_tmem_load_0[8])), "=r"(*reinterpret_cast(&_tmem_load_0[9])), "=r"(*reinterpret_cast(&_tmem_load_0[10])), "=r"(*reinterpret_cast(&_tmem_load_0[11])), "=r"(*reinterpret_cast(&_tmem_load_0[12])), "=r"(*reinterpret_cast(&_tmem_load_0[13])), "=r"(*reinterpret_cast(&_tmem_load_0[14])), "=r"(*reinterpret_cast(&_tmem_load_0[15])), "=r"(*reinterpret_cast(&_tmem_load_0[16])), "=r"(*reinterpret_cast(&_tmem_load_0[17])), "=r"(*reinterpret_cast(&_tmem_load_0[18])), "=r"(*reinterpret_cast(&_tmem_load_0[19])), "=r"(*reinterpret_cast(&_tmem_load_0[20])), "=r"(*reinterpret_cast(&_tmem_load_0[21])), "=r"(*reinterpret_cast(&_tmem_load_0[22])), "=r"(*reinterpret_cast(&_tmem_load_0[23])), "=r"(*reinterpret_cast(&_tmem_load_0[24])), "=r"(*reinterpret_cast(&_tmem_load_0[25])), "=r"(*reinterpret_cast(&_tmem_load_0[26])), "=r"(*reinterpret_cast(&_tmem_load_0[27])), "=r"(*reinterpret_cast(&_tmem_load_0[28])), "=r"(*reinterpret_cast(&_tmem_load_0[29])), "=r"(*reinterpret_cast(&_tmem_load_0[30])), "=r"(*reinterpret_cast(&_tmem_load_0[31])) - : "r"(taddr + (unsigned int)score_col)); - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_0[32])), "=r"(*reinterpret_cast(&_tmem_load_0[33])), "=r"(*reinterpret_cast(&_tmem_load_0[34])), "=r"(*reinterpret_cast(&_tmem_load_0[35])), "=r"(*reinterpret_cast(&_tmem_load_0[36])), "=r"(*reinterpret_cast(&_tmem_load_0[37])), "=r"(*reinterpret_cast(&_tmem_load_0[38])), "=r"(*reinterpret_cast(&_tmem_load_0[39])), "=r"(*reinterpret_cast(&_tmem_load_0[40])), "=r"(*reinterpret_cast(&_tmem_load_0[41])), "=r"(*reinterpret_cast(&_tmem_load_0[42])), "=r"(*reinterpret_cast(&_tmem_load_0[43])), "=r"(*reinterpret_cast(&_tmem_load_0[44])), "=r"(*reinterpret_cast(&_tmem_load_0[45])), "=r"(*reinterpret_cast(&_tmem_load_0[46])), "=r"(*reinterpret_cast(&_tmem_load_0[47])), "=r"(*reinterpret_cast(&_tmem_load_0[48])), "=r"(*reinterpret_cast(&_tmem_load_0[49])), "=r"(*reinterpret_cast(&_tmem_load_0[50])), "=r"(*reinterpret_cast(&_tmem_load_0[51])), "=r"(*reinterpret_cast(&_tmem_load_0[52])), "=r"(*reinterpret_cast(&_tmem_load_0[53])), "=r"(*reinterpret_cast(&_tmem_load_0[54])), "=r"(*reinterpret_cast(&_tmem_load_0[55])), "=r"(*reinterpret_cast(&_tmem_load_0[56])), "=r"(*reinterpret_cast(&_tmem_load_0[57])), "=r"(*reinterpret_cast(&_tmem_load_0[58])), "=r"(*reinterpret_cast(&_tmem_load_0[59])), "=r"(*reinterpret_cast(&_tmem_load_0[60])), "=r"(*reinterpret_cast(&_tmem_load_0[61])), "=r"(*reinterpret_cast(&_tmem_load_0[62])), "=r"(*reinterpret_cast(&_tmem_load_0[63])) - : "r"(taddr + (unsigned int)score_col + 32)); - int valid_cols = active_topk - tile * 128 - col_half * 64; - if (valid_cols < 0) { - valid_cols = 0; - } - if (valid_cols > 64) { - valid_cols = 64; - } - uint32_t _slice_lo_mask_0; - { - int _lim_0 = valid_cols; - if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } - else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); - } - } - if (!(_slice_lo_mask_0 & (1u << 0))) _tmem_load_0[0] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 1))) _tmem_load_0[1] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 2))) _tmem_load_0[2] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 3))) _tmem_load_0[3] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 4))) _tmem_load_0[4] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 5))) _tmem_load_0[5] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 6))) _tmem_load_0[6] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 7))) _tmem_load_0[7] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 8))) _tmem_load_0[8] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 9))) _tmem_load_0[9] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 10))) _tmem_load_0[10] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 11))) _tmem_load_0[11] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 12))) _tmem_load_0[12] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 13))) _tmem_load_0[13] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 14))) _tmem_load_0[14] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 15))) _tmem_load_0[15] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 16))) _tmem_load_0[16] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 17))) _tmem_load_0[17] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 18))) _tmem_load_0[18] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 19))) _tmem_load_0[19] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 20))) _tmem_load_0[20] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 21))) _tmem_load_0[21] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 22))) _tmem_load_0[22] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 23))) _tmem_load_0[23] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 24))) _tmem_load_0[24] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 25))) _tmem_load_0[25] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 26))) _tmem_load_0[26] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 27))) _tmem_load_0[27] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 28))) _tmem_load_0[28] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 29))) _tmem_load_0[29] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 30))) _tmem_load_0[30] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 31))) _tmem_load_0[31] = -CAKE_INF; - uint32_t _slice_lo_mask_1; - { - int _lim_1 = valid_cols - 32; - if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } - else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); - } - } - if (!(_slice_lo_mask_1 & (1u << 0))) _tmem_load_0[32] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 1))) _tmem_load_0[33] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 2))) _tmem_load_0[34] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 3))) _tmem_load_0[35] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 4))) _tmem_load_0[36] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 5))) _tmem_load_0[37] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 6))) _tmem_load_0[38] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 7))) _tmem_load_0[39] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 8))) _tmem_load_0[40] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 9))) _tmem_load_0[41] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 10))) _tmem_load_0[42] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 11))) _tmem_load_0[43] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 12))) _tmem_load_0[44] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 13))) _tmem_load_0[45] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 14))) _tmem_load_0[46] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 15))) _tmem_load_0[47] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 16))) _tmem_load_0[48] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 17))) _tmem_load_0[49] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 18))) _tmem_load_0[50] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 19))) _tmem_load_0[51] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 20))) _tmem_load_0[52] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 21))) _tmem_load_0[53] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 22))) _tmem_load_0[54] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 23))) _tmem_load_0[55] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 24))) _tmem_load_0[56] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 25))) _tmem_load_0[57] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 26))) _tmem_load_0[58] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 27))) _tmem_load_0[59] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 28))) _tmem_load_0[60] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 29))) _tmem_load_0[61] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 30))) _tmem_load_0[62] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 31))) _tmem_load_0[63] = -CAKE_INF; - float2 _reg_reduce_max2_2 = {-CAKE_INF, -CAKE_INF}; - row_max_x32_accum(&_tmem_load_0[0], _reg_reduce_max2_2); - row_max_x32_accum(&_tmem_load_0[32], _reg_reduce_max2_2); - float _tmem_load_0_max = row_max_reduce(_reg_reduce_max2_2); - float _max_5 = max_noftz(row_max, _tmem_load_0_max); - float new_max = _max_5; - float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, new_max, 16); - float _max_6 = max_noftz(new_max, _shfl_xor_0); - new_max = _max_6; - float _fma_0 = __fmaf_rn(row_max, softmax_rescale_log2, (-new_max) * softmax_rescale_log2); - float delta = _fma_0; - float _exp2_0 = approx_exp2(delta); - float acc_scale = ((row_max > -CAKE_INF) ? _exp2_0 : 1.0f); - int stats_tile = stats_cursor + tile; - int stats_stage = stats_tile & 1; - int stats_phase = stats_tile >> 1 & 1; - mbarrier_wait(stats_empty_addr + (stats_stage) * 8, 1 - stats_phase); - float stats_pair[2]; - stats_pair[0] = row_max; - stats_pair[1] = new_max; - asm volatile( - "tcgen05.st.sync.aligned.32x32b.x2.b32" - " [%0], {%1, %2};" - :: "r"(taddr + (unsigned int)(stats_stage * 128)), "f"(stats_pair[0]), "f"(stats_pair[1])); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(stats_full_addr + (stats_stage) * 8); - float safe_max = ((new_max == -CAKE_INF) ? 0.0f : new_max); - float max_scaled = safe_max * softmax_scale_log2; - const float2 _fma_b2_3 = {softmax_scale_log2, softmax_scale_log2}; - const float2 _fma_c2_4 = {-max_scaled + 8.8073549f, -max_scaled + 8.8073549f}; - #pragma unroll - for (int _lf = 0; _lf < 32; _lf++) - fma_f32x2_inplace(&reinterpret_cast(_tmem_load_0)[_lf], _fma_b2_3, _fma_c2_4); - #pragma unroll - for (int _le = 0; _le < 64; _le++) { - _tmem_load_0[_le] = approx_exp2(_tmem_load_0[_le]); - } - float2 _reg_reduce_sum2_5 = make_float2(0.0f, 0.0f); - softmax_block_sum(&_tmem_load_0[0], &_reg_reduce_sum2_5); - softmax_block_sum(&_tmem_load_0[32], &_reg_reduce_sum2_5); - float _tmem_load_0_sum = _reg_reduce_sum2_5.x + _reg_reduce_sum2_5.y; - float block_sum = _tmem_load_0_sum; - float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, block_sum, 16); - block_sum = block_sum + _shfl_xor_1; - float _fma_1 = __fmaf_rn(row_sum, acc_scale, block_sum); - row_sum = _fma_1; - uint32_t _fp8_0[16]; - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_0[0]), "f"(_tmem_load_0[1]), - "f"(_tmem_load_0[2]), "f"(_tmem_load_0[3])); - _fp8_0[0] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_0[4]), "f"(_tmem_load_0[5]), - "f"(_tmem_load_0[6]), "f"(_tmem_load_0[7])); - _fp8_0[1] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_0[8]), "f"(_tmem_load_0[9]), - "f"(_tmem_load_0[10]), "f"(_tmem_load_0[11])); - _fp8_0[2] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_0[12]), "f"(_tmem_load_0[13]), - "f"(_tmem_load_0[14]), "f"(_tmem_load_0[15])); - _fp8_0[3] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_0[16]), "f"(_tmem_load_0[17]), - "f"(_tmem_load_0[18]), "f"(_tmem_load_0[19])); - _fp8_0[4] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_0[20]), "f"(_tmem_load_0[21]), - "f"(_tmem_load_0[22]), "f"(_tmem_load_0[23])); - _fp8_0[5] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_0[24]), "f"(_tmem_load_0[25]), - "f"(_tmem_load_0[26]), "f"(_tmem_load_0[27])); - _fp8_0[6] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_0[28]), "f"(_tmem_load_0[29]), - "f"(_tmem_load_0[30]), "f"(_tmem_load_0[31])); - _fp8_0[7] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_0[32]), "f"(_tmem_load_0[33]), - "f"(_tmem_load_0[34]), "f"(_tmem_load_0[35])); - _fp8_0[8] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_0[36]), "f"(_tmem_load_0[37]), - "f"(_tmem_load_0[38]), "f"(_tmem_load_0[39])); - _fp8_0[9] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_0[40]), "f"(_tmem_load_0[41]), - "f"(_tmem_load_0[42]), "f"(_tmem_load_0[43])); - _fp8_0[10] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_0[44]), "f"(_tmem_load_0[45]), - "f"(_tmem_load_0[46]), "f"(_tmem_load_0[47])); - _fp8_0[11] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_0[48]), "f"(_tmem_load_0[49]), - "f"(_tmem_load_0[50]), "f"(_tmem_load_0[51])); - _fp8_0[12] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_0[52]), "f"(_tmem_load_0[53]), - "f"(_tmem_load_0[54]), "f"(_tmem_load_0[55])); - _fp8_0[13] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_0[56]), "f"(_tmem_load_0[57]), - "f"(_tmem_load_0[58]), "f"(_tmem_load_0[59])); - _fp8_0[14] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_0[60]), "f"(_tmem_load_0[61]), - "f"(_tmem_load_0[62]), "f"(_tmem_load_0[63])); - _fp8_0[15] = _packed; - } - int p_col = ((score_stage != 0) ? 160 : 32); - asm volatile( - "tcgen05.st.sync.aligned.16x32bx2.x16.b32" - " [%0], 16, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" - :: "r"(taddr + (unsigned int)p_col), "r"(*reinterpret_cast(&_fp8_0[0])), "r"(*reinterpret_cast(&_fp8_0[1])), "r"(*reinterpret_cast(&_fp8_0[2])), "r"(*reinterpret_cast(&_fp8_0[3])), "r"(*reinterpret_cast(&_fp8_0[4])), "r"(*reinterpret_cast(&_fp8_0[5])), "r"(*reinterpret_cast(&_fp8_0[6])), "r"(*reinterpret_cast(&_fp8_0[7])), "r"(*reinterpret_cast(&_fp8_0[8])), "r"(*reinterpret_cast(&_fp8_0[9])), "r"(*reinterpret_cast(&_fp8_0[10])), "r"(*reinterpret_cast(&_fp8_0[11])), "r"(*reinterpret_cast(&_fp8_0[12])), "r"(*reinterpret_cast(&_fp8_0[13])), "r"(*reinterpret_cast(&_fp8_0[14])), "r"(*reinterpret_cast(&_fp8_0[15]))); - asm volatile( - "tcgen05.st.sync.aligned.16x32bx2.x16.b32" - " [%0], 16, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" - :: "r"(taddr + 1048576 + (unsigned int)p_col), "r"(*reinterpret_cast(&_fp8_0[0])), "r"(*reinterpret_cast(&_fp8_0[1])), "r"(*reinterpret_cast(&_fp8_0[2])), "r"(*reinterpret_cast(&_fp8_0[3])), "r"(*reinterpret_cast(&_fp8_0[4])), "r"(*reinterpret_cast(&_fp8_0[5])), "r"(*reinterpret_cast(&_fp8_0[6])), "r"(*reinterpret_cast(&_fp8_0[7])), "r"(*reinterpret_cast(&_fp8_0[8])), "r"(*reinterpret_cast(&_fp8_0[9])), "r"(*reinterpret_cast(&_fp8_0[10])), "r"(*reinterpret_cast(&_fp8_0[11])), "r"(*reinterpret_cast(&_fp8_0[12])), "r"(*reinterpret_cast(&_fp8_0[13])), "r"(*reinterpret_cast(&_fp8_0[14])), "r"(*reinterpret_cast(&_fp8_0[15]))); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - if (num_tiles > tile + 1) { - mbarrier_arrive(score_empty_addr + (score_stage) * 8); - } - row_max = new_max; - } - int final_stats_tile = stats_cursor + num_tiles; - int final_stats_stage = final_stats_tile & 1; - int final_stats_phase = final_stats_tile >> 1 & 1; - mbarrier_wait(stats_empty_addr + (final_stats_stage) * 8, 1 - final_stats_phase); - float final_stats[2]; - final_stats[0] = row_sum; - final_stats[1] = row_max; - asm volatile( - "tcgen05.st.sync.aligned.32x32b.x2.b32" - " [%0], {%1, %2};" - :: "r"(taddr + (unsigned int)(final_stats_stage * 128)), "f"(final_stats[0]), "f"(final_stats[1])); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(stats_full_addr + (final_stats_stage) * 8); - int tail_stats_tile = final_stats_tile + 1; - int tail_stats_stage = tail_stats_tile & 1; - int tail_stats_phase = tail_stats_tile >> 1 & 1; - mbarrier_wait(stats_empty_addr + (tail_stats_stage) * 8, 1 - tail_stats_phase); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(stats_full_addr + (tail_stats_stage) * 8); - int last_score_tile = score_cursor + num_tiles - 1; - mbarrier_arrive(score_empty_addr + (last_score_tile & 1) * 8); - stats_cursor = stats_cursor + num_tiles + 2; - score_cursor = score_cursor + num_tiles; - } - mbarrier_wait(work_full_addr + (softmax_work_stage) * 8, _phase_work_full); - uint32_t _clc_valid_3 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_3) - : "r"(work_response_addr + softmax_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_9 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_9) - : "r"(work_response_addr + softmax_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_10 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_10) - : "r"(work_response_addr + softmax_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_11 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_11) - : "r"(work_response_addr + softmax_work_stage * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(work_empty_addr + (softmax_work_stage) * 8); - softmax_work_stage += 1; - if (softmax_work_stage == 2) { softmax_work_stage = 0; _phase_work_full ^= 1; } - if (_clc_valid_3 == 0) { - break; - } - softmax_q = _clc_ctaid_9; - softmax_head = _clc_ctaid_10; - softmax_batch = _clc_ctaid_11; - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - } - } - // ---- Role: correction ---- - if (warp >= 4 && warp <= 7) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 144;"); - { // correction_main - const int warp_rank_1 = warp - 4; - const int logical_row_origin_1 = warp_rank_1 * 16; - const int my_row_1 = logical_row_origin_1 + lane % 16; - const int col_half_1 = lane / 16; - int score_cursor_1 = 0; - int stats_cursor_1 = 0; - unsigned int correction_work_stage = 0; - unsigned int correction_q = blockIdx.x; - unsigned int correction_head = blockIdx.y; - unsigned int correction_batch = blockIdx.z; - unsigned int _phase_work_full_1 = 0; - #pragma unroll 1 - for (unsigned int _work_1 = 0; _work_1 < total_work_items; _work_1++) { - int q_begin_1 = cum_seq_lens_q[(int)correction_batch]; - int q_end_1 = cum_seq_lens_q[(int)(correction_batch + 1)]; - int q_local_1 = (int)correction_q; - int work_valid_1 = q_local_1 < q_end_1 - q_begin_1; - asm volatile("griddepcontrol.wait;" ::: "memory"); - float softmax_scale_log2_1 = bmm1_scale[0] * 1.4426950408889634f; - float output_scale = bmm2_scale[0]; - if (work_valid_1 != 0) { - int query_idx_1 = q_begin_1 + q_local_1; - int head_base = (int)correction_head * 64; - int head_idx = head_base + my_row_1; - int _max_7 = ((sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) : (0)); - int _min_4 = ((_max_7) < (sparse_topk) ? (_max_7) : (sparse_topk)); - int active_topk_1 = _min_4; - int num_tiles_1 = (active_topk_1 + 128 - 1) / 128; - int first_stats_tile = stats_cursor_1; - int first_stats_stage = first_stats_tile & 1; - int first_stats_phase = first_stats_tile >> 1 & 1; - mbarrier_wait(stats_full_addr + (first_stats_stage) * 8, first_stats_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - mbarrier_arrive(stats_empty_addr + (first_stats_stage) * 8); - #pragma unroll 1 - for (int tile_1 = 1; tile_1 < num_tiles_1; tile_1++) { - int stats_tile_1 = stats_cursor_1 + tile_1; - int stats_stage_1 = stats_tile_1 & 1; - int stats_phase_1 = stats_tile_1 >> 1 & 1; - mbarrier_wait(stats_full_addr + (stats_stage_1) * 8, stats_phase_1); - asm volatile("tcgen05.fence::after_thread_sync;"); - float _tmem_load_1[2]; - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x2.b32" - " {%0, %1}, [%2];" - : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]) - : "r"(taddr + (unsigned int)(stats_stage_1 * 128))); - asm volatile("tcgen05.wait::ld.sync.aligned;"); - mbarrier_arrive(stats_empty_addr + (stats_stage_1) * 8); - int output_tile = score_cursor_1 + tile_1 - 1; - mbarrier_wait(o_full_addr, output_tile & 1); - asm volatile("tcgen05.fence::after_thread_sync;"); - float max_delta = _tmem_load_1[0] - _tmem_load_1[1]; - float _exp2_1 = approx_exp2(max_delta * softmax_scale_log2_1); - float acc_scale_1 = ((max_delta != 0.0f) ? _exp2_1 : 1.0f); - int _vote_0 = __any_sync(0xFFFFFFFF, acc_scale_1 < 1.0f); - int any_rescale = _vote_0; - if (any_rescale != 0) { - #pragma unroll - for (int output_half = 0; output_half < 2; output_half++) { - int row_bank = ((output_half != 0) ? 1048576 : 0); - #pragma unroll - for (int col = 0; col < 128; col += 64) { - float _tmem_load_2[64]; - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 128;" - : "=r"(*reinterpret_cast(&_tmem_load_2[0])), "=r"(*reinterpret_cast(&_tmem_load_2[1])), "=r"(*reinterpret_cast(&_tmem_load_2[2])), "=r"(*reinterpret_cast(&_tmem_load_2[3])), "=r"(*reinterpret_cast(&_tmem_load_2[4])), "=r"(*reinterpret_cast(&_tmem_load_2[5])), "=r"(*reinterpret_cast(&_tmem_load_2[6])), "=r"(*reinterpret_cast(&_tmem_load_2[7])), "=r"(*reinterpret_cast(&_tmem_load_2[8])), "=r"(*reinterpret_cast(&_tmem_load_2[9])), "=r"(*reinterpret_cast(&_tmem_load_2[10])), "=r"(*reinterpret_cast(&_tmem_load_2[11])), "=r"(*reinterpret_cast(&_tmem_load_2[12])), "=r"(*reinterpret_cast(&_tmem_load_2[13])), "=r"(*reinterpret_cast(&_tmem_load_2[14])), "=r"(*reinterpret_cast(&_tmem_load_2[15])), "=r"(*reinterpret_cast(&_tmem_load_2[16])), "=r"(*reinterpret_cast(&_tmem_load_2[17])), "=r"(*reinterpret_cast(&_tmem_load_2[18])), "=r"(*reinterpret_cast(&_tmem_load_2[19])), "=r"(*reinterpret_cast(&_tmem_load_2[20])), "=r"(*reinterpret_cast(&_tmem_load_2[21])), "=r"(*reinterpret_cast(&_tmem_load_2[22])), "=r"(*reinterpret_cast(&_tmem_load_2[23])), "=r"(*reinterpret_cast(&_tmem_load_2[24])), "=r"(*reinterpret_cast(&_tmem_load_2[25])), "=r"(*reinterpret_cast(&_tmem_load_2[26])), "=r"(*reinterpret_cast(&_tmem_load_2[27])), "=r"(*reinterpret_cast(&_tmem_load_2[28])), "=r"(*reinterpret_cast(&_tmem_load_2[29])), "=r"(*reinterpret_cast(&_tmem_load_2[30])), "=r"(*reinterpret_cast(&_tmem_load_2[31])) - : "r"(taddr + 256 + (unsigned int)row_bank + (unsigned int)col)); - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 128;" - : "=r"(*reinterpret_cast(&_tmem_load_2[32])), "=r"(*reinterpret_cast(&_tmem_load_2[33])), "=r"(*reinterpret_cast(&_tmem_load_2[34])), "=r"(*reinterpret_cast(&_tmem_load_2[35])), "=r"(*reinterpret_cast(&_tmem_load_2[36])), "=r"(*reinterpret_cast(&_tmem_load_2[37])), "=r"(*reinterpret_cast(&_tmem_load_2[38])), "=r"(*reinterpret_cast(&_tmem_load_2[39])), "=r"(*reinterpret_cast(&_tmem_load_2[40])), "=r"(*reinterpret_cast(&_tmem_load_2[41])), "=r"(*reinterpret_cast(&_tmem_load_2[42])), "=r"(*reinterpret_cast(&_tmem_load_2[43])), "=r"(*reinterpret_cast(&_tmem_load_2[44])), "=r"(*reinterpret_cast(&_tmem_load_2[45])), "=r"(*reinterpret_cast(&_tmem_load_2[46])), "=r"(*reinterpret_cast(&_tmem_load_2[47])), "=r"(*reinterpret_cast(&_tmem_load_2[48])), "=r"(*reinterpret_cast(&_tmem_load_2[49])), "=r"(*reinterpret_cast(&_tmem_load_2[50])), "=r"(*reinterpret_cast(&_tmem_load_2[51])), "=r"(*reinterpret_cast(&_tmem_load_2[52])), "=r"(*reinterpret_cast(&_tmem_load_2[53])), "=r"(*reinterpret_cast(&_tmem_load_2[54])), "=r"(*reinterpret_cast(&_tmem_load_2[55])), "=r"(*reinterpret_cast(&_tmem_load_2[56])), "=r"(*reinterpret_cast(&_tmem_load_2[57])), "=r"(*reinterpret_cast(&_tmem_load_2[58])), "=r"(*reinterpret_cast(&_tmem_load_2[59])), "=r"(*reinterpret_cast(&_tmem_load_2[60])), "=r"(*reinterpret_cast(&_tmem_load_2[61])), "=r"(*reinterpret_cast(&_tmem_load_2[62])), "=r"(*reinterpret_cast(&_tmem_load_2[63])) - : "r"(taddr + 256 + (unsigned int)row_bank + (unsigned int)col + 32)); - const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; - #pragma unroll - for (int _ls = 0; _ls < 32; _ls++) - mul_f32x2_inplace(&reinterpret_cast(_tmem_load_2)[_ls], _scale2_0); - asm volatile( - "tcgen05.st.sync.aligned.16x32bx2.x64.b32" - " [%0], 128, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64};" - :: "r"(taddr + 256 + (unsigned int)row_bank + (unsigned int)col), "r"(*reinterpret_cast(&_tmem_load_2[0])), "r"(*reinterpret_cast(&_tmem_load_2[1])), "r"(*reinterpret_cast(&_tmem_load_2[2])), "r"(*reinterpret_cast(&_tmem_load_2[3])), "r"(*reinterpret_cast(&_tmem_load_2[4])), "r"(*reinterpret_cast(&_tmem_load_2[5])), "r"(*reinterpret_cast(&_tmem_load_2[6])), "r"(*reinterpret_cast(&_tmem_load_2[7])), "r"(*reinterpret_cast(&_tmem_load_2[8])), "r"(*reinterpret_cast(&_tmem_load_2[9])), "r"(*reinterpret_cast(&_tmem_load_2[10])), "r"(*reinterpret_cast(&_tmem_load_2[11])), "r"(*reinterpret_cast(&_tmem_load_2[12])), "r"(*reinterpret_cast(&_tmem_load_2[13])), "r"(*reinterpret_cast(&_tmem_load_2[14])), "r"(*reinterpret_cast(&_tmem_load_2[15])), "r"(*reinterpret_cast(&_tmem_load_2[16])), "r"(*reinterpret_cast(&_tmem_load_2[17])), "r"(*reinterpret_cast(&_tmem_load_2[18])), "r"(*reinterpret_cast(&_tmem_load_2[19])), "r"(*reinterpret_cast(&_tmem_load_2[20])), "r"(*reinterpret_cast(&_tmem_load_2[21])), "r"(*reinterpret_cast(&_tmem_load_2[22])), "r"(*reinterpret_cast(&_tmem_load_2[23])), "r"(*reinterpret_cast(&_tmem_load_2[24])), "r"(*reinterpret_cast(&_tmem_load_2[25])), "r"(*reinterpret_cast(&_tmem_load_2[26])), "r"(*reinterpret_cast(&_tmem_load_2[27])), "r"(*reinterpret_cast(&_tmem_load_2[28])), "r"(*reinterpret_cast(&_tmem_load_2[29])), "r"(*reinterpret_cast(&_tmem_load_2[30])), "r"(*reinterpret_cast(&_tmem_load_2[31])), "r"(*reinterpret_cast(&_tmem_load_2[32])), "r"(*reinterpret_cast(&_tmem_load_2[33])), "r"(*reinterpret_cast(&_tmem_load_2[34])), "r"(*reinterpret_cast(&_tmem_load_2[35])), "r"(*reinterpret_cast(&_tmem_load_2[36])), "r"(*reinterpret_cast(&_tmem_load_2[37])), "r"(*reinterpret_cast(&_tmem_load_2[38])), "r"(*reinterpret_cast(&_tmem_load_2[39])), "r"(*reinterpret_cast(&_tmem_load_2[40])), "r"(*reinterpret_cast(&_tmem_load_2[41])), "r"(*reinterpret_cast(&_tmem_load_2[42])), "r"(*reinterpret_cast(&_tmem_load_2[43])), "r"(*reinterpret_cast(&_tmem_load_2[44])), "r"(*reinterpret_cast(&_tmem_load_2[45])), "r"(*reinterpret_cast(&_tmem_load_2[46])), "r"(*reinterpret_cast(&_tmem_load_2[47])), "r"(*reinterpret_cast(&_tmem_load_2[48])), "r"(*reinterpret_cast(&_tmem_load_2[49])), "r"(*reinterpret_cast(&_tmem_load_2[50])), "r"(*reinterpret_cast(&_tmem_load_2[51])), "r"(*reinterpret_cast(&_tmem_load_2[52])), "r"(*reinterpret_cast(&_tmem_load_2[53])), "r"(*reinterpret_cast(&_tmem_load_2[54])), "r"(*reinterpret_cast(&_tmem_load_2[55])), "r"(*reinterpret_cast(&_tmem_load_2[56])), "r"(*reinterpret_cast(&_tmem_load_2[57])), "r"(*reinterpret_cast(&_tmem_load_2[58])), "r"(*reinterpret_cast(&_tmem_load_2[59])), "r"(*reinterpret_cast(&_tmem_load_2[60])), "r"(*reinterpret_cast(&_tmem_load_2[61])), "r"(*reinterpret_cast(&_tmem_load_2[62])), "r"(*reinterpret_cast(&_tmem_load_2[63]))); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - } - } - } - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(o_empty_addr); - } - int final_stats_tile_1 = stats_cursor_1 + num_tiles_1; - int final_stats_stage_1 = final_stats_tile_1 & 1; - int final_stats_phase_1 = final_stats_tile_1 >> 1 & 1; - mbarrier_wait(stats_full_addr + (final_stats_stage_1) * 8, final_stats_phase_1); - asm volatile("tcgen05.fence::after_thread_sync;"); - float _tmem_load_3[2]; - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x2.b32" - " {%0, %1}, [%2];" - : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]) - : "r"(taddr + (unsigned int)(final_stats_stage_1 * 128))); - asm volatile("tcgen05.wait::ld.sync.aligned;"); - mbarrier_arrive(stats_empty_addr + (final_stats_stage_1) * 8); - int tail_stats_tile_1 = final_stats_tile_1 + 1; - int tail_stats_stage_1 = tail_stats_tile_1 & 1; - int tail_stats_phase_1 = tail_stats_tile_1 >> 1 & 1; - mbarrier_wait(stats_full_addr + (tail_stats_stage_1) * 8, tail_stats_phase_1); - asm volatile("tcgen05.fence::after_thread_sync;"); - mbarrier_arrive(stats_empty_addr + (tail_stats_stage_1) * 8); - int last_output_tile = score_cursor_1 + num_tiles_1 - 1; - mbarrier_wait(o_full_addr, last_output_tile & 1); - asm volatile("tcgen05.fence::after_thread_sync;"); - float total_sum = _tmem_load_3[0]; - float final_max = _tmem_load_3[1]; - int output_in_bounds = head_idx < num_heads; - if (((int)(has_sinks != 0) & output_in_bounds) != 0) { - float _exp2_2 = approx_exp2(sinks[head_idx] * 1.4426950408889634f - final_max * softmax_scale_log2_1); - total_sum = total_sum + _exp2_2 * 448.0f; - } - float _rcp_0 = approx_rcp(total_sum); - float inv_sum = ((total_sum > 0.0f) ? output_scale * _rcp_0 : 0.0f); - int output_base = (query_idx_1 * num_heads + head_idx) * 512; - #pragma unroll - for (int output_half_1 = 0; output_half_1 < 2; output_half_1++) { - int row_bank_1 = ((output_half_1 != 0) ? 1048576 : 0); - int col_base = col_half_1 * 128; - #pragma unroll 2 - for (int col_1 = 0; col_1 < 128; col_1 += 8) { - float _tmem_load_4[8]; - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x8.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8], 128;" - : "=r"(*reinterpret_cast(&_tmem_load_4[0])), "=r"(*reinterpret_cast(&_tmem_load_4[1])), "=r"(*reinterpret_cast(&_tmem_load_4[2])), "=r"(*reinterpret_cast(&_tmem_load_4[3])), "=r"(*reinterpret_cast(&_tmem_load_4[4])), "=r"(*reinterpret_cast(&_tmem_load_4[5])), "=r"(*reinterpret_cast(&_tmem_load_4[6])), "=r"(*reinterpret_cast(&_tmem_load_4[7])) - : "r"(taddr + 256 + (unsigned int)row_bank_1 + (unsigned int)col_1)); - const float2 _scale2_1 = {inv_sum, inv_sum}; - #pragma unroll - for (int _ls = 0; _ls < 4; _ls++) - mul_f32x2_inplace(&reinterpret_cast(_tmem_load_4)[_ls], _scale2_1); - uint32_t _tmem_load_4_bf16[4]; - #pragma unroll - for (int _lp = 0; _lp < 4; _lp++) { - __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_4[_lp*2 + 0], _tmem_load_4[_lp*2+1 + 0])); - _tmem_load_4_bf16[_lp] = *(uint32_t*)&_bf2; - } - if (output_in_bounds != 0) { - reinterpret_cast(O + (output_base + output_half_1 * 256 + col_base + col_1))[0] = reinterpret_cast(_tmem_load_4_bf16)[0]; - } - } - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(o_empty_addr); - stats_cursor_1 = stats_cursor_1 + num_tiles_1 + 2; - score_cursor_1 = score_cursor_1 + num_tiles_1; - } - mbarrier_wait(work_full_addr + (correction_work_stage) * 8, _phase_work_full_1); - uint32_t _clc_valid_4 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_4) - : "r"(work_response_addr + correction_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_12 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_12) - : "r"(work_response_addr + correction_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_13 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_13) - : "r"(work_response_addr + correction_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_14 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_14) - : "r"(work_response_addr + correction_work_stage * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(work_empty_addr + (correction_work_stage) * 8); - correction_work_stage += 1; - if (correction_work_stage == 2) { correction_work_stage = 0; _phase_work_full_1 ^= 1; } - if (_clc_valid_4 == 0) { - break; - } - correction_q = _clc_ctaid_12; - correction_head = _clc_ctaid_13; - correction_batch = _clc_ctaid_14; - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile("barrier.sync 7, 128;" ::: "memory"); - if (warp == 4) { - int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); - asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); - } - } - } - // ---- Role: mma ---- - if (warp == 8) { - { // mma_main - CakePipelineState<4> _kv_cons_state_state{0u, 0u, 0u}; - unsigned int q_cons_stage = 0; - int score_prod_stage = 0; - int score_prod_phase = 0; - int p_cons_stage = 0; - int p_cons_phase = 0; - int o_prod_phase = 0; - unsigned int mma_work_stage = 0; - unsigned int mma_q = blockIdx.x; - unsigned int mma_head = blockIdx.y; - unsigned int mma_batch = blockIdx.z; - #pragma unroll - for (int score_lookahead = 0; score_lookahead < 2; score_lookahead++) { - int lookahead_stage = score_lookahead & 1; - int lookahead_phase = score_lookahead >> 1 & 1; - mbarrier_wait(score_empty_addr + (lookahead_stage) * 8, 1 - lookahead_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - } - unsigned int _phase_q_full = 0; - unsigned int _phase_work_full_2 = 0; - #pragma unroll 1 - for (unsigned int _work_2 = 0; _work_2 < total_work_items; _work_2++) { - int q_begin_2 = cum_seq_lens_q[(int)mma_batch]; - int q_end_2 = cum_seq_lens_q[(int)(mma_batch + 1)]; - int q_local_2 = (int)mma_q; - int work_valid_2 = q_local_2 < q_end_2 - q_begin_2; - if (work_valid_2 != 0) { - int query_idx_2 = q_begin_2 + q_local_2; - int _max_3 = ((sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) : (0)); - int _min_2 = ((_max_3) < (sparse_topk) ? (_max_3) : (sparse_topk)); - int active_topk_2 = _min_2; - int num_tiles_2 = (active_topk_2 + 128 - 1) / 128; - if (num_tiles_2 > 0) { - mbarrier_wait(q_full_addr + (q_cons_stage) * 8, _phase_q_full); - asm volatile("tcgen05.fence::after_thread_sync;"); - int score_col_1 = ((score_prod_stage != 0) ? 128 : 0); - mbarrier_wait(kv_full_addr + (_kv_cons_state_state.index()) * 8, _kv_cons_state_state.phase()); - asm volatile("tcgen05.fence::after_thread_sync;"); - int _mma_a_lo_0 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (q_cons_stage) * 2048); - int _mma_b_lo_0 = make_warp_uniform((((smem_k_addr) >> 4) & 0x3FFF) + (_kv_cons_state_state.index()) * 2048); - { - uint64_t _mma_ss_a_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_0); - uint64_t _mma_ss_b_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_0); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69206032, 0); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 506U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 1018U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69206032, 1); - } - } - elect_commit(kv_empty_addr + (_kv_cons_state_state.index()) * 8); - ++_kv_cons_state_state; - mbarrier_wait(kv_full_addr + (_kv_cons_state_state.index()) * 8, _kv_cons_state_state.phase()); - asm volatile("tcgen05.fence::after_thread_sync;"); - int _mma_a_lo_1 = make_warp_uniform((((smem_q_addr + 16384) >> 4) & 0x3FFF) + (q_cons_stage) * 2048); - int _mma_b_lo_1 = make_warp_uniform((((smem_k_addr) >> 4) & 0x3FFF) + (_kv_cons_state_state.index()) * 2048); - { - uint64_t _mma_ss_a_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_1); - uint64_t _mma_ss_b_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_1); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 506U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 1018U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69206032, 1); - } - } - elect_commit(score_full_addr + (score_prod_stage) * 8); - elect_commit(kv_empty_addr + (_kv_cons_state_state.index()) * 8); - ++_kv_cons_state_state; - score_prod_stage ^= 1; - if (score_prod_stage == 0) { - score_prod_phase ^= 1; - } - int first_pv = 1; - #pragma unroll 1 - for (int _steady = 0; _steady < num_tiles_2 - 1; _steady++) { - score_col_1 = ((score_prod_stage != 0) ? 128 : 0); - mbarrier_wait(kv_full_addr + (_kv_cons_state_state.index()) * 8, _kv_cons_state_state.phase()); - asm volatile("tcgen05.fence::after_thread_sync;"); - int _mma_a_lo_2 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (q_cons_stage) * 2048); - int _mma_b_lo_2 = make_warp_uniform((((smem_k_addr) >> 4) & 0x3FFF) + (_kv_cons_state_state.index()) * 2048); - { - uint64_t _mma_ss_a_desc_2 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_2); - uint64_t _mma_ss_b_desc_2 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_2); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 69206032, 0); - } - incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_2, 506U); - incr_smem_desc_lo(_mma_ss_b_desc_2, 1018U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 69206032, 1); - } - } - elect_commit(kv_empty_addr + (_kv_cons_state_state.index()) * 8); - ++_kv_cons_state_state; - mbarrier_wait(kv_full_addr + (_kv_cons_state_state.index()) * 8, _kv_cons_state_state.phase()); - asm volatile("tcgen05.fence::after_thread_sync;"); - int _mma_a_lo_3 = make_warp_uniform((((smem_q_addr + 16384) >> 4) & 0x3FFF) + (q_cons_stage) * 2048); - int _mma_b_lo_3 = make_warp_uniform((((smem_k_addr) >> 4) & 0x3FFF) + (_kv_cons_state_state.index()) * 2048); - { - uint64_t _mma_ss_a_desc_3 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_3); - uint64_t _mma_ss_b_desc_3 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_3); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_3, 506U); - incr_smem_desc_lo(_mma_ss_b_desc_3, 1018U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 69206032, 1); - } - } - elect_commit(score_full_addr + (score_prod_stage) * 8); - elect_commit(kv_empty_addr + (_kv_cons_state_state.index()) * 8); - ++_kv_cons_state_state; - score_prod_stage ^= 1; - if (score_prod_stage == 0) { - score_prod_phase ^= 1; - } - mbarrier_wait(score_empty_addr + (score_prod_stage) * 8, 1 - score_prod_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - mbarrier_wait(o_empty_addr, 1 - o_prod_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - #pragma unroll - for (int v_half = 0; v_half < 2; v_half++) { - mbarrier_wait(kv_full_addr + (_kv_cons_state_state.index()) * 8, _kv_cons_state_state.phase()); - asm volatile("tcgen05.fence::after_thread_sync;"); - int row_bank_2 = ((v_half != 0) ? 1048576 : 0); - int p_col_1 = ((p_cons_stage != 0) ? 160 : 32) + row_bank_2; - int _mma_b_lo_4 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (_kv_cons_state_state.index()) * 2048); - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 dhi, blo, ta, id;\n\t" - ".reg .b64 db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "" - "mov.b32 dhi, 0x40004040;\n\t" - "mov.b32 id, 71368720;\n\t" - "mov.b32 ta, %2;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p0;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" - "}\n" - :: "r"((tmem_tmem + (256 + row_bank_2))), "r"(_mma_b_lo_4), "r"(tmem_tmem + p_col_1), "r"(((first_pv) ? 0 : 1))); - elect_commit(kv_empty_addr + (_kv_cons_state_state.index()) * 8); - ++_kv_cons_state_state; - } - first_pv = 0; - elect_commit(o_full_addr); - p_cons_stage ^= 1; - if (p_cons_stage == 0) { - p_cons_phase ^= 1; - } - o_prod_phase ^= 1; - } - elect_commit(q_empty_addr + (q_cons_stage) * 8); - q_cons_stage += 1; - if (q_cons_stage == 2) { q_cons_stage = 0; _phase_q_full ^= 1; } - int tail_score_stage = score_prod_stage ^ 1; - int tail_score_phase = score_prod_phase ^ score_prod_stage; - mbarrier_wait(score_empty_addr + (tail_score_stage) * 8, 1 - tail_score_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - mbarrier_wait(o_empty_addr, 1 - o_prod_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - #pragma unroll - for (int v_half_1 = 0; v_half_1 < 2; v_half_1++) { - mbarrier_wait(kv_full_addr + (_kv_cons_state_state.index()) * 8, _kv_cons_state_state.phase()); - asm volatile("tcgen05.fence::after_thread_sync;"); - int row_bank_3 = ((v_half_1 != 0) ? 1048576 : 0); - int p_col_2 = ((p_cons_stage != 0) ? 160 : 32) + row_bank_3; - int _mma_b_lo_5 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (_kv_cons_state_state.index()) * 2048); - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 dhi, blo, ta, id;\n\t" - ".reg .b64 db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "" - "mov.b32 dhi, 0x40004040;\n\t" - "mov.b32 id, 71368720;\n\t" - "mov.b32 ta, %2;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p0;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" - "}\n" - :: "r"((tmem_tmem + (256 + row_bank_3))), "r"(_mma_b_lo_5), "r"(tmem_tmem + p_col_2), "r"(((first_pv) ? 0 : 1))); - elect_commit(kv_empty_addr + (_kv_cons_state_state.index()) * 8); - ++_kv_cons_state_state; - } - elect_commit(o_full_addr); - p_cons_stage ^= 1; - if (p_cons_stage == 0) { - p_cons_phase ^= 1; - } - o_prod_phase ^= 1; - } - } - mbarrier_wait(work_full_addr + (mma_work_stage) * 8, _phase_work_full_2); - uint32_t _clc_valid_2 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_2) - : "r"(work_response_addr + mma_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_6 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_6) - : "r"(work_response_addr + mma_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_7 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_7) - : "r"(work_response_addr + mma_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_8 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_8) - : "r"(work_response_addr + mma_work_stage * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(work_empty_addr + (mma_work_stage) * 8); - mma_work_stage += 1; - if (mma_work_stage == 2) { mma_work_stage = 0; _phase_work_full_2 ^= 1; } - if (_clc_valid_2 == 0) { - break; - } - mma_q = _clc_ctaid_6; - mma_head = _clc_ctaid_7; - mma_batch = _clc_ctaid_8; - } - } - } - // ---- Role: page_loader ---- - if (warp == 9) { - { // page_loader_main - unsigned int page_prod_stage = 0; - unsigned int page_prod_phase = 1; - unsigned int page_work_stage = 0; - unsigned int page_q = blockIdx.x; - unsigned int page_head = blockIdx.y; - unsigned int page_batch = blockIdx.z; - unsigned int _phase_work_full_3 = 0; - #pragma unroll 1 - for (unsigned int _work_3 = 0; _work_3 < total_work_items; _work_3++) { - int q_begin_3 = cum_seq_lens_q[(int)page_batch]; - int q_end_3 = cum_seq_lens_q[(int)(page_batch + 1)]; - int q_local_3 = (int)page_q; - int work_valid_3 = q_local_3 < q_end_3 - q_begin_3; - if (work_valid_3 != 0) { - int query_idx_3 = q_begin_3 + q_local_3; - int _max_0 = ((sparse_topk_lens[query_idx_3] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_3] + sparse_topk_lens_offset) : (0)); - int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); - int active_topk_3 = _min_0; - int num_tiles_3 = (active_topk_3 + 128 - 1) / 128; - int num_page_passes = (num_tiles_3 + 1) / 2; - int* swa_row = swa_indices + (query_idx_3 * swa_index_stride); - int* compressed_row = compressed_indices + (query_idx_3 * compressed_index_stride); - int _max_1 = (((active_topk_3 - 1) / 4 * 4) > (0) ? ((active_topk_3 - 1) / 4 * 4) : (0)); - int last_vec_base = _max_1; - #pragma unroll 1 - for (int page_pass = 0; page_pass < num_page_passes; page_pass++) { - #pragma unroll - for (int _duplicate = 0; _duplicate < 2; _duplicate++) { - mbarrier_wait(page_empty_addr + (page_prod_stage) * 8, page_prod_phase); - int page_dst = smem_page_offsets_addr + page_prod_stage * 1024; - #pragma unroll - for (int page_half = 0; page_half < 2; page_half++) { - int page_offset = page_pass * 256 + page_half * 128 + lane * 4; - int clamped_offset = ((page_offset < active_topk_3) ? page_offset : last_vec_base); - int* page_src = ((clamped_offset < 128) ? swa_row + clamped_offset : compressed_row + (clamped_offset - 128)); - asm volatile("cp.async.cg.shared::cta.global.L2::128B [%0], [%1], 16;" - :: "r"(page_dst + (page_half * 128 + lane * 4) * 4), "l"(page_src)); - } - asm volatile("cp.async.commit_group;"); - asm volatile("cp.async.wait_group 0;"); - mbarrier_arrive(page_full_addr + (page_prod_stage) * 8); - page_prod_stage += 1; - if (page_prod_stage == 6) { page_prod_stage = 0; page_prod_phase ^= 1; } - } - } - } - mbarrier_wait(work_full_addr + (page_work_stage) * 8, _phase_work_full_3); - uint32_t _clc_valid_0 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_0) - : "r"(work_response_addr + page_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_0 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_0) - : "r"(work_response_addr + page_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_1 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_1) - : "r"(work_response_addr + page_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_2 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_2) - : "r"(work_response_addr + page_work_stage * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(work_empty_addr + (page_work_stage) * 8); - page_work_stage += 1; - if (page_work_stage == 2) { page_work_stage = 0; _phase_work_full_3 ^= 1; } - if (_clc_valid_0 == 0) { - break; - } - page_q = _clc_ctaid_0; - page_head = _clc_ctaid_1; - page_batch = _clc_ctaid_2; - } - } - } - // ---- Role: scheduler ---- - if (warp == 10) { - { // scheduler_main - unsigned int scheduler_work_stage = 0; - unsigned int scheduler_throttle_stage = 0; - unsigned int _phase_throttle_full = 0; - unsigned int _phase_work_empty = 1; - unsigned int _phase_work_full_4 = 0; - #pragma unroll 1 - for (unsigned int _work_4 = 0; _work_4 < total_work_items; _work_4++) { - mbarrier_wait(throttle_full_addr + (scheduler_throttle_stage) * 8, _phase_throttle_full); - mbarrier_arrive(throttle_empty_addr + (scheduler_throttle_stage) * 8); - scheduler_throttle_stage += 1; - if (scheduler_throttle_stage == 2) { scheduler_throttle_stage = 0; _phase_throttle_full ^= 1; } - if (elect_sync()) { - mbarrier_wait(work_empty_addr + (scheduler_work_stage) * 8, _phase_work_empty); - mbarrier_arrive_expect_tx(work_full_addr + (scheduler_work_stage) * 8, 16); - asm volatile( - "fence.proxy.async.shared::cta;\n\t" - "clusterlaunchcontrol.try_cancel.async.shared::cta" - ".mbarrier::complete_tx::bytes.b128" - " [%0], [%1];" - :: "r"(work_response_addr + scheduler_work_stage * 16 + 0 * 16), "r"(work_full_addr + scheduler_work_stage * 8) - : "memory"); - } - mbarrier_wait(work_full_addr + (scheduler_work_stage) * 8, _phase_work_full_4); - uint32_t _clc_valid_5 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_5) - : "r"(work_response_addr + scheduler_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_15 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_15) - : "r"(work_response_addr + scheduler_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_16 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_16) - : "r"(work_response_addr + scheduler_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_17 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_17) - : "r"(work_response_addr + scheduler_work_stage * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(work_empty_addr + (scheduler_work_stage) * 8); - scheduler_work_stage += 1; - if (scheduler_work_stage == 2) { scheduler_work_stage = 0; _phase_work_empty ^= 1; _phase_work_full_4 ^= 1; } - if (_clc_valid_5 == 0) { - break; - } - } - } - } - // ---- Role: padding ---- - if (warp == 11) { - { // padding_main - unsigned int padding_work_stage = 0; - unsigned int _phase_work_full_5 = 0; - #pragma unroll 1 - for (unsigned int _work_5 = 0; _work_5 < total_work_items; _work_5++) { - mbarrier_wait(work_full_addr + (padding_work_stage) * 8, _phase_work_full_5); - uint32_t _clc_valid_6 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_6) - : "r"(work_response_addr + padding_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_18 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_18) - : "r"(work_response_addr + padding_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_19 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_19) - : "r"(work_response_addr + padding_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_20 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_20) - : "r"(work_response_addr + padding_work_stage * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(work_empty_addr + (padding_work_stage) * 8); - padding_work_stage += 1; - if (padding_work_stage == 2) { padding_work_stage = 0; _phase_work_full_5 ^= 1; } - if (_clc_valid_6 == 0) { - break; - } - } - } - } - // ---- Role: load ---- - if (warp >= 12 && warp <= 15) { - asm volatile("setmaxnreg.dec.sync.aligned.u32 72;"); - { // load_main - const int load_rank = warp - 12; - unsigned int q_prod_stage = 0; - unsigned int kv_prod_stage = 0; - unsigned int page_k_stage = 0; - unsigned int page_k_phase = 0; - unsigned int page_v_stage = 1; - unsigned int page_v_phase = 0; - unsigned int load_work_stage = 0; - unsigned int throttle_stage = 0; - unsigned int load_q = blockIdx.x; - unsigned int load_head = blockIdx.y; - unsigned int load_batch = blockIdx.z; - unsigned int _phase_throttle_empty = 1; - unsigned int _phase_q_empty = 1; - unsigned int _phase_kv_empty = 1; - unsigned int _phase_work_full_6 = 0; - #pragma unroll 1 - for (unsigned int _work_6 = 0; _work_6 < total_work_items; _work_6++) { - mbarrier_wait(throttle_empty_addr + (throttle_stage) * 8, _phase_throttle_empty); - mbarrier_arrive(throttle_full_addr + (throttle_stage) * 8); - throttle_stage += 1; - if (throttle_stage == 2) { throttle_stage = 0; _phase_throttle_empty ^= 1; } - int q_begin_4 = cum_seq_lens_q[(int)load_batch]; - int q_end_4 = cum_seq_lens_q[(int)(load_batch + 1)]; - int q_local_4 = (int)load_q; - int work_valid_4 = q_local_4 < q_end_4 - q_begin_4; - if (work_valid_4 != 0) { - int query_idx_4 = q_begin_4 + q_local_4; - int _max_2 = ((sparse_topk_lens[query_idx_4] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_4] + sparse_topk_lens_offset) : (0)); - int _min_1 = ((_max_2) < (sparse_topk) ? (_max_2) : (sparse_topk)); - int active_topk_4 = _min_1; - int num_tiles_4 = (active_topk_4 + 128 - 1) / 128; - mbarrier_wait(q_empty_addr + (q_prod_stage) * 8, _phase_q_empty); - if (load_rank == 0) { - if (elect_sync()) { - mbarrier_arrive_expect_tx(q_full_addr + (q_prod_stage) * 8, 32768); - } - } - asm volatile("griddepcontrol.wait;" ::: "memory"); - if (load_rank == 0) { - if (elect_sync()) { - #pragma unroll - for (int q_chunk = 0; q_chunk < 4; q_chunk++) { - asm volatile( - "cp.async.bulk.tensor.4d.shared::cluster.global.tile.mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5}], [%6];" - :: "r"(smem_q_addr + q_prod_stage * 32768 + (unsigned int)(q_chunk * 8192)), "l"((&tmap_q)), "r"(q_chunk * 128), "r"(0), "r"((int)load_head), "r"(query_idx_4), "r"(q_full_addr + (q_prod_stage) * 8) : "memory"); - } - } - } - q_prod_stage += 1; - if (q_prod_stage == 2) { q_prod_stage = 0; _phase_q_empty ^= 1; } - #pragma unroll 1 - for (int tile_2 = 0; tile_2 < num_tiles_4; tile_2++) { - if ((tile_2 & 1) == 0) { - mbarrier_wait(page_full_addr + (page_k_stage) * 8, page_k_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - } - int page_k_base = smem_page_offsets_addr + page_k_stage * 1024; - #pragma unroll - for (int k_half = 0; k_half < 2; k_half++) { - mbarrier_wait(kv_empty_addr + (kv_prod_stage) * 8, _phase_kv_empty); - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (kv_prod_stage) * 8, 8192); - if (tile_2 == 0) { - #pragma unroll - for (int local_group = 0; local_group < 8; local_group++) { - int group = load_rank * 8 + local_group; - int group_offset = (tile_2 & 1) * 128 + group * 4; - int rows[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&rows[0])), "=r"(*reinterpret_cast(&rows[(0) + 1])), "=r"(*reinterpret_cast(&rows[(0) + 2])), "=r"(*reinterpret_cast(&rows[(0) + 3])) - : "r"(page_k_base + group_offset * 4)); - int k_dst = smem_k_addr + kv_prod_stage * 32768; - tma_gather4_gmem2smem(k_dst + group * 512, (&tmap_swa_kv), k_half * 256, rows[0], rows[1], rows[2], rows[3], kv_full_addr + (kv_prod_stage) * 8); - tma_gather4_gmem2smem(k_dst + 16384 + group * 512, (&tmap_swa_kv), k_half * 256 + 128, rows[0], rows[1], rows[2], rows[3], kv_full_addr + (kv_prod_stage) * 8); - } - } else { - #pragma unroll - for (int local_group_1 = 0; local_group_1 < 8; local_group_1++) { - int group_1 = load_rank * 8 + local_group_1; - int group_offset_1 = (tile_2 & 1) * 128 + group_1 * 4; - int rows_1[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&rows_1[0])), "=r"(*reinterpret_cast(&rows_1[(0) + 1])), "=r"(*reinterpret_cast(&rows_1[(0) + 2])), "=r"(*reinterpret_cast(&rows_1[(0) + 3])) - : "r"(page_k_base + group_offset_1 * 4)); - int k_dst_1 = smem_k_addr + kv_prod_stage * 32768; - tma_gather4_gmem2smem(k_dst_1 + group_1 * 512, (&tmap_compressed_kv), k_half * 256, rows_1[0], rows_1[1], rows_1[2], rows_1[3], kv_full_addr + (kv_prod_stage) * 8); - tma_gather4_gmem2smem(k_dst_1 + 16384 + group_1 * 512, (&tmap_compressed_kv), k_half * 256 + 128, rows_1[0], rows_1[1], rows_1[2], rows_1[3], kv_full_addr + (kv_prod_stage) * 8); - } - } - } - kv_prod_stage += 1; - if (kv_prod_stage == 4) { kv_prod_stage = 0; _phase_kv_empty ^= 1; } - } - if ((tile_2 & 1) != 0 || tile_2 == num_tiles_4 - 1) { - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(page_empty_addr + (page_k_stage) * 8); - page_k_stage += 1; - if (page_k_stage == 6) { page_k_stage = 0; page_k_phase ^= 1; } - page_k_stage += 1; - if (page_k_stage == 6) { page_k_stage = 0; page_k_phase ^= 1; } - } - if (tile_2 > 0) { - int prev_tile = tile_2 - 1; - if ((prev_tile & 1) == 0) { - mbarrier_wait(page_full_addr + (page_v_stage) * 8, page_v_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - } - int page_v_base = smem_page_offsets_addr + page_v_stage * 1024; - #pragma unroll - for (int v_half_2 = 0; v_half_2 < 2; v_half_2++) { - mbarrier_wait(kv_empty_addr + (kv_prod_stage) * 8, _phase_kv_empty); - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (kv_prod_stage) * 8, 8192); - if (prev_tile == 0) { - #pragma unroll - for (int local_group_2 = 0; local_group_2 < 8; local_group_2++) { - int group_2 = load_rank * 8 + local_group_2; - int group_offset_2 = (prev_tile & 1) * 128 + group_2 * 4; - int rows_2[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&rows_2[0])), "=r"(*reinterpret_cast(&rows_2[(0) + 1])), "=r"(*reinterpret_cast(&rows_2[(0) + 2])), "=r"(*reinterpret_cast(&rows_2[(0) + 3])) - : "r"(page_v_base + group_offset_2 * 4)); - int v_dst = smem_v_addr + kv_prod_stage * 32768; - tma_gather4_gmem2smem(v_dst + group_2 * 512, (&tmap_swa_kv), v_half_2 * 256, rows_2[0], rows_2[1], rows_2[2], rows_2[3], kv_full_addr + (kv_prod_stage) * 8); - tma_gather4_gmem2smem(v_dst + 16384 + group_2 * 512, (&tmap_swa_kv), v_half_2 * 256 + 128, rows_2[0], rows_2[1], rows_2[2], rows_2[3], kv_full_addr + (kv_prod_stage) * 8); - } - } else { - #pragma unroll - for (int local_group_3 = 0; local_group_3 < 8; local_group_3++) { - int group_3 = load_rank * 8 + local_group_3; - int group_offset_3 = (prev_tile & 1) * 128 + group_3 * 4; - int rows_3[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&rows_3[0])), "=r"(*reinterpret_cast(&rows_3[(0) + 1])), "=r"(*reinterpret_cast(&rows_3[(0) + 2])), "=r"(*reinterpret_cast(&rows_3[(0) + 3])) - : "r"(page_v_base + group_offset_3 * 4)); - int v_dst_1 = smem_v_addr + kv_prod_stage * 32768; - tma_gather4_gmem2smem(v_dst_1 + group_3 * 512, (&tmap_compressed_kv), v_half_2 * 256, rows_3[0], rows_3[1], rows_3[2], rows_3[3], kv_full_addr + (kv_prod_stage) * 8); - tma_gather4_gmem2smem(v_dst_1 + 16384 + group_3 * 512, (&tmap_compressed_kv), v_half_2 * 256 + 128, rows_3[0], rows_3[1], rows_3[2], rows_3[3], kv_full_addr + (kv_prod_stage) * 8); - } - } - } - kv_prod_stage += 1; - if (kv_prod_stage == 4) { kv_prod_stage = 0; _phase_kv_empty ^= 1; } - } - if ((prev_tile & 1) != 0) { - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(page_empty_addr + (page_v_stage) * 8); - page_v_stage += 1; - if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } - page_v_stage += 1; - if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } - } - } - } - int last_tile = num_tiles_4 - 1; - mbarrier_wait(page_full_addr + (page_v_stage) * 8, page_v_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - int page_v_base_1 = smem_page_offsets_addr + page_v_stage * 1024; - #pragma unroll - for (int v_half_3 = 0; v_half_3 < 2; v_half_3++) { - mbarrier_wait(kv_empty_addr + (kv_prod_stage) * 8, _phase_kv_empty); - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (kv_prod_stage) * 8, 8192); - if (last_tile == 0) { - #pragma unroll - for (int local_group_4 = 0; local_group_4 < 8; local_group_4++) { - int group_4 = load_rank * 8 + local_group_4; - int group_offset_4 = (last_tile & 1) * 128 + group_4 * 4; - int rows_4[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&rows_4[0])), "=r"(*reinterpret_cast(&rows_4[(0) + 1])), "=r"(*reinterpret_cast(&rows_4[(0) + 2])), "=r"(*reinterpret_cast(&rows_4[(0) + 3])) - : "r"(page_v_base_1 + group_offset_4 * 4)); - int v_dst_2 = smem_v_addr + kv_prod_stage * 32768; - tma_gather4_gmem2smem(v_dst_2 + group_4 * 512, (&tmap_swa_kv), v_half_3 * 256, rows_4[0], rows_4[1], rows_4[2], rows_4[3], kv_full_addr + (kv_prod_stage) * 8); - tma_gather4_gmem2smem(v_dst_2 + 16384 + group_4 * 512, (&tmap_swa_kv), v_half_3 * 256 + 128, rows_4[0], rows_4[1], rows_4[2], rows_4[3], kv_full_addr + (kv_prod_stage) * 8); - } - } else { - #pragma unroll - for (int local_group_5 = 0; local_group_5 < 8; local_group_5++) { - int group_5 = load_rank * 8 + local_group_5; - int group_offset_5 = (last_tile & 1) * 128 + group_5 * 4; - int rows_5[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&rows_5[0])), "=r"(*reinterpret_cast(&rows_5[(0) + 1])), "=r"(*reinterpret_cast(&rows_5[(0) + 2])), "=r"(*reinterpret_cast(&rows_5[(0) + 3])) - : "r"(page_v_base_1 + group_offset_5 * 4)); - int v_dst_3 = smem_v_addr + kv_prod_stage * 32768; - tma_gather4_gmem2smem(v_dst_3 + group_5 * 512, (&tmap_compressed_kv), v_half_3 * 256, rows_5[0], rows_5[1], rows_5[2], rows_5[3], kv_full_addr + (kv_prod_stage) * 8); - tma_gather4_gmem2smem(v_dst_3 + 16384 + group_5 * 512, (&tmap_compressed_kv), v_half_3 * 256 + 128, rows_5[0], rows_5[1], rows_5[2], rows_5[3], kv_full_addr + (kv_prod_stage) * 8); - } - } - } - kv_prod_stage += 1; - if (kv_prod_stage == 4) { kv_prod_stage = 0; _phase_kv_empty ^= 1; } - } - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(page_empty_addr + (page_v_stage) * 8); - page_v_stage += 1; - if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } - page_v_stage += 1; - if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } - } - mbarrier_wait(work_full_addr + (load_work_stage) * 8, _phase_work_full_6); - uint32_t _clc_valid_1 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_1) - : "r"(work_response_addr + load_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_3 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_3) - : "r"(work_response_addr + load_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_4 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_4) - : "r"(work_response_addr + load_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_5 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_5) - : "r"(work_response_addr + load_work_stage * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(work_empty_addr + (load_work_stage) * 8); - load_work_stage += 1; - if (load_work_stage == 2) { load_work_stage = 0; _phase_work_full_6 ^= 1; } - if (_clc_valid_1 == 0) { - break; - } - load_q = _clc_ctaid_3; - load_head = _clc_ctaid_4; - load_batch = _clc_ctaid_5; - } - } - } - - // Cleanup -} - -} // extern "C" diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_75827bd38b443de66f67_binding.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_75827bd38b443de66f67_binding.cu deleted file mode 100644 index 7cd6edd23c7..00000000000 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_75827bd38b443de66f67_binding.cu +++ /dev/null @@ -1,445 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -// Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_75827bd38b443de66f67'. -#include -#include -#include -#include - -#include "tvm_ffi_utils.h" - -#include -#include -#include -#include -#include -#include -#include - -extern "C" __global__ void kernel_cake_dsv4_75827bd38b443de66f67(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items); - - -namespace cake_host_shim_fefc853e0bc1ea14 { - -using tvm::ffi::Optional; -using tvm::ffi::TensorView; - -class ScopedCudaDevice { - public: - explicit ScopedCudaDevice(int device_id) { - cudaError_t error = cudaGetDevice(&previous_device_); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed before host-shim launch: cudaError=" - << static_cast(error); - if (previous_device_ != device_id) { - error = cudaSetDevice(device_id); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaSetDevice failed before host-shim launch for cuda:" - << device_id << ": cudaError=" << static_cast(error); - restore_ = true; - } - } - - ScopedCudaDevice(const ScopedCudaDevice&) = delete; - ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; - - ~ScopedCudaDevice() noexcept { - if (restore_) { - (void)cudaSetDevice(previous_device_); - } - } - - private: - int previous_device_ = -1; - bool restore_ = false; -}; - -inline int64_t CakeDeviceMultiprocessorCount(int device_id) { - constexpr int kMaxCachedCudaDevices = 64; - TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) - << "physical-SM-count cache does not cover cuda:" << device_id; - static std::atomic count_by_device[kMaxCachedCudaDevices]{}; - int cached = count_by_device[device_id].load(std::memory_order_acquire); - if (cached > 0) return cached; - - int count = 0; - cudaError_t error = cudaDeviceGetAttribute( - &count, cudaDevAttrMultiProcessorCount, device_id); - TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) - << "querying multiProcessorCount failed for cuda:" << device_id - << ": cudaError=" << static_cast(error) << ", count=" << count; - // Concurrent first launches may repeat the immutable device query, but all - // publication is atomic and every later hot-path lookup is one acquire load. - count_by_device[device_id].store(count, std::memory_order_release); - return count; -} - -inline void CheckCudaTensor(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) - << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; -} - -inline void CheckSameCudaDevice( - const TensorView& t, - const TensorView& reference, - const char* name, - const char* reference_name) { - TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) - << name << " must be on the same CUDA device as " << reference_name - << ": got cuda:" << t.device().device_id - << " versus cuda:" << reference.device().device_id; -} - -inline void CheckCurrentCudaDevice( - const TensorView& reference, - const char* reference_name) { - int current_device = -1; - cudaError_t error = cudaGetDevice(¤t_device); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed while validating " << reference_name - << ": cudaError=" << static_cast(error); - TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) - << "current CUDA device must match " << reference_name - << ": current=cuda:" << current_device - << ", tensor=cuda:" << reference.device().device_id; -} - -inline void CheckContiguous(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; -} - -inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { - DLDataType d = t.dtype(); - TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) - << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits - << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits - << ", lanes=" << (int)d.lanes << ")"; -} - -// A logical axis.outer(trailing) folds every source dim above the trailing -// dimensions. Shape products are independent of physical strides, so verify -// the leading dimensions form one dense row-major chain instead of inventing -// a "folded stride". The descriptor reads its exact adjacent physical step -// separately through stride[-(trailing + 1)]. -inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { - TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) - << name << " cannot fold leading dimensions above " << trailing - << " trailing dims from ndim=" << t.ndim(); - int outer_last = t.ndim() - trailing - 1; - if (outer_last <= 0) { - return; - } - int64_t step = t.stride(outer_last); - TVM_FFI_CHECK(step > 0, ValueError) - << name << " physical strides must be positive"; - int64_t expected = step; - for (int axis = outer_last - 1; axis >= 0; --axis) { - expected *= t.size(axis + 1); - if (t.size(axis) > 1) { - TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) - << name << " leading dims are not physically foldable above " << trailing - << " trailing dims: stride(" << axis << ")=" << t.stride(axis) - << ", expected " << expected; - } - } -} - -// 4D TMA descriptor for buffer 'tmap_q' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 3, ValueError) - << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - int64_t d2 = t.size(t.ndim() - 2); - int64_t d3 = t.size(t.ndim() - 3); - TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) - << "TMA source 'tmap_q' trailing dims must be positive"; - TVM_FFI_CHECK(d1 % 128 == 0, ValueError) - << "TMA source 'tmap_q' extent " << d1 - << " must divide exactly by " << 128; - uint64_t global_dim[4] = {(uint64_t)(128), (uint64_t)(d2), (uint64_t)((d1 / 128)), (uint64_t)(d3)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; - TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) - << "TMA box (128, 128, 1, 1) exceeds resolved global dims for 'tmap_q'"; - int64_t carrier_stride_0 = d1; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; - int64_t carrier_stride_1 = 128; - TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; - TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; - TVM_FFI_CHECK((carrier_stride_1 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; - int64_t carrier_stride_2 = (d2 * d1); - TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; - TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; - TVM_FFI_CHECK((carrier_stride_2 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; - uint64_t global_strides[3] = { - (uint64_t)((carrier_stride_0 * 8) / 8), - (uint64_t)((carrier_stride_1 * 8) / 8), - (uint64_t)((carrier_stride_2 * 8) / 8), - }; - uint32_t box_dim[4] = {128u, 128u, 1u, 1u}; - uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_swa_kv' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_swa_kv' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; - TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (128, 1) exceeds resolved global dims for 'tmap_swa_kv'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 8) / 8), - }; - uint32_t box_dim[2] = {128u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_compressed_kv' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_compressed_kv(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_compressed_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_compressed_kv' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_compressed_kv' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_compressed_kv"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_compressed_kv' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved a non-positive global dim"; - TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (128, 1) exceeds resolved global dims for 'tmap_compressed_kv'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 8) / 8), - }; - uint32_t box_dim[2] = {128u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_kv') failed: CUresult=" << (int)r; - return tm; -} - -void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_compressed_kv, TensorView arg_O, TensorView arg_partial_lse, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_num_query_tokens, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, int64_t grid_x, int64_t grid_y, int64_t grid_z) { - DLDevice dev = arg_tmap_q.device(); - ScopedCudaDevice device_guard(dev.device_id); - CheckCudaTensor(arg_tmap_q, "tmap_q"); - CheckDtype(arg_tmap_q, "tmap_q", 1, 8, 1); - CheckContiguous(arg_tmap_q, "tmap_q"); - CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); - CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 1, 8, 1); - CheckCudaTensor(arg_tmap_compressed_kv, "tmap_compressed_kv"); - CheckDtype(arg_tmap_compressed_kv, "tmap_compressed_kv", 1, 8, 1); - CheckCudaTensor(arg_O, "O"); - CheckDtype(arg_O, "O", 4, 16, 1); - CheckContiguous(arg_O, "O"); - CheckCudaTensor(arg_partial_lse, "partial_lse"); - CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); - CheckContiguous(arg_partial_lse, "partial_lse"); - CheckCudaTensor(arg_swa_indices, "swa_indices"); - CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); - CheckContiguous(arg_swa_indices, "swa_indices"); - CheckCudaTensor(arg_compressed_indices, "compressed_indices"); - CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); - CheckContiguous(arg_compressed_indices, "compressed_indices"); - CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); - CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckCudaTensor(arg_sinks, "sinks"); - CheckDtype(arg_sinks, "sinks", 2, 32, 1); - CheckContiguous(arg_sinks, "sinks"); - CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); - CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); - CheckContiguous(arg_bmm1_scale, "bmm1_scale"); - CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); - CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); - CheckContiguous(arg_bmm2_scale, "bmm2_scale"); - TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) - << "scalar 'num_heads' value " << arg_num_heads - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) - << "scalar 'swa_index_stride' value " << arg_swa_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) - << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) - << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_num_query_tokens >= -2147483648LL && arg_num_query_tokens <= 2147483647LL, ValueError) - << "scalar 'num_query_tokens' value " << arg_num_query_tokens - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) - << "scalar 'sparse_topk' value " << arg_sparse_topk - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) - << "scalar 'has_sinks' value " << arg_has_sinks - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) - << "scalar 'total_work_items' value " << arg_total_work_items - << " is outside i32 range [-2147483648, 2147483647]"; - CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); - CheckSameCudaDevice(arg_tmap_compressed_kv, arg_tmap_q, "tmap_compressed_kv", "tmap_q"); - CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); - CheckSameCudaDevice(arg_partial_lse, arg_tmap_q, "partial_lse", "tmap_q"); - CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); - CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); - CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); - CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); - CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); - CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); - TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) - << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y - << ", " << grid_z << ")"; - - cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); - CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); - CUtensorMap p_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); - CUtensorMap p_tmap_compressed_kv = EncodeTma_tmap_compressed_kv(arg_tmap_compressed_kv); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); - float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); - int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); - int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); - int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); - float* p_sinks = static_cast(arg_sinks.data_ptr()); - float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); - float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); - int32_t v_num_heads = (int32_t)arg_num_heads; - int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; - int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; - int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; - int32_t v_num_query_tokens = (int32_t)arg_num_query_tokens; - int32_t v_sparse_topk = (int32_t)arg_sparse_topk; - int32_t v_has_sinks = (int32_t)arg_has_sinks; - int32_t v_total_work_items = (int32_t)arg_total_work_items; - void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_tmap_compressed_kv, &p_O, &p_partial_lse, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_num_query_tokens, &v_sparse_topk, &v_has_sinks, &v_total_work_items}; - - static std::mutex smem_mu; - static auto* smem_status_by_device = - new std::unordered_map(); - cudaError_t smem_status = cudaSuccess; - { - std::lock_guard lock(smem_mu); - auto it = smem_status_by_device->find(dev.device_id); - if (it == smem_status_by_device->end()) { - int current_device = -1; - smem_status = cudaGetDevice(¤t_device); - if (smem_status == cudaSuccess && current_device != dev.device_id) { - smem_status = cudaErrorInvalidDevice; - } - if (smem_status == cudaSuccess) { - smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_75827bd38b443de66f67), - cudaFuncAttributeMaxDynamicSharedMemorySize, - 205824); - } - smem_status_by_device->emplace(dev.device_id, smem_status); - } else { - smem_status = it->second; - } - } - TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_75827bd38b443de66f67 failed: " - << cudaGetErrorString(smem_status); - dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); - dim3 block(512u, 1u, 1u); - - cudaLaunchConfig_t config{}; - config.gridDim = grid; - config.blockDim = block; - config.dynamicSmemBytes = 205824u; - config.stream = stream; - cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_75827bd38b443de66f67), kargs); - TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_75827bd38b443de66f67 failed: " - << cudaGetErrorString(launch_status); -} - -} // namespace cake_host_shim_fefc853e0bc1ea14 - -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_fefc853e0bc1ea14::Run); diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_75827bd38b443de66f67_kernel.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_75827bd38b443de66f67_kernel.cu deleted file mode 100644 index fb67b80f0b2..00000000000 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_75827bd38b443de66f67_kernel.cu +++ /dev/null @@ -1,2392 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -typedef signed char int8_t; -typedef unsigned char uint8_t; -typedef unsigned short uint16_t; -typedef unsigned int uint32_t; -#if defined(__CUDACC_RTC__) -typedef unsigned long long uint64_t; -#else -typedef unsigned long uint64_t; -#endif -static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); -typedef signed int int32_t; -typedef short int int16_t; -struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; -struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; -static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); -static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); - -#if defined(__CUDACC_RTC__) -typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; -#else -#include -#endif - -static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); -static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); -#include -#include - -#define CAKE_INF CUDART_INF_F -#define TMEM_NCOLS 512 -#define TMEM_TMEM_SCRATCH_OFFSET 0 -#define NUM_K_PIPE_STAGES 1 -#define NUM_V_PIPE_STAGES 1 -#define NUM_KV_PIPE_STAGES 8 -#define NUM_INDEX_PIPE_STAGES 1 -#define SMEM_SMEM_Q_OFF 1024 -#define SMEM_SMEM_Q_STAGE_BYTES 16384 -#define SMEM_SMEM_Q_STRIDE 16384 -#define SMEM_SMEM_KV_OFF 66560 -#define SMEM_SMEM_KV_STAGE_BYTES 16384 -#define SMEM_SMEM_KV_STRIDE 16384 -#define SMEM_SMEM_V_OFF 66560 -#define SMEM_SMEM_V_STAGE_BYTES 16384 -#define SMEM_SMEM_V_STRIDE 16384 -#define SMEM_SMEM_Q_FULL_OFF 1024 -#define SMEM_SMEM_Q_FULL_STAGE_BYTES 32768 -#define SMEM_SMEM_Q_FULL_STRIDE 32768 -#define SMEM_SMEM_K_FULL_OFF 66560 -#define SMEM_SMEM_K_FULL_STAGE_BYTES 32768 -#define SMEM_SMEM_K_FULL_STRIDE 32768 -#define SMEM_SMEM_V_FULL_OFF 132096 -#define SMEM_SMEM_V_FULL_STAGE_BYTES 32768 -#define SMEM_SMEM_V_FULL_STRIDE 32768 -#define SMEM_SMEM_STATS_MAX_OFF 197632 -#define SMEM_SMEM_STATS_MAX_STAGE_BYTES 1024 -#define SMEM_SMEM_STATS_MAX_STRIDE 1024 -#define SMEM_SMEM_STATS_SUM_OFF 198656 -#define SMEM_SMEM_STATS_SUM_STAGE_BYTES 512 -#define SMEM_SMEM_STATS_SUM_STRIDE 512 -#define SMEM_SMEM_STATS_FINAL_MAX_OFF 199168 -#define SMEM_SMEM_STATS_FINAL_MAX_STAGE_BYTES 512 -#define SMEM_SMEM_STATS_FINAL_MAX_STRIDE 512 -#define SMEM_SMEM_SOFTMAX_WARP_PAIR_EXCHANGE_OFF 199680 -#define SMEM_SMEM_SOFTMAX_WARP_PAIR_EXCHANGE_STAGE_BYTES 1024 -#define SMEM_SMEM_SOFTMAX_WARP_PAIR_EXCHANGE_STRIDE 1024 -#define SMEM_SMEM_CORR_WARP_PAIR_EXCHANGE_OFF 200704 -#define SMEM_SMEM_CORR_WARP_PAIR_EXCHANGE_STAGE_BYTES 512 -#define SMEM_SMEM_CORR_WARP_PAIR_EXCHANGE_STRIDE 512 -#define SMEM_SMEM_SPARSE_INDICES_OFF 201216 -#define SMEM_SMEM_SPARSE_INDICES_STAGE_BYTES 4608 -#define SMEM_SMEM_SPARSE_INDICES_STRIDE 4608 -#define SMEM_SMEM_P_FP8_OFF 1024 -#define SMEM_SMEM_P_FP8_STAGE_BYTES 8192 -#define SMEM_SMEM_P_FP8_STRIDE 8192 -#define SMEM_TOTAL 205824 -#define THREADS 512 - -#include - -__device__ __forceinline__ uint32_t elect_sync() { - uint32_t pred = 0; - asm volatile( - "{\n\t" - ".reg .pred %%px;\n\t" - "elect.sync _|%%px, %1;\n\t" - "@%%px mov.s32 %0, 1;\n\t" - "}\n" - : "+r"(pred) - : "r"(0xFFFFFFFF)); - return pred; -} - - -__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { - asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" - :: "r"(mbar_addr), "r"(count) : "memory"); -} - -__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { - asm volatile("mbarrier.init.b64 [%0], %1;" - :: "l"(mbar_addr), "r"(count) : "memory"); -} - - -__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - - -__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { - uint32_t ticks = 0x989680; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE;\n\t" - "bra.uni LAB_WAIT;\n\t" - "DONE:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); -} - -// Source-faithful relaxed CTA wait used only by a typed protocol that does -// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. -__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, 10000000;\n\t" - "@P1 bra.uni DONE_RELAXED;\n\t" - "bra.uni LAB_WAIT_RELAXED;\n\t" - "DONE_RELAXED:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Exact source ports may request the PTX suspendTimeHint operand explicitly. -// The hint is expressed in nanoseconds and is kept separate from the canonical -// no-hint CTA helper so unrelated schedules retain their existing retry path. -__device__ __forceinline__ void mbarrier_wait_suspend( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_SUSPEND:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_SUSPEND;\n\t" - "bra.uni LAB_WAIT_SUSPEND;\n\t" - "DONE_SUSPEND:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE_CLUSTER;\n\t" - "bra.uni LAB_WAIT_CLUSTER;\n\t" - "DONE_CLUSTER:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - ".reg .u32 WAIT_ADDR;\n\t" - "mov.u32 WAIT_ADDR, %0;\n\t" - "LAB_WAIT_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [WAIT_ADDR], %1, %2;\n\t" - "@P1 bra.uni DONE_HINT;\n\t" - "bra.uni LAB_WAIT_HINT;\n\t" - "DONE_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -// Exact unqualified CTA wait used by source schedules whose PTX intentionally -// omits the acquire qualifier while retaining a typed suspendTimeHint operand. -__device__ __forceinline__ void mbarrier_wait_relaxed_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED_HINT:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra DONE_RELAXED_HINT;\n\t" - "bra LAB_WAIT_RELAXED_HINT;\n\t" - "DONE_RELAXED_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); -} - -__device__ __forceinline__ void mbarrier_wait_cluster_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" - "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" - "DONE_CLUSTER_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_suspend( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait_cluster(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); - } -} - - -__device__ __forceinline__ void tcgen05_mma_f8f6f4( - int taddr, uint64_t a_desc, uint64_t b_desc, - uint32_t i_desc, int enable_input_d) { - asm volatile( - "{\n\t" - ".reg .pred p;\n\t" - "setp.ne.b32 p, %4, 0;\n\t" - "tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], %1, %2, %3, p;\n\t" - "}\n" - :: "r"(taddr), "l"(a_desc), "l"(b_desc), - "r"(i_desc), "r"(enable_input_d) - : "memory"); -} - - -__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { - return (x & 0x3FFFFULL) >> 4ULL; -} - - -union MmaSmemDesc { - uint64_t u64; - uint32_t u32[2]; -}; - -__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { - MmaSmemDesc tmp; - tmp.u64 = smem_desc; - tmp.u32[0] += offset; - smem_desc = tmp.u64; -} - - -__device__ __forceinline__ void elect_commit(int mbar_addr) { - asm volatile( - "{\n\t" - ".reg .pred leader;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" - ".shared::cluster.b64 [%0];\n\t" - "}\n" - :: "r"(mbar_addr)); -} - - -__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { - asm volatile( - "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" - :: "r"(mbar_addr), "r"(bytes) : "memory"); -} - - -__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7," - " %8, %9, %10, %11, %12, %13, %14, %15," - " %16, %17, %18, %19, %20, %21, %22, %23," - " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), - "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), - "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), - "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), - "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), - "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), - "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x16.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7," - " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), - "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), - "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ void tmem_ld_x4(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x4.b32" - " {%0, %1, %2, %3}, [%4];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { - asm volatile( - "tcgen05.st.sync.aligned.32x32b.x16.b32" - " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," - " %9, %10, %11, %12, %13, %14, %15, %16};" - :: "r"(tmem_addr), - "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), - "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), - "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), - "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); -} - - -__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { - asm volatile( - "tcgen05.st.sync.aligned.32x32b.x4.b32" - " [%0], {%1, %2, %3, %4};" - :: "r"(tmem_addr), - "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); -} - - -__device__ __forceinline__ float approx_exp2(float x) { - float y; - asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float approx_rcp(float x) { - float y; - asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float max_noftz(float a, float b) { - float c; - asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); - return c; -} - - -__device__ __forceinline__ float warp_reduce_max(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); - return val; -} - - -__device__ __forceinline__ float warp_reduce_sum(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val += __shfl_xor_sync(0xFFFFFFFF, val, offset); - return val; -} - - -__device__ __forceinline__ float row_max_reduce(float2 acc) { - return max_noftz(acc.x, acc.y); -} - - -__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { - #pragma unroll - for (int j = 0; j < 16; j++) { - if (j % 2 == 0) - acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); - else - acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); - } -} - - -__device__ __forceinline__ float2 ex2_emulation_f32x2_value(float2 value) { - const float c0 = 1.0f, c1 = 0.695146143436431884765625f; - const float c2 = 0.227564394474029541015625f, c3 = 0.077119089663028717041015625f; - const float magic = 12582912.0f; - float x0 = max_noftz(value.x, -127.0f), x1 = max_noftz(value.y, -127.0f); - float2 xc2 = make_float2(x0, x1), magic2 = make_float2(magic, magic); - float2 xr2; - asm("add.rm.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xr2) - : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&magic2)); - float2 c3_2 = make_float2(c3, c3), c2_2 = make_float2(c2, c2); - float2 c1_2 = make_float2(c1, c1), c0_2 = make_float2(c0, c0); - float2 xrb2, xfrac2; - asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xrb2) - : "l"(*(unsigned long long*)&xr2), "l"(*(unsigned long long*)&magic2)); - asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xfrac2) - : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&xrb2)); - float2 poly2; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&c3_2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c2_2)); - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c1_2)); - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c0_2)); - int x0r_i, x1r_i, p0_i, p1_i; - asm("mov.b64 {%0, %1}, %2;" : "=r"(x0r_i), "=r"(x1r_i) : "l"(*(unsigned long long*)&xr2)); - asm("mov.b64 {%0, %1}, %2;" : "=r"(p0_i), "=r"(p1_i) : "l"(*(unsigned long long*)&poly2)); - float r0, r1; - asm("mov.b32 %0, %1;" : "=f"(r0) : "r"((x0r_i << 23) + p0_i)); - asm("mov.b32 %0, %1;" : "=f"(r1) : "r"((x1r_i << 23) + p1_i)); - return make_float2(r0, r1); -} - -__device__ __forceinline__ void ex2_emulation_f32x2(float* x0_ptr, float* x1_ptr) { - float2 result = ex2_emulation_f32x2_value(make_float2(*x0_ptr, *x1_ptr)); - *x0_ptr = result.x; *x1_ptr = result.y; -} - -__device__ __forceinline__ void softmax_frag_exp2_cast( - float* sv, uint32_t* pv, int use_emu) -{ - #pragma unroll - for (int j = 0; j < 16; j++) { - if (use_emu && j >= 12) - ex2_emulation_f32x2(&sv[j*2], &sv[j*2+1]); - else { - sv[j*2] = approx_exp2(sv[j*2]); - sv[j*2+1] = approx_exp2(sv[j*2+1]); - } - } - #pragma unroll - for (int j = 0; j < 16; j++) { - __nv_bfloat162 bf = __float22bfloat162_rn({sv[j*2], sv[j*2+1]}); - pv[j] = reinterpret_cast(bf); - } -} - - - -__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { - const float2* sv2 = reinterpret_cast(sv); - #pragma unroll - for (int j = 0; j < 16; j++) { - asm("add.f32x2 %0, %1, %2;" - : "+l"(reinterpret_cast(*acc)) - : "l"(reinterpret_cast(*acc)), - "l"(reinterpret_cast(sv2[j]))); - } -} - - -__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { - asm("mul.rn.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { - asm("mul.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { - asm("add.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { - asm("add.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { - asm("sub.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { - asm("sub.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("add.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { - float2 r; - asm("sub.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("sub.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ void fma_scale_x32( - float* sv, const float2* scale2, const float2* neg_max2) -{ - float2* sv_2 = reinterpret_cast(sv); - #pragma unroll - for (int j = 0; j < 16; j++) - fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); -} - -__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("mul.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - - -// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) - -__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("add.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("add.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("add.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("add.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("add.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("add.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("add.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - - -__device__ __forceinline__ void fence_async_shared() { - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); -} - - -__device__ __forceinline__ uint64_t make_smem_desc(int addr) { - const int SBO = 1024; - return desc_encode(addr) - | (desc_encode(SBO) << 32ULL) - | (1ULL << 46ULL) - | (2ULL << 61ULL); -} - - -__device__ __forceinline__ void tma_2d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3}], [%4];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_4d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.4d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5}], [%6];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr) { - // Canonical .shared::cta form for non-multicast gather4, matching - // trtllm-gen / cuda_ptx and the PTX ISA qualifier order - // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, - // .shared::cluster is reserved for the multicast variant (ctaMask). - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem_cta2( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr) { - // Canonical .shared::cta form; see tma_gather4_gmem2smem above. - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" - ".mbarrier::complete_tx::bytes.cta_group::2" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem_mc( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr, unsigned short cta_mask) { - // Multicast variant: the PTX grammar ties the .shared::cluster - // destination to .multicast::cluster + ctaMask (cf. cuda_ptx / - // SM100_TMA_LOAD_MULTICAST_2D_GATHER4 in CUTLASS). - asm volatile( - "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" - ".mbarrier::complete_tx::bytes.multicast::cluster" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr), "h"(cta_mask) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem_mc_cta2( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr, unsigned short cta_mask) { - // Multicast + cta_group::2 variant; see tma_gather4_gmem2smem_mc. - asm volatile( - "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" - ".mbarrier::complete_tx::bytes.multicast::cluster.cta_group::2" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr), "h"(cta_mask) : "memory"); -} - - -__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { - asm volatile( - "tcgen05.commit.cta_group::1.mbarrier::arrive::one" - ".shared::cluster.b64 [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tmem_st_x8_u32(int addr, const uint32_t* src) { - asm volatile( - "tcgen05.st.sync.aligned.32x32b.x8.b32" - " [%0], {%1,%2,%3,%4,%5,%6,%7,%8};" - :: "r"(addr), - "r"(src[0]), "r"(src[1]), "r"(src[2]), "r"(src[3]), - "r"(src[4]), "r"(src[5]), "r"(src[6]), "r"(src[7])); -} - - -__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { - uint32_t result; - asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" - : "=r"(result) : "r"(val)); - return result; -} - -extern "C" { - -__global__ __launch_bounds__(512, 1) void -kernel_cake_dsv4_75827bd38b443de66f67(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items) -{ - const int tid = threadIdx.x; - const int warp = make_warp_uniform(tid / 32); - const int lane = tid % 32; - - extern __shared__ __align__(1024) char smem_raw[]; - int smem; - asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); - smem = make_warp_uniform(smem); - - const int mbar_base = smem; - #define q_full_addr (mbar_base + 0) - #define q_empty_addr (mbar_base + 8) - #define k_full_addr (mbar_base + 16) - #define k_empty_addr (mbar_base + 24) - #define v_full_addr (mbar_base + 32) - #define v_empty_addr (mbar_base + 40) - #define kv_full_addr (mbar_base + 48) - #define kv_empty_addr (mbar_base + 112) - #define index_full_addr (mbar_base + 176) - #define index_empty_addr (mbar_base + 184) - #define s_full_addr (mbar_base + 192) - #define p_full_addr (mbar_base + 208) - #define corr_done_addr (mbar_base + 224) - #define stats_addr (mbar_base + 240) - #define sum_ready_addr (mbar_base + 256) - #define o_done_addr (mbar_base + 264) - #define pv_done_addr (mbar_base + 272) - #define s_seeded_addr (mbar_base + 280) - #define q_pair_ready_addr (mbar_base + 288) - #define kv_pair_ready_addr (mbar_base + 296) - #define pv_pair_ready_addr (mbar_base + 360) - #define tmem_dealloc_addr (mbar_base + 376) - #define tmem_dealloc_peer_addr (mbar_base + 384) - - const int bid = blockIdx.x; - const int num_bids = gridDim.x; - - const int cta_rank = 0; - - // Kernel setup ops - uint8_t* smem_q = reinterpret_cast(smem_raw + 1024); - const int smem_q_addr = smem + 1024; - uint8_t* smem_kv = reinterpret_cast(smem_raw + 66560); - const int smem_kv_addr = smem + 66560; - uint8_t* smem_v = reinterpret_cast(smem_raw + 66560); - const int smem_v_addr = smem + 66560; - uint8_t* smem_q_full = reinterpret_cast(smem_raw + 1024); - const int smem_q_full_addr = smem + 1024; - uint8_t* smem_k_full = reinterpret_cast(smem_raw + 66560); - const int smem_k_full_addr = smem + 66560; - uint8_t* smem_v_full = reinterpret_cast(smem_raw + 132096); - const int smem_v_full_addr = smem + 132096; - float* smem_stats_max = reinterpret_cast(smem_raw + 197632); - const int smem_stats_max_addr = smem + 197632; - float* smem_stats_sum = reinterpret_cast(smem_raw + 198656); - const int smem_stats_sum_addr = smem + 198656; - float* smem_stats_final_max = reinterpret_cast(smem_raw + 199168); - const int smem_stats_final_max_addr = smem + 199168; - float* smem_softmax_warp_pair_exchange = reinterpret_cast(smem_raw + 199680); - const int smem_softmax_warp_pair_exchange_addr = smem + 199680; - float* smem_corr_warp_pair_exchange = reinterpret_cast(smem_raw + 200704); - const int smem_corr_warp_pair_exchange_addr = smem + 200704; - int* smem_sparse_indices = reinterpret_cast(smem_raw + 201216); - const int smem_sparse_indices_addr = smem + 201216; - uint8_t* smem_p_fp8 = reinterpret_cast(smem_raw + 1024); - const int smem_p_fp8_addr = smem + 1024; - - // Mbarrier init (23 pipeline groups, 0 ordered-sequence groups, 49 barriers) - // Mbarriers at smem_raw[0..392) - - if (warp == 0) { - uint32_t leader = elect_sync(); - if (leader) { - // q_full: 1 barriers, init_count=1 - mbarrier_init(smem + 0, 1); - // q_empty: 1 barriers, init_count=1 - mbarrier_init(smem + 8, 1); - // --- pipeline 'k_pipe' --- - // k_full: 1 barriers, init_count=1 - mbarrier_init(smem + 16, 1); - // k_empty: 1 barriers, init_count=1 - mbarrier_init(smem + 24, 1); - // --- pipeline 'v_pipe' --- - // v_full: 1 barriers, init_count=1 - mbarrier_init(smem + 32, 1); - // v_empty: 1 barriers, init_count=1 - mbarrier_init(smem + 40, 1); - // --- pipeline 'kv_pipe' --- - // kv_full: 8 barriers, init_count=1 - mbarrier_init(smem + 48, 1); - mbarrier_init(smem + 56, 1); - mbarrier_init(smem + 64, 1); - mbarrier_init(smem + 72, 1); - mbarrier_init(smem + 80, 1); - mbarrier_init(smem + 88, 1); - mbarrier_init(smem + 96, 1); - mbarrier_init(smem + 104, 1); - // kv_empty: 8 barriers, init_count=1 - mbarrier_init(smem + 112, 1); - mbarrier_init(smem + 120, 1); - mbarrier_init(smem + 128, 1); - mbarrier_init(smem + 136, 1); - mbarrier_init(smem + 144, 1); - mbarrier_init(smem + 152, 1); - mbarrier_init(smem + 160, 1); - mbarrier_init(smem + 168, 1); - // --- pipeline 'index_pipe' --- - // index_full: 1 barriers, init_count=1 - mbarrier_init(smem + 176, 1); - // index_empty: 1 barriers, init_count=1 - mbarrier_init(smem + 184, 1); - // s_full: 2 barriers, init_count=1 - mbarrier_init(smem + 192, 1); - mbarrier_init(smem + 200, 1); - // p_full: 2 barriers, init_count=128 - mbarrier_init(smem + 208, 128); - mbarrier_init(smem + 216, 128); - // corr_done: 2 barriers, init_count=128 - mbarrier_init(smem + 224, 128); - mbarrier_init(smem + 232, 128); - // stats: 2 barriers, init_count=128 - mbarrier_init(smem + 240, 128); - mbarrier_init(smem + 248, 128); - // sum_ready: 1 barriers, init_count=128 - mbarrier_init(smem + 256, 128); - // o_done: 1 barriers, init_count=1 - mbarrier_init(smem + 264, 1); - // pv_done: 1 barriers, init_count=1 - mbarrier_init(smem + 272, 1); - // s_seeded: 1 barriers, init_count=128 - mbarrier_init(smem + 280, 128); - // q_pair_ready: 1 barriers, init_count=64 - mbarrier_init(smem + 288, 64); - // --- pipeline 'kv_pipe' --- - // kv_pair_ready: 8 barriers, init_count=64 - mbarrier_init(smem + 296, 64); - mbarrier_init(smem + 304, 64); - mbarrier_init(smem + 312, 64); - mbarrier_init(smem + 320, 64); - mbarrier_init(smem + 328, 64); - mbarrier_init(smem + 336, 64); - mbarrier_init(smem + 344, 64); - mbarrier_init(smem + 352, 64); - // pv_pair_ready: 2 barriers, init_count=64 - mbarrier_init(smem + 360, 64); - mbarrier_init(smem + 368, 64); - // tmem_dealloc: 1 barriers, init_count=416 - mbarrier_init(smem + 376, 416); - // tmem_dealloc_peer: 1 barriers, init_count=32 - mbarrier_init(smem + 384, 32); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - - __syncwarp(); - - // TMEM alloc (512 columns, 512 used) - volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 392); - if (warp == 0) { - int _tmem_hold = smem + 392; - asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); - __syncwarp(); - asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); - } - - __syncthreads(); - asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); - - const int taddr = tmem_addr_storage[0]; - - // Kernel post-init ops - const int tmem_tmem_scratch = taddr; - - // ---- Ordered hardware-WG register redistribution ---- - // Dec phase frees registers before any WG attempts inc. - if (warp >= 8 && warp <= 15) { - asm volatile("setmaxnreg.dec.sync.aligned.u32 64;"); - } - - // ---- Role: index_warp ---- - if (warp == 13) { - { // index_warp_main - const int index_dummy = 0; - unsigned int _phase_index_empty = 1; - unsigned int _phase_q_full_0 = 0; - { - #pragma unroll 1 - for (unsigned int work_idx = bid; work_idx < total_work_items; work_idx += num_bids) { - mbarrier_wait(q_full_addr, _phase_q_full_0); - _phase_q_full_0 ^= 1; - } - } - } - } - // ---- Role: load_warp ---- - if (warp >= 9 && warp <= 12) { - { // load_warp_main - const int wg2_dummy = 0; - const int load_warp_rank = warp - ((0) ? 12 : 9); - int all_num_kv_tiles = (sparse_topk + 128 - 1) / 128; - unsigned int load_k_stage = 0; - unsigned int load_v_stage = 0; - unsigned int load_kv_stage = 0; - unsigned int load_k_index_stage = 0; - unsigned int load_v_index_stage = 0; - int k_cta_offset = 0; - unsigned int _phase_q_empty_0 = 1; - unsigned int _phase_index_full = 0; - unsigned int _phase_k_empty = 1; - unsigned int _phase_kv_empty = 1; - unsigned int _phase_v_empty = 1; - #pragma unroll 1 - for (unsigned int work_idx_1 = bid; work_idx_1 < total_work_items; work_idx_1 += num_bids) { - int split_idx = 0; - int query_idx = work_idx_1 >> 1; - int v_chunk = work_idx_1 & 1; - int tiles_per_split = all_num_kv_tiles; - int first_tile = split_idx * tiles_per_split; - int num_kv_tiles = tiles_per_split; - int sparse_extent = num_kv_tiles * 128; - mbarrier_wait(q_empty_addr, _phase_q_empty_0); - _phase_q_empty_0 ^= 1; - if (load_warp_rank == 0) { - if (elect_sync()) { - mbarrier_arrive_expect_tx(q_full_addr, 65536); - #pragma unroll - for (int q_stage = 0; q_stage < 4; q_stage++) { - tma_4d_gmem2smem(smem_q_addr + (unsigned int)(q_stage * 16384), (&tmap_q), 0, 0, q_stage, query_idx, q_full_addr); - } - } - } - { - int num_index_passes = sparse_extent / 128; - #pragma unroll 1 - for (int index_pass = load_warp_rank; index_pass < num_index_passes; index_pass += 4) { - int index_offset = index_pass * 128 + lane * 4; - int tile_col_base = (first_tile + index_pass) * 128; - int* row_tile_ptr = ((tile_col_base < 128) ? swa_indices + (query_idx * swa_index_stride) : compressed_indices + (query_idx * compressed_index_stride + (tile_col_base - 128))); - int lane_col = tile_col_base + lane * 4; - int sparse_rows[4]; - #pragma unroll - for (int row_i = 0; row_i < 4; row_i++) { - sparse_rows[row_i] = -1; - } - if (lane_col + 4 <= sparse_topk) { - int _vec_load_0[4]; - { - const int4* _ivptr_0 = reinterpret_cast(row_tile_ptr + (lane * 4) + 0); - int4 _ivld_0; - _ivld_0 = *_ivptr_0; - _vec_load_0[0 + 0] = _ivld_0.x; - _vec_load_0[0 + 1] = _ivld_0.y; - _vec_load_0[0 + 2] = _ivld_0.z; - _vec_load_0[0 + 3] = _ivld_0.w; - } - #pragma unroll - for (int row_i_1 = 0; row_i_1 < 4; row_i_1++) { - sparse_rows[row_i_1] = _vec_load_0[row_i_1]; - } - } else { - #pragma unroll - for (int row_i_2 = 0; row_i_2 < 4; row_i_2++) { - if (lane_col + row_i_2 < sparse_topk) { - sparse_rows[row_i_2] = row_tile_ptr[lane * 4 + row_i_2]; - } - } - } - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(smem_sparse_indices_addr + (unsigned int)(index_offset * 4)), "r"(sparse_rows[0]), "r"(sparse_rows[1]), "r"(sparse_rows[2]), "r"(sparse_rows[3]) : "memory"); - } - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - asm volatile("barrier.sync 8, 128;" ::: "memory"); - } - int k_index_stage_base = smem_sparse_indices_addr; - { - #pragma unroll - for (int qk_stage = 0; qk_stage < 4; qk_stage++) { - if (qk_stage == load_warp_rank) { - mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (load_kv_stage) * 8, 16384); - } - int k_dst = smem_kv_addr + load_kv_stage * 16384; - int group = lane; - if (group < 32) { - int group_offset = k_cta_offset + group * 4; - int raw_rows[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows[0])), "=r"(*reinterpret_cast(&raw_rows[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows[(0) + 3])) - : "r"(k_index_stage_base + group_offset * 4)); - int row0 = ((raw_rows[0] >= 0) ? raw_rows[0] : 0); - int row1 = ((raw_rows[1] >= 0) ? raw_rows[1] : 0); - int row2 = ((raw_rows[2] >= 0) ? raw_rows[2] : 0); - int row3 = ((raw_rows[3] >= 0) ? raw_rows[3] : 0); - if (first_tile == 0) { - tma_gather4_gmem2smem(k_dst + group * 512, (&tmap_swa_kv), qk_stage * 128, row0, row1, row2, row3, kv_full_addr + (load_kv_stage) * 8); - } else { - tma_gather4_gmem2smem(k_dst + group * 512, (&tmap_compressed_kv), qk_stage * 128, row0, row1, row2, row3, kv_full_addr + (load_kv_stage) * 8); - } - } - } - load_kv_stage += 1; - if (load_kv_stage == 8) { load_kv_stage = 0; _phase_kv_empty ^= 1; } - } - } - #pragma unroll 1 - for (int tile = 1; tile < num_kv_tiles; tile++) { - if ((0 & (int)((tile & 1) == 0)) != 0) { - mbarrier_wait(index_full_addr + (load_k_index_stage) * 8, _phase_index_full); - asm volatile("tcgen05.fence::after_thread_sync;"); - } - k_index_stage_base = smem_sparse_indices_addr; - { - #pragma unroll - for (int qk_stage_1 = 0; qk_stage_1 < 4; qk_stage_1++) { - if (qk_stage_1 == load_warp_rank) { - mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (load_kv_stage) * 8, 16384); - } - int k_dst_1 = smem_kv_addr + load_kv_stage * 16384; - int group_1 = lane; - if (group_1 < 32) { - int group_offset_1 = tile * 128 + k_cta_offset + group_1 * 4; - int raw_rows_1[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_1[0])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 3])) - : "r"(k_index_stage_base + group_offset_1 * 4)); - int row0_1 = ((raw_rows_1[0] >= 0) ? raw_rows_1[0] : 0); - int row1_1 = ((raw_rows_1[1] >= 0) ? raw_rows_1[1] : 0); - int row2_1 = ((raw_rows_1[2] >= 0) ? raw_rows_1[2] : 0); - int row3_1 = ((raw_rows_1[3] >= 0) ? raw_rows_1[3] : 0); - tma_gather4_gmem2smem(k_dst_1 + group_1 * 512, (&tmap_compressed_kv), qk_stage_1 * 128, row0_1, row1_1, row2_1, row3_1, kv_full_addr + (load_kv_stage) * 8); - } - } - load_kv_stage += 1; - if (load_kv_stage == 8) { load_kv_stage = 0; _phase_kv_empty ^= 1; } - } - } - int prev_tile = tile - 1; - int global_prev_tile = first_tile + prev_tile; - int v_index_stage_base = smem_sparse_indices_addr; - { - #pragma unroll - for (int pv_stage = 0; pv_stage < 2; pv_stage++) { - if (pv_stage == load_warp_rank) { - mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (load_kv_stage) * 8, 16384); - } - int v_dst = smem_kv_addr + load_kv_stage * 16384; - int v_col = v_chunk * 256 + pv_stage * 128; - int group_2 = lane; - int group_offset_2 = prev_tile * 128 + group_2 * 4; - int raw_rows_2[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_2[0])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 3])) - : "r"(v_index_stage_base + group_offset_2 * 4)); - int row0_2 = ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : 0); - int row1_2 = ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : 0); - int row2_2 = ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : 0); - int row3_2 = ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : 0); - if (global_prev_tile == 0) { - tma_gather4_gmem2smem(v_dst + group_2 * 512, (&tmap_swa_kv), v_col, row0_2, row1_2, row2_2, row3_2, kv_full_addr + (load_kv_stage) * 8); - } else { - tma_gather4_gmem2smem(v_dst + group_2 * 512, (&tmap_compressed_kv), v_col, row0_2, row1_2, row2_2, row3_2, kv_full_addr + (load_kv_stage) * 8); - } - } - load_kv_stage += 1; - if (load_kv_stage == 8) { load_kv_stage = 0; _phase_kv_empty ^= 1; } - } - } - } - int last_tile = num_kv_tiles - 1; - int global_last_tile = first_tile + last_tile; - { - #pragma unroll - for (int pv_stage_1 = 0; pv_stage_1 < 2; pv_stage_1++) { - if (pv_stage_1 == load_warp_rank) { - mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (load_kv_stage) * 8, 16384); - } - int v_dst_1 = smem_kv_addr + load_kv_stage * 16384; - int v_col_1 = v_chunk * 256 + pv_stage_1 * 128; - int group_3 = lane; - int group_offset_3 = last_tile * 128 + group_3 * 4; - int raw_rows_3[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_3[0])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 3])) - : "r"(smem_sparse_indices_addr + (unsigned int)(group_offset_3 * 4))); - int row0_3 = ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : 0); - int row1_3 = ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : 0); - int row2_3 = ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : 0); - int row3_3 = ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : 0); - if (global_last_tile == 0) { - tma_gather4_gmem2smem(v_dst_1 + group_3 * 512, (&tmap_swa_kv), v_col_1, row0_3, row1_3, row2_3, row3_3, kv_full_addr + (load_kv_stage) * 8); - } else { - tma_gather4_gmem2smem(v_dst_1 + group_3 * 512, (&tmap_compressed_kv), v_col_1, row0_3, row1_3, row2_3, row3_3, kv_full_addr + (load_kv_stage) * 8); - } - } - load_kv_stage += 1; - if (load_kv_stage == 8) { load_kv_stage = 0; _phase_kv_empty ^= 1; } - } - } - } - mbarrier_arrive(tmem_dealloc_addr); - } - } - // ---- Role: softmax_wg ---- - if (warp <= 3) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 184;"); - { // softmax_wg_main - float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; - const int wg_dummy_inc = 0; - int all_num_kv_tiles_1 = (sparse_topk + 128 - 1) / 128; - const int tmem_row_base = ((0) ? warp % 2 * 32 : warp % 4 * 32); - const int tmem_score_row_base = ((0) ? (int)(warp % 4 * 32) : tmem_row_base); - const int n_half = ((0) ? (int)(warp % 4 / 2) : 0); - const int my_row = tmem_row_base + lane; - const int stats_row = n_half * 64 + my_row; - int softmax_tile_cursor = 0; - unsigned int softmax_index_stage = 0; - unsigned int _phase_q_full_0_1 = 0; - #pragma unroll 1 - for (unsigned int work_idx_2 = bid; work_idx_2 < total_work_items; work_idx_2 += num_bids) { - int split_idx_1 = 0; - int query_idx_1 = work_idx_2 >> 1; - int tiles_per_split_1 = all_num_kv_tiles_1; - int first_tile_1 = split_idx_1 * tiles_per_split_1; - int num_kv_tiles_1 = tiles_per_split_1; - int _max_0 = ((sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) : (0)); - int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); - int active_topk = _min_0; - { - float seed_zero[4]; - #pragma unroll - for (int seed_c4 = 0; seed_c4 < 4; seed_c4++) { - seed_zero[seed_c4] = 0.0f; - } - #pragma unroll - for (int seed_half = 0; seed_half < 2; seed_half++) { - #pragma unroll - for (int seed_c = 0; seed_c < 8; seed_c++) { - int seed_addr = taddr + (unsigned int)(((seed_half == 0) ? 96 : 224)) + (unsigned int)(seed_c * 4) + (unsigned int)(tmem_row_base << 16); - tmem_st_x4_f32(seed_addr, seed_zero); - } - } - { - #pragma unroll - for (int seed_half_1 = 0; seed_half_1 < 2; seed_half_1++) { - #pragma unroll - for (int seed_c_1 = 16; seed_c_1 < 24; seed_c_1++) { - int seed_addr_1 = taddr + (unsigned int)(((seed_half_1 == 0) ? 0 : 128)) + (unsigned int)(seed_c_1 * 4) + (unsigned int)(tmem_row_base << 16); - tmem_st_x4_f32(seed_addr_1, seed_zero); - } - } - } - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(s_seeded_addr); - } - { - mbarrier_wait(q_full_addr, _phase_q_full_0_1); - _phase_q_full_0_1 ^= 1; - } - float row_max_val = -CAKE_INF; - float row_sum_val = 0.0f; - int sink_head = ((0) ? my_row : my_row); - if (has_sinks != 0 && sink_head < num_heads && split_idx_1 == 0) { - row_max_val = sinks[sink_head] * 1.4426950408889634f / softmax_scale_log2; - row_sum_val = 1.0f; - } - #pragma unroll 1 - for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { - int pipeline_tile = softmax_tile_cursor + tile_1; - int phase = pipeline_tile & 1; - int s_wait_phase = pipeline_tile >> 1 & 1; - mbarrier_wait(s_full_addr + (phase) * 8, s_wait_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - int s_off = ((phase != 0) ? 128 : 0); - int s_base = taddr + (unsigned int)s_off + (unsigned int)(tmem_score_row_base << 16); - float new_max = row_max_val; - int valid_sparse_cols = ((active_topk < sparse_topk) ? active_topk : sparse_topk); - valid_sparse_cols = valid_sparse_cols - (first_tile_1 + tile_1) * 128 - n_half * 64; - if (valid_sparse_cols < 0) { - valid_sparse_cols = 0; - } - if (valid_sparse_cols > 128) { - valid_sparse_cols = 128; - } - int index_tile_addr = smem_sparse_indices_addr + (unsigned int)(tile_1 * 512); - int index_lane_addr = index_tile_addr + (n_half * 64 + lane) * 4; - int staged_index[4]; - unsigned int invalid_cols[4]; - unsigned int any_invalid = 0; - #pragma unroll - for (int w = 0; w < 4; w++) { - asm volatile("ld.shared.b32 %0, [%1];" : "=r"(*reinterpret_cast(&staged_index[w])) : "r"(index_lane_addr + w * 128)); - unsigned int _vote_0 = __ballot_sync(0xFFFFFFFF, staged_index[w] < 0); - invalid_cols[w] = _vote_0; - any_invalid = any_invalid | invalid_cols[w]; - } - float _tmem_load_0[4]; - tmem_ld_x4(&_tmem_load_0[0], s_base); - float _tmem_load_1[4]; - tmem_ld_x4(&_tmem_load_1[0], s_base); - { - #pragma unroll - for (int frag = 0; frag < 4; frag++) { - float _tmem_load_2[32]; - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]) - : "r"(s_base + frag * 32)); - int frag_valid = valid_sparse_cols - frag * 32; - uint32_t _slice_lo_mask_1; - { - int _lim_0 = frag_valid; - if (_lim_0 <= 0) { _slice_lo_mask_1 = 0u; } - else if (_lim_0 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_0)); - } - } - if (!(_slice_lo_mask_1 & (1u << 0))) _tmem_load_2[0] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 1))) _tmem_load_2[1] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 2))) _tmem_load_2[2] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 3))) _tmem_load_2[3] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 4))) _tmem_load_2[4] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 5))) _tmem_load_2[5] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 6))) _tmem_load_2[6] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 7))) _tmem_load_2[7] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 8))) _tmem_load_2[8] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 9))) _tmem_load_2[9] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 10))) _tmem_load_2[10] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 11))) _tmem_load_2[11] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 12))) _tmem_load_2[12] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 13))) _tmem_load_2[13] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 14))) _tmem_load_2[14] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 15))) _tmem_load_2[15] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 16))) _tmem_load_2[16] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 17))) _tmem_load_2[17] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 18))) _tmem_load_2[18] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 19))) _tmem_load_2[19] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 20))) _tmem_load_2[20] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 21))) _tmem_load_2[21] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 22))) _tmem_load_2[22] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 23))) _tmem_load_2[23] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 24))) _tmem_load_2[24] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 25))) _tmem_load_2[25] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 26))) _tmem_load_2[26] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 27))) _tmem_load_2[27] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 28))) _tmem_load_2[28] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 29))) _tmem_load_2[29] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 30))) _tmem_load_2[30] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 31))) _tmem_load_2[31] = -CAKE_INF; - if (any_invalid != 0) { - unsigned int frag_mask_word = invalid_cols[frag]; - _tmem_load_2[0] = (((frag_mask_word & 1) != 0) ? -CAKE_INF : _tmem_load_2[0]); - _tmem_load_2[1] = (((frag_mask_word >> 1 & 1) != 0) ? -CAKE_INF : _tmem_load_2[1]); - _tmem_load_2[2] = (((frag_mask_word >> 2 & 1) != 0) ? -CAKE_INF : _tmem_load_2[2]); - _tmem_load_2[3] = (((frag_mask_word >> 3 & 1) != 0) ? -CAKE_INF : _tmem_load_2[3]); - _tmem_load_2[4] = (((frag_mask_word >> 4 & 1) != 0) ? -CAKE_INF : _tmem_load_2[4]); - _tmem_load_2[5] = (((frag_mask_word >> 5 & 1) != 0) ? -CAKE_INF : _tmem_load_2[5]); - _tmem_load_2[6] = (((frag_mask_word >> 6 & 1) != 0) ? -CAKE_INF : _tmem_load_2[6]); - _tmem_load_2[7] = (((frag_mask_word >> 7 & 1) != 0) ? -CAKE_INF : _tmem_load_2[7]); - _tmem_load_2[8] = (((frag_mask_word >> 8 & 1) != 0) ? -CAKE_INF : _tmem_load_2[8]); - _tmem_load_2[9] = (((frag_mask_word >> 9 & 1) != 0) ? -CAKE_INF : _tmem_load_2[9]); - _tmem_load_2[10] = (((frag_mask_word >> 10 & 1) != 0) ? -CAKE_INF : _tmem_load_2[10]); - _tmem_load_2[11] = (((frag_mask_word >> 11 & 1) != 0) ? -CAKE_INF : _tmem_load_2[11]); - _tmem_load_2[12] = (((frag_mask_word >> 12 & 1) != 0) ? -CAKE_INF : _tmem_load_2[12]); - _tmem_load_2[13] = (((frag_mask_word >> 13 & 1) != 0) ? -CAKE_INF : _tmem_load_2[13]); - _tmem_load_2[14] = (((frag_mask_word >> 14 & 1) != 0) ? -CAKE_INF : _tmem_load_2[14]); - _tmem_load_2[15] = (((frag_mask_word >> 15 & 1) != 0) ? -CAKE_INF : _tmem_load_2[15]); - _tmem_load_2[16] = (((frag_mask_word >> 16 & 1) != 0) ? -CAKE_INF : _tmem_load_2[16]); - _tmem_load_2[17] = (((frag_mask_word >> 17 & 1) != 0) ? -CAKE_INF : _tmem_load_2[17]); - _tmem_load_2[18] = (((frag_mask_word >> 18 & 1) != 0) ? -CAKE_INF : _tmem_load_2[18]); - _tmem_load_2[19] = (((frag_mask_word >> 19 & 1) != 0) ? -CAKE_INF : _tmem_load_2[19]); - _tmem_load_2[20] = (((frag_mask_word >> 20 & 1) != 0) ? -CAKE_INF : _tmem_load_2[20]); - _tmem_load_2[21] = (((frag_mask_word >> 21 & 1) != 0) ? -CAKE_INF : _tmem_load_2[21]); - _tmem_load_2[22] = (((frag_mask_word >> 22 & 1) != 0) ? -CAKE_INF : _tmem_load_2[22]); - _tmem_load_2[23] = (((frag_mask_word >> 23 & 1) != 0) ? -CAKE_INF : _tmem_load_2[23]); - _tmem_load_2[24] = (((frag_mask_word >> 24 & 1) != 0) ? -CAKE_INF : _tmem_load_2[24]); - _tmem_load_2[25] = (((frag_mask_word >> 25 & 1) != 0) ? -CAKE_INF : _tmem_load_2[25]); - _tmem_load_2[26] = (((frag_mask_word >> 26 & 1) != 0) ? -CAKE_INF : _tmem_load_2[26]); - _tmem_load_2[27] = (((frag_mask_word >> 27 & 1) != 0) ? -CAKE_INF : _tmem_load_2[27]); - _tmem_load_2[28] = (((frag_mask_word >> 28 & 1) != 0) ? -CAKE_INF : _tmem_load_2[28]); - _tmem_load_2[29] = (((frag_mask_word >> 29 & 1) != 0) ? -CAKE_INF : _tmem_load_2[29]); - _tmem_load_2[30] = (((frag_mask_word >> 30 & 1) != 0) ? -CAKE_INF : _tmem_load_2[30]); - _tmem_load_2[31] = (((frag_mask_word >> 31 & 1) != 0) ? -CAKE_INF : _tmem_load_2[31]); - } - float2 _reg_reduce_max2_1 = {-CAKE_INF, -CAKE_INF}; - _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_2[0], _tmem_load_2[1])); - _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_2[2], _tmem_load_2[3])); - _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_2[4], _tmem_load_2[5])); - _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_2[6], _tmem_load_2[7])); - _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_2[8], _tmem_load_2[9])); - _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_2[10], _tmem_load_2[11])); - _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_2[12], _tmem_load_2[13])); - _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_2[14], _tmem_load_2[15])); - _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_2[16], _tmem_load_2[17])); - _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_2[18], _tmem_load_2[19])); - _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_2[20], _tmem_load_2[21])); - _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_2[22], _tmem_load_2[23])); - _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_2[24], _tmem_load_2[25])); - _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_2[26], _tmem_load_2[27])); - _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_2[28], _tmem_load_2[29])); - _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_2[30], _tmem_load_2[31])); - float _tmem_load_2_max = row_max_reduce(_reg_reduce_max2_1); - float _max_3 = max_noftz(new_max, _tmem_load_2_max); - new_max = _max_3; - } - } - float _fma_0 = __fmaf_rn(row_max_val, softmax_scale_log2, (-new_max) * softmax_scale_log2); - float delta = _fma_0; - float _exp2_0 = approx_exp2(delta); - float exp_delta = _exp2_0; - float acc_scale = ((row_max_val > -CAKE_INF) ? exp_delta : 1.0f); - smem_stats_max[phase * 128 + stats_row] = acc_scale; - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(stats_addr + (phase) * 8); - row_max_val = new_max; - float safe_max = ((new_max == -CAKE_INF) ? 0.0f : new_max); - float max_scaled = safe_max * softmax_scale_log2; - float block_sum = 0.0f; - { - int p_off = ((phase != 0) ? 224 : 96); - int p_base = taddr + (unsigned int)p_off + (unsigned int)(tmem_row_base << 16); - #pragma unroll - for (int frag_1 = 0; frag_1 < 4; frag_1++) { - float _tmem_load_3[32]; - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]) - : "r"(s_base + frag_1 * 32)); - int frag_valid_1 = valid_sparse_cols - frag_1 * 32; - uint32_t _slice_lo_mask_3; - { - int _lim_2 = frag_valid_1; - if (_lim_2 <= 0) { _slice_lo_mask_3 = 0u; } - else if (_lim_2 >= 32) { _slice_lo_mask_3 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_slice_lo_mask_3) : "r"(_lim_2)); - } - } - if (!(_slice_lo_mask_3 & (1u << 0))) _tmem_load_3[0] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 1))) _tmem_load_3[1] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 2))) _tmem_load_3[2] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 3))) _tmem_load_3[3] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 4))) _tmem_load_3[4] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 5))) _tmem_load_3[5] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 6))) _tmem_load_3[6] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 7))) _tmem_load_3[7] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 8))) _tmem_load_3[8] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 9))) _tmem_load_3[9] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 10))) _tmem_load_3[10] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 11))) _tmem_load_3[11] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 12))) _tmem_load_3[12] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 13))) _tmem_load_3[13] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 14))) _tmem_load_3[14] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 15))) _tmem_load_3[15] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 16))) _tmem_load_3[16] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 17))) _tmem_load_3[17] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 18))) _tmem_load_3[18] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 19))) _tmem_load_3[19] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 20))) _tmem_load_3[20] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 21))) _tmem_load_3[21] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 22))) _tmem_load_3[22] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 23))) _tmem_load_3[23] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 24))) _tmem_load_3[24] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 25))) _tmem_load_3[25] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 26))) _tmem_load_3[26] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 27))) _tmem_load_3[27] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 28))) _tmem_load_3[28] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 29))) _tmem_load_3[29] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 30))) _tmem_load_3[30] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 31))) _tmem_load_3[31] = -CAKE_INF; - if (any_invalid != 0) { - unsigned int frag_mask_word_1 = invalid_cols[frag_1]; - _tmem_load_3[0] = (((frag_mask_word_1 & 1) != 0) ? -CAKE_INF : _tmem_load_3[0]); - _tmem_load_3[1] = (((frag_mask_word_1 >> 1 & 1) != 0) ? -CAKE_INF : _tmem_load_3[1]); - _tmem_load_3[2] = (((frag_mask_word_1 >> 2 & 1) != 0) ? -CAKE_INF : _tmem_load_3[2]); - _tmem_load_3[3] = (((frag_mask_word_1 >> 3 & 1) != 0) ? -CAKE_INF : _tmem_load_3[3]); - _tmem_load_3[4] = (((frag_mask_word_1 >> 4 & 1) != 0) ? -CAKE_INF : _tmem_load_3[4]); - _tmem_load_3[5] = (((frag_mask_word_1 >> 5 & 1) != 0) ? -CAKE_INF : _tmem_load_3[5]); - _tmem_load_3[6] = (((frag_mask_word_1 >> 6 & 1) != 0) ? -CAKE_INF : _tmem_load_3[6]); - _tmem_load_3[7] = (((frag_mask_word_1 >> 7 & 1) != 0) ? -CAKE_INF : _tmem_load_3[7]); - _tmem_load_3[8] = (((frag_mask_word_1 >> 8 & 1) != 0) ? -CAKE_INF : _tmem_load_3[8]); - _tmem_load_3[9] = (((frag_mask_word_1 >> 9 & 1) != 0) ? -CAKE_INF : _tmem_load_3[9]); - _tmem_load_3[10] = (((frag_mask_word_1 >> 10 & 1) != 0) ? -CAKE_INF : _tmem_load_3[10]); - _tmem_load_3[11] = (((frag_mask_word_1 >> 11 & 1) != 0) ? -CAKE_INF : _tmem_load_3[11]); - _tmem_load_3[12] = (((frag_mask_word_1 >> 12 & 1) != 0) ? -CAKE_INF : _tmem_load_3[12]); - _tmem_load_3[13] = (((frag_mask_word_1 >> 13 & 1) != 0) ? -CAKE_INF : _tmem_load_3[13]); - _tmem_load_3[14] = (((frag_mask_word_1 >> 14 & 1) != 0) ? -CAKE_INF : _tmem_load_3[14]); - _tmem_load_3[15] = (((frag_mask_word_1 >> 15 & 1) != 0) ? -CAKE_INF : _tmem_load_3[15]); - _tmem_load_3[16] = (((frag_mask_word_1 >> 16 & 1) != 0) ? -CAKE_INF : _tmem_load_3[16]); - _tmem_load_3[17] = (((frag_mask_word_1 >> 17 & 1) != 0) ? -CAKE_INF : _tmem_load_3[17]); - _tmem_load_3[18] = (((frag_mask_word_1 >> 18 & 1) != 0) ? -CAKE_INF : _tmem_load_3[18]); - _tmem_load_3[19] = (((frag_mask_word_1 >> 19 & 1) != 0) ? -CAKE_INF : _tmem_load_3[19]); - _tmem_load_3[20] = (((frag_mask_word_1 >> 20 & 1) != 0) ? -CAKE_INF : _tmem_load_3[20]); - _tmem_load_3[21] = (((frag_mask_word_1 >> 21 & 1) != 0) ? -CAKE_INF : _tmem_load_3[21]); - _tmem_load_3[22] = (((frag_mask_word_1 >> 22 & 1) != 0) ? -CAKE_INF : _tmem_load_3[22]); - _tmem_load_3[23] = (((frag_mask_word_1 >> 23 & 1) != 0) ? -CAKE_INF : _tmem_load_3[23]); - _tmem_load_3[24] = (((frag_mask_word_1 >> 24 & 1) != 0) ? -CAKE_INF : _tmem_load_3[24]); - _tmem_load_3[25] = (((frag_mask_word_1 >> 25 & 1) != 0) ? -CAKE_INF : _tmem_load_3[25]); - _tmem_load_3[26] = (((frag_mask_word_1 >> 26 & 1) != 0) ? -CAKE_INF : _tmem_load_3[26]); - _tmem_load_3[27] = (((frag_mask_word_1 >> 27 & 1) != 0) ? -CAKE_INF : _tmem_load_3[27]); - _tmem_load_3[28] = (((frag_mask_word_1 >> 28 & 1) != 0) ? -CAKE_INF : _tmem_load_3[28]); - _tmem_load_3[29] = (((frag_mask_word_1 >> 29 & 1) != 0) ? -CAKE_INF : _tmem_load_3[29]); - _tmem_load_3[30] = (((frag_mask_word_1 >> 30 & 1) != 0) ? -CAKE_INF : _tmem_load_3[30]); - _tmem_load_3[31] = (((frag_mask_word_1 >> 31 & 1) != 0) ? -CAKE_INF : _tmem_load_3[31]); - } - const float2 _fma_b2_3 = {softmax_scale_log2, softmax_scale_log2}; - const float2 _fma_c2_4 = {-max_scaled, -max_scaled}; - float2 _fma_pair_5 = fma_f32x2(make_float2(_tmem_load_3[0], _tmem_load_3[1]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[0] = _fma_pair_5.x; - _tmem_load_3[1] = _fma_pair_5.y; - float2 _fma_pair_6 = fma_f32x2(make_float2(_tmem_load_3[2], _tmem_load_3[3]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[2] = _fma_pair_6.x; - _tmem_load_3[3] = _fma_pair_6.y; - float2 _fma_pair_7 = fma_f32x2(make_float2(_tmem_load_3[4], _tmem_load_3[5]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[4] = _fma_pair_7.x; - _tmem_load_3[5] = _fma_pair_7.y; - float2 _fma_pair_8 = fma_f32x2(make_float2(_tmem_load_3[6], _tmem_load_3[7]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[6] = _fma_pair_8.x; - _tmem_load_3[7] = _fma_pair_8.y; - float2 _fma_pair_9 = fma_f32x2(make_float2(_tmem_load_3[8], _tmem_load_3[9]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[8] = _fma_pair_9.x; - _tmem_load_3[9] = _fma_pair_9.y; - float2 _fma_pair_10 = fma_f32x2(make_float2(_tmem_load_3[10], _tmem_load_3[11]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[10] = _fma_pair_10.x; - _tmem_load_3[11] = _fma_pair_10.y; - float2 _fma_pair_11 = fma_f32x2(make_float2(_tmem_load_3[12], _tmem_load_3[13]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[12] = _fma_pair_11.x; - _tmem_load_3[13] = _fma_pair_11.y; - float2 _fma_pair_12 = fma_f32x2(make_float2(_tmem_load_3[14], _tmem_load_3[15]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[14] = _fma_pair_12.x; - _tmem_load_3[15] = _fma_pair_12.y; - float2 _fma_pair_13 = fma_f32x2(make_float2(_tmem_load_3[16], _tmem_load_3[17]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[16] = _fma_pair_13.x; - _tmem_load_3[17] = _fma_pair_13.y; - float2 _fma_pair_14 = fma_f32x2(make_float2(_tmem_load_3[18], _tmem_load_3[19]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[18] = _fma_pair_14.x; - _tmem_load_3[19] = _fma_pair_14.y; - float2 _fma_pair_15 = fma_f32x2(make_float2(_tmem_load_3[20], _tmem_load_3[21]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[20] = _fma_pair_15.x; - _tmem_load_3[21] = _fma_pair_15.y; - float2 _fma_pair_16 = fma_f32x2(make_float2(_tmem_load_3[22], _tmem_load_3[23]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[22] = _fma_pair_16.x; - _tmem_load_3[23] = _fma_pair_16.y; - float2 _fma_pair_17 = fma_f32x2(make_float2(_tmem_load_3[24], _tmem_load_3[25]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[24] = _fma_pair_17.x; - _tmem_load_3[25] = _fma_pair_17.y; - float2 _fma_pair_18 = fma_f32x2(make_float2(_tmem_load_3[26], _tmem_load_3[27]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[26] = _fma_pair_18.x; - _tmem_load_3[27] = _fma_pair_18.y; - float2 _fma_pair_19 = fma_f32x2(make_float2(_tmem_load_3[28], _tmem_load_3[29]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[28] = _fma_pair_19.x; - _tmem_load_3[29] = _fma_pair_19.y; - float2 _fma_pair_20 = fma_f32x2(make_float2(_tmem_load_3[30], _tmem_load_3[31]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[30] = _fma_pair_20.x; - _tmem_load_3[31] = _fma_pair_20.y; - #pragma unroll - for (int _le = 0; _le < 32; _le++) { - _tmem_load_3[_le] = approx_exp2(_tmem_load_3[_le]); - } - float2 _reg_reduce_sum2_21 = make_float2(0.0f, 0.0f); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[0], _tmem_load_3[1])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[2], _tmem_load_3[3])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[4], _tmem_load_3[5])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[6], _tmem_load_3[7])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[8], _tmem_load_3[9])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[10], _tmem_load_3[11])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[12], _tmem_load_3[13])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[14], _tmem_load_3[15])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[16], _tmem_load_3[17])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[18], _tmem_load_3[19])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[20], _tmem_load_3[21])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[22], _tmem_load_3[23])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[24], _tmem_load_3[25])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[26], _tmem_load_3[27])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[28], _tmem_load_3[29])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[30], _tmem_load_3[31])); - float _tmem_load_3_sum = _reg_reduce_sum2_21.x + _reg_reduce_sum2_21.y; - block_sum = block_sum + _tmem_load_3_sum; - { - uint32_t _pv_packed[8]; - #pragma unroll - for (int _j = 0; _j < 8; _j++) { - uint32_t _pk; - asm("{\n\t" - ".reg .b16 _lo, _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}\n" - : "=r"(_pk) : "f"(_tmem_load_3[0 + _j * 4]), "f"(_tmem_load_3[0 + _j * 4 + 1]), - "f"(_tmem_load_3[0 + _j * 4 + 2]), "f"(_tmem_load_3[0 + _j * 4 + 3])); - _pv_packed[_j] = _pk; - } - tmem_st_x8_u32(p_base + frag_1 * 8, _pv_packed); - } - } - } - { - float _fma_1 = __fmaf_rn(row_sum_val, acc_scale, block_sum); - row_sum_val = _fma_1; - } - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(p_full_addr + (phase) * 8); - } - smem_stats_sum[stats_row] = row_sum_val; - smem_stats_final_max[stats_row] = row_max_val; - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(sum_ready_addr); - softmax_tile_cursor = softmax_tile_cursor + num_kv_tiles_1; - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(tmem_dealloc_addr); - } - } - // ---- Role: correction_wg ---- - if (warp >= 4 && warp <= 7) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 184;"); - { // correction_wg_main - float softmax_scale_log2_1 = bmm1_scale[0] * 1.4426950408889634f; - float output_scale = bmm2_scale[0]; - const int wg_dummy_inc_1 = 0; - int all_num_kv_tiles_2 = (sparse_topk + 128 - 1) / 128; - const int tmem_row_base_1 = ((0) ? warp % 2 * 32 : warp % 4 * 32); - const int n_half_1 = ((0) ? (int)(warp % 4 / 2) : 0); - const int my_row_1 = tmem_row_base_1 + lane; - const int stats_row_1 = n_half_1 * 64 + my_row_1; - const int corr_row = tmem_row_base_1 << 16; - int correction_tile_cursor = 0; - unsigned int _phase_q_full_0_2 = 0; - unsigned int _phase_pv_done_0 = 0; - unsigned int _phase_o_done_0 = 0; - unsigned int _phase_sum_ready_0 = 0; - #pragma unroll 1 - for (unsigned int work_idx_3 = bid; work_idx_3 < total_work_items; work_idx_3 += num_bids) { - int split_idx_2 = 0; - int query_idx_2 = work_idx_3 >> 1; - int v_chunk_1 = work_idx_3 & 1; - int tiles_per_split_2 = all_num_kv_tiles_2; - int first_tile_2 = split_idx_2 * tiles_per_split_2; - int num_kv_tiles_2 = tiles_per_split_2; - { - mbarrier_wait(q_full_addr, _phase_q_full_0_2); - _phase_q_full_0_2 ^= 1; - } - #pragma unroll 1 - for (int tile_2 = 0; tile_2 < num_kv_tiles_2; tile_2++) { - int pipeline_tile_1 = correction_tile_cursor + tile_2; - int phase_1 = pipeline_tile_1 & 1; - int stats_wait_phase = pipeline_tile_1 >> 1 & 1; - mbarrier_wait(stats_addr + (phase_1) * 8, stats_wait_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - float acc_scale_1 = smem_stats_max[phase_1 * 128 + stats_row_1]; - if (tile_2 > 0) { - mbarrier_wait(pv_done_addr, _phase_pv_done_0); - _phase_pv_done_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - } - if (tile_2 > 0) { - int _vote_1 = __any_sync(0xFFFFFFFF, acc_scale_1 < 1.0f); - int any_rescale = _vote_1; - if (any_rescale != 0) { - #pragma unroll - for (int vs = 0; vs < 2; vs++) { - int o_base = taddr + 256 + (unsigned int)(vs * 128) + (unsigned int)corr_row; - #pragma unroll - for (int c = 0; c < 128; c += 16) { - float _tmem_load_4[16]; - tmem_ld_x16(&_tmem_load_4[0], o_base + c); - const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; - #pragma unroll - for (int _ls = 0; _ls < 8; _ls++) - mul_f32x2_inplace(&reinterpret_cast(_tmem_load_4)[_ls], _scale2_0); - tmem_st_x16_f32(o_base + c, _tmem_load_4); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - } - } - } - } - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(corr_done_addr + (phase_1) * 8); - } - mbarrier_wait(o_done_addr, _phase_o_done_0); - _phase_o_done_0 ^= 1; - mbarrier_wait(sum_ready_addr, _phase_sum_ready_0); - _phase_sum_ready_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - float total_sum = smem_stats_sum[stats_row_1]; - float final_max = smem_stats_final_max[stats_row_1]; - float _rcp_0 = approx_rcp(total_sum); - float inv_sum = ((total_sum > 0.0f) ? _rcp_0 : 0.0f); - int head_idx = my_row_1; - int direct_o_offset = (query_idx_2 * num_heads + head_idx) * 512 + v_chunk_1 * 256; - int partial_o_offset = (query_idx_2 * num_heads + head_idx + split_idx_2) * 512 + ((0) ? v_chunk_1 : n_half_1) * 256; - int o_offset = direct_o_offset; - if ((0 & (int)(n_half_1 == 0) & (int)(head_idx < num_heads)) != 0) { - int stat_offset = query_idx_2 * num_heads + head_idx + split_idx_2; - float _log2_0; - asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(total_sum)); - partial_lse[stat_offset] = ((total_sum > 0.0f) ? final_max * softmax_scale_log2_1 + _log2_0 : -CAKE_INF); - } - #pragma unroll - for (int vs_1 = 0; vs_1 < 2; vs_1++) { - int o_base_epi = taddr + 256 + (unsigned int)(vs_1 * 128) + (unsigned int)corr_row; - #pragma unroll - for (int c_1 = 0; c_1 < 128; c_1 += 32) { - float _tmem_load_5[32]; - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(_tmem_load_5[0]), "=f"(_tmem_load_5[1]), "=f"(_tmem_load_5[2]), "=f"(_tmem_load_5[3]), "=f"(_tmem_load_5[4]), "=f"(_tmem_load_5[5]), "=f"(_tmem_load_5[6]), "=f"(_tmem_load_5[7]), "=f"(_tmem_load_5[8]), "=f"(_tmem_load_5[9]), "=f"(_tmem_load_5[10]), "=f"(_tmem_load_5[11]), "=f"(_tmem_load_5[12]), "=f"(_tmem_load_5[13]), "=f"(_tmem_load_5[14]), "=f"(_tmem_load_5[15]), "=f"(_tmem_load_5[16]), "=f"(_tmem_load_5[17]), "=f"(_tmem_load_5[18]), "=f"(_tmem_load_5[19]), "=f"(_tmem_load_5[20]), "=f"(_tmem_load_5[21]), "=f"(_tmem_load_5[22]), "=f"(_tmem_load_5[23]), "=f"(_tmem_load_5[24]), "=f"(_tmem_load_5[25]), "=f"(_tmem_load_5[26]), "=f"(_tmem_load_5[27]), "=f"(_tmem_load_5[28]), "=f"(_tmem_load_5[29]), "=f"(_tmem_load_5[30]), "=f"(_tmem_load_5[31]) - : "r"(o_base_epi + c_1)); - int gmem_base = o_offset + vs_1 * 128 + c_1; - if (head_idx < num_heads) { - #pragma unroll - for (int j = 0; j < 32; j += 16) { - { - const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; - #if __CUDA_ARCH__ >= 1000 - #pragma unroll - for (int _ps = 0; _ps < 8; _ps++) - mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_5[j])[_ps], _prescale2_1); - #else - #pragma unroll - for (int _ps = 0; _ps < 16; _ps++) - _tmem_load_5[j + _ps] *= inv_sum * output_scale; - #endif - { - __nv_bfloat162 _pk0 = __floats2bfloat162_rn(_tmem_load_5[j + 0], _tmem_load_5[j + 1]); - unsigned _pk_u0 = *reinterpret_cast(&_pk0); - __nv_bfloat162 _pk1 = __floats2bfloat162_rn(_tmem_load_5[j + 2], _tmem_load_5[j + 3]); - unsigned _pk_u1 = *reinterpret_cast(&_pk1); - __nv_bfloat162 _pk2 = __floats2bfloat162_rn(_tmem_load_5[j + 4], _tmem_load_5[j + 5]); - unsigned _pk_u2 = *reinterpret_cast(&_pk2); - __nv_bfloat162 _pk3 = __floats2bfloat162_rn(_tmem_load_5[j + 6], _tmem_load_5[j + 7]); - unsigned _pk_u3 = *reinterpret_cast(&_pk3); - __nv_bfloat162 _pk4 = __floats2bfloat162_rn(_tmem_load_5[j + 8], _tmem_load_5[j + 9]); - unsigned _pk_u4 = *reinterpret_cast(&_pk4); - __nv_bfloat162 _pk5 = __floats2bfloat162_rn(_tmem_load_5[j + 10], _tmem_load_5[j + 11]); - unsigned _pk_u5 = *reinterpret_cast(&_pk5); - __nv_bfloat162 _pk6 = __floats2bfloat162_rn(_tmem_load_5[j + 12], _tmem_load_5[j + 13]); - unsigned _pk_u6 = *reinterpret_cast(&_pk6); - __nv_bfloat162 _pk7 = __floats2bfloat162_rn(_tmem_load_5[j + 14], _tmem_load_5[j + 15]); - unsigned _pk_u7 = *reinterpret_cast(&_pk7); - asm volatile( - "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" - :: "l"((void*)(&((__nv_bfloat16*)(O + (gmem_base + j)))[0])), "r"(_pk_u0), "r"(_pk_u1), "r"(_pk_u2), "r"(_pk_u3), "r"(_pk_u4), "r"(_pk_u5), "r"(_pk_u6), "r"(_pk_u7) : "memory"); - } - } - } - } - } - } - correction_tile_cursor = correction_tile_cursor + num_kv_tiles_2; - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(tmem_dealloc_addr); - } - } - // ---- Role: mma_warp ---- - if (warp == 8) { - { // mma_warp_main - const int wg2_dummy_1 = 0; - int all_num_kv_tiles_3 = (sparse_topk + 128 - 1) / 128; - unsigned int mma_k_stage = 0; - unsigned int mma_v_stage = 0; - unsigned int mma_kv_stage = 0; - int mma_tile_cursor = 0; - unsigned int _phase_s_seeded_0 = 0; - unsigned int _phase_q_full_0_3 = 0; - unsigned int _phase_q_pair_ready_0 = 0; - unsigned int _phase_k_full = 0; - unsigned int _phase_kv_full = 0; - unsigned int _phase_kv_pair_ready = 0; - unsigned int _phase_v_full = 0; - #pragma unroll 1 - for (unsigned int work_idx_4 = bid; work_idx_4 < total_work_items; work_idx_4 += num_bids) { - int split_idx_3 = 0; - int tiles_per_split_3 = all_num_kv_tiles_3; - int first_tile_3 = split_idx_3 * tiles_per_split_3; - int num_kv_tiles_3 = tiles_per_split_3; - { - { - mbarrier_wait(s_seeded_addr, _phase_s_seeded_0); - _phase_s_seeded_0 ^= 1; - } - asm volatile("tcgen05.fence::after_thread_sync;"); - } - { - mbarrier_wait(q_full_addr, _phase_q_full_0_3); - _phase_q_full_0_3 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - } - int first_pv = 1; - #pragma unroll 1 - for (int tile_3 = 0; tile_3 < num_kv_tiles_3; tile_3++) { - int pipeline_tile_2 = mma_tile_cursor + tile_3; - int phase_2 = pipeline_tile_2 & 1; - int score_col = ((phase_2 != 0) ? 128 : 0); - { - #pragma unroll 1 - for (int n = 0; n < 4; n++) { - mbarrier_wait(kv_full_addr + (mma_kv_stage) * 8, _phase_kv_full); - asm volatile("tcgen05.fence::after_thread_sync;"); - { - int _mma_a_lo_4 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (n) * 1024); - int _mma_b_lo_4 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 1024); - { - uint64_t _mma_ss_a_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_4); - uint64_t _mma_ss_b_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_4); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem_scratch + (score_col)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 136314896, ((n == 0) ? 0 : 1)); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem_scratch + (score_col)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 136314896, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem_scratch + (score_col)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 136314896, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem_scratch + (score_col)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 136314896, 1); - } - } - if (n == 3) { - elect_commit(s_full_addr + (phase_2) * 8); - } - elect_commit(kv_empty_addr + (mma_kv_stage) * 8); - } - mma_kv_stage += 1; - if (mma_kv_stage == 8) { mma_kv_stage = 0; _phase_kv_full ^= 1; _phase_kv_pair_ready ^= 1; } - } - } - if (tile_3 > 0) { - int prev_pipeline_tile = pipeline_tile_2 - 1; - int prev_phase = prev_pipeline_tile & 1; - int pv_wait_phase = prev_pipeline_tile >> 1 & 1; - mbarrier_wait(p_full_addr + (prev_phase) * 8, pv_wait_phase); - mbarrier_wait(corr_done_addr + (prev_phase) * 8, pv_wait_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - int p_col = ((prev_phase != 0) ? 224 : 96); - { - #pragma unroll - for (int vs_2 = 0; vs_2 < 2; vs_2++) { - mbarrier_wait(kv_full_addr + (mma_kv_stage) * 8, _phase_kv_full); - asm volatile("tcgen05.fence::after_thread_sync;"); - int output_col = 256 + vs_2 * 128; - { - int _mma_b_lo_8 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (mma_kv_stage) * 1024); - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 dhi, blo, ta, id;\n\t" - ".reg .b64 db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "" - "mov.b32 dhi, 0x40004040;\n\t" - "mov.b32 id, 136380432;\n\t" - "mov.b32 ta, %2;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p0;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" - "}\n" - :: "r"((tmem_tmem_scratch + (output_col))), "r"(_mma_b_lo_8), "r"(tmem_tmem_scratch + p_col), "r"(((first_pv) ? 0 : 1))); - elect_commit(kv_empty_addr + (mma_kv_stage) * 8); - } - mma_kv_stage += 1; - if (mma_kv_stage == 8) { mma_kv_stage = 0; _phase_kv_full ^= 1; _phase_kv_pair_ready ^= 1; } - } - } - first_pv = 0; - elect_commit(pv_done_addr); - } - } - int last_pipeline_tile = mma_tile_cursor + num_kv_tiles_3 - 1; - int last_phase = last_pipeline_tile & 1; - int drain_wait_phase = last_pipeline_tile >> 1 & 1; - mbarrier_wait(p_full_addr + (last_phase) * 8, drain_wait_phase); - mbarrier_wait(corr_done_addr + (last_phase) * 8, drain_wait_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - int p_col_last = ((last_phase != 0) ? 224 : 96); - { - #pragma unroll - for (int vs_3 = 0; vs_3 < 2; vs_3++) { - mbarrier_wait(kv_full_addr + (mma_kv_stage) * 8, _phase_kv_full); - asm volatile("tcgen05.fence::after_thread_sync;"); - int output_col_d = 256 + vs_3 * 128; - { - int _mma_b_lo_12 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (mma_kv_stage) * 1024); - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 dhi, blo, ta, id;\n\t" - ".reg .b64 db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "" - "mov.b32 dhi, 0x40004040;\n\t" - "mov.b32 id, 136380432;\n\t" - "mov.b32 ta, %2;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p0;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" - "}\n" - :: "r"((tmem_tmem_scratch + (output_col_d))), "r"(_mma_b_lo_12), "r"(tmem_tmem_scratch + p_col_last), "r"(((first_pv) ? 0 : 1))); - elect_commit(kv_empty_addr + (mma_kv_stage) * 8); - } - mma_kv_stage += 1; - if (mma_kv_stage == 8) { mma_kv_stage = 0; _phase_kv_full ^= 1; _phase_kv_pair_ready ^= 1; } - } - } - elect_commit(q_empty_addr); - elect_commit(o_done_addr); - mma_tile_cursor = mma_tile_cursor + num_kv_tiles_3; - } - mbarrier_arrive(tmem_dealloc_addr); - unsigned int _phase_tmem_dealloc_0 = 0; - mbarrier_wait(tmem_dealloc_addr, _phase_tmem_dealloc_0); - _phase_tmem_dealloc_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - { - int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); - asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); - } - } - } - // ---- Role: empty1 ---- - if (warp >= 14 && warp <= 15) { - { // empty1_main - const int wg2_dummy_2 = 0; - unsigned int _phase_q_full_0_4 = 0; - { - #pragma unroll 1 - for (unsigned int work_idx_5 = bid; work_idx_5 < total_work_items; work_idx_5 += num_bids) { - mbarrier_wait(q_full_addr, _phase_q_full_0_4); - _phase_q_full_0_4 ^= 1; - } - } - } - } - - // Cleanup -} - -} // extern "C" diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_7771bcc5c13a5ff83890_binding.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_7771bcc5c13a5ff83890_binding.cu deleted file mode 100644 index 2df363fef17..00000000000 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_7771bcc5c13a5ff83890_binding.cu +++ /dev/null @@ -1,461 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -// Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_7771bcc5c13a5ff83890'. -#include -#include -#include -#include - -#include "tvm_ffi_utils.h" - -#include -#include -#include -#include -#include -#include -#include - -extern "C" __global__ void kernel_cake_dsv4_7771bcc5c13a5ff83890(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items); - - -namespace cake_host_shim_392ec7fc32b7041f { - -using tvm::ffi::Optional; -using tvm::ffi::TensorView; - -class ScopedCudaDevice { - public: - explicit ScopedCudaDevice(int device_id) { - cudaError_t error = cudaGetDevice(&previous_device_); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed before host-shim launch: cudaError=" - << static_cast(error); - if (previous_device_ != device_id) { - error = cudaSetDevice(device_id); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaSetDevice failed before host-shim launch for cuda:" - << device_id << ": cudaError=" << static_cast(error); - restore_ = true; - } - } - - ScopedCudaDevice(const ScopedCudaDevice&) = delete; - ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; - - ~ScopedCudaDevice() noexcept { - if (restore_) { - (void)cudaSetDevice(previous_device_); - } - } - - private: - int previous_device_ = -1; - bool restore_ = false; -}; - -inline int64_t CakeDeviceMultiprocessorCount(int device_id) { - constexpr int kMaxCachedCudaDevices = 64; - TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) - << "physical-SM-count cache does not cover cuda:" << device_id; - static std::atomic count_by_device[kMaxCachedCudaDevices]{}; - int cached = count_by_device[device_id].load(std::memory_order_acquire); - if (cached > 0) return cached; - - int count = 0; - cudaError_t error = cudaDeviceGetAttribute( - &count, cudaDevAttrMultiProcessorCount, device_id); - TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) - << "querying multiProcessorCount failed for cuda:" << device_id - << ": cudaError=" << static_cast(error) << ", count=" << count; - // Concurrent first launches may repeat the immutable device query, but all - // publication is atomic and every later hot-path lookup is one acquire load. - count_by_device[device_id].store(count, std::memory_order_release); - return count; -} - -inline void CheckCudaTensor(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) - << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; -} - -inline void CheckSameCudaDevice( - const TensorView& t, - const TensorView& reference, - const char* name, - const char* reference_name) { - TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) - << name << " must be on the same CUDA device as " << reference_name - << ": got cuda:" << t.device().device_id - << " versus cuda:" << reference.device().device_id; -} - -inline void CheckCurrentCudaDevice( - const TensorView& reference, - const char* reference_name) { - int current_device = -1; - cudaError_t error = cudaGetDevice(¤t_device); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed while validating " << reference_name - << ": cudaError=" << static_cast(error); - TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) - << "current CUDA device must match " << reference_name - << ": current=cuda:" << current_device - << ", tensor=cuda:" << reference.device().device_id; -} - -inline void CheckContiguous(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; -} - -inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { - DLDataType d = t.dtype(); - TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) - << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits - << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits - << ", lanes=" << (int)d.lanes << ")"; -} - -// A logical axis.outer(trailing) folds every source dim above the trailing -// dimensions. Shape products are independent of physical strides, so verify -// the leading dimensions form one dense row-major chain instead of inventing -// a "folded stride". The descriptor reads its exact adjacent physical step -// separately through stride[-(trailing + 1)]. -inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { - TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) - << name << " cannot fold leading dimensions above " << trailing - << " trailing dims from ndim=" << t.ndim(); - int outer_last = t.ndim() - trailing - 1; - if (outer_last <= 0) { - return; - } - int64_t step = t.stride(outer_last); - TVM_FFI_CHECK(step > 0, ValueError) - << name << " physical strides must be positive"; - int64_t expected = step; - for (int axis = outer_last - 1; axis >= 0; --axis) { - expected *= t.size(axis + 1); - if (t.size(axis) > 1) { - TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) - << name << " leading dims are not physically foldable above " << trailing - << " trailing dims: stride(" << axis << ")=" << t.stride(axis) - << ", expected " << expected; - } - } -} - -// 4D TMA descriptor for buffer 'tmap_q' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 3, ValueError) - << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - int64_t d2 = t.size(t.ndim() - 2); - int64_t d3 = t.size(t.ndim() - 3); - TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) - << "TMA source 'tmap_q' trailing dims must be positive"; - TVM_FFI_CHECK(d1 % 128 == 0, ValueError) - << "TMA source 'tmap_q' extent " << d1 - << " must divide exactly by " << 128; - uint64_t global_dim[4] = {(uint64_t)(128), (uint64_t)(d2), (uint64_t)((d1 / 128)), (uint64_t)(d3)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; - TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) - << "TMA box (128, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; - int64_t carrier_stride_0 = d1; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; - int64_t carrier_stride_1 = 128; - TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; - TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; - TVM_FFI_CHECK((carrier_stride_1 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; - int64_t carrier_stride_2 = (d2 * d1); - TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; - TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; - TVM_FFI_CHECK((carrier_stride_2 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; - uint64_t global_strides[3] = { - (uint64_t)((carrier_stride_0 * 8) / 8), - (uint64_t)((carrier_stride_1 * 8) / 8), - (uint64_t)((carrier_stride_2 * 8) / 8), - }; - uint32_t box_dim[4] = {128u, 64u, 1u, 1u}; - uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_swa_kv' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_swa_kv' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; - TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (128, 1) exceeds resolved global dims for 'tmap_swa_kv'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 8) / 8), - }; - uint32_t box_dim[2] = {128u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_compressed_kv' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_compressed_kv(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_compressed_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_compressed_kv' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_compressed_kv' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_compressed_kv"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_compressed_kv' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved a non-positive global dim"; - TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (128, 1) exceeds resolved global dims for 'tmap_compressed_kv'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 8) / 8), - }; - uint32_t box_dim[2] = {128u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_kv') failed: CUresult=" << (int)r; - return tm; -} - -void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_compressed_kv, TensorView arg_O, TensorView arg_partial_lse, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_num_query_tokens, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, int64_t grid_x, int64_t grid_y, int64_t grid_z) { - DLDevice dev = arg_tmap_q.device(); - ScopedCudaDevice device_guard(dev.device_id); - CheckCudaTensor(arg_tmap_q, "tmap_q"); - CheckDtype(arg_tmap_q, "tmap_q", 1, 8, 1); - CheckContiguous(arg_tmap_q, "tmap_q"); - CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); - CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 1, 8, 1); - CheckCudaTensor(arg_tmap_compressed_kv, "tmap_compressed_kv"); - CheckDtype(arg_tmap_compressed_kv, "tmap_compressed_kv", 1, 8, 1); - CheckCudaTensor(arg_O, "O"); - CheckDtype(arg_O, "O", 4, 16, 1); - CheckContiguous(arg_O, "O"); - CheckCudaTensor(arg_partial_lse, "partial_lse"); - CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); - CheckContiguous(arg_partial_lse, "partial_lse"); - CheckCudaTensor(arg_swa_indices, "swa_indices"); - CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); - CheckContiguous(arg_swa_indices, "swa_indices"); - CheckCudaTensor(arg_compressed_indices, "compressed_indices"); - CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); - CheckContiguous(arg_compressed_indices, "compressed_indices"); - CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); - CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckCudaTensor(arg_sinks, "sinks"); - CheckDtype(arg_sinks, "sinks", 2, 32, 1); - CheckContiguous(arg_sinks, "sinks"); - CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); - CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); - CheckContiguous(arg_bmm1_scale, "bmm1_scale"); - CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); - CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); - CheckContiguous(arg_bmm2_scale, "bmm2_scale"); - TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) - << "scalar 'num_heads' value " << arg_num_heads - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) - << "scalar 'swa_index_stride' value " << arg_swa_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) - << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) - << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_num_query_tokens >= -2147483648LL && arg_num_query_tokens <= 2147483647LL, ValueError) - << "scalar 'num_query_tokens' value " << arg_num_query_tokens - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) - << "scalar 'sparse_topk' value " << arg_sparse_topk - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) - << "scalar 'has_sinks' value " << arg_has_sinks - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) - << "scalar 'total_work_items' value " << arg_total_work_items - << " is outside i32 range [-2147483648, 2147483647]"; - CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); - CheckSameCudaDevice(arg_tmap_compressed_kv, arg_tmap_q, "tmap_compressed_kv", "tmap_q"); - CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); - CheckSameCudaDevice(arg_partial_lse, arg_tmap_q, "partial_lse", "tmap_q"); - CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); - CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); - CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); - CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); - CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); - CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); - TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) - << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y - << ", " << grid_z << ")"; - TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) - << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z - << ") must be divisible by cluster dims (2, 1, 1)"; - - cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); - CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); - CUtensorMap p_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); - CUtensorMap p_tmap_compressed_kv = EncodeTma_tmap_compressed_kv(arg_tmap_compressed_kv); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); - float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); - int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); - int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); - int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); - float* p_sinks = static_cast(arg_sinks.data_ptr()); - float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); - float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); - int32_t v_num_heads = (int32_t)arg_num_heads; - int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; - int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; - int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; - int32_t v_num_query_tokens = (int32_t)arg_num_query_tokens; - int32_t v_sparse_topk = (int32_t)arg_sparse_topk; - int32_t v_has_sinks = (int32_t)arg_has_sinks; - int32_t v_total_work_items = (int32_t)arg_total_work_items; - void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_tmap_compressed_kv, &p_O, &p_partial_lse, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_num_query_tokens, &v_sparse_topk, &v_has_sinks, &v_total_work_items}; - - static std::mutex smem_mu; - static auto* smem_status_by_device = - new std::unordered_map(); - cudaError_t smem_status = cudaSuccess; - { - std::lock_guard lock(smem_mu); - auto it = smem_status_by_device->find(dev.device_id); - if (it == smem_status_by_device->end()) { - int current_device = -1; - smem_status = cudaGetDevice(¤t_device); - if (smem_status == cudaSuccess && current_device != dev.device_id) { - smem_status = cudaErrorInvalidDevice; - } - if (smem_status == cudaSuccess) { - smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_7771bcc5c13a5ff83890), - cudaFuncAttributeMaxDynamicSharedMemorySize, - 223744); - } - smem_status_by_device->emplace(dev.device_id, smem_status); - } else { - smem_status = it->second; - } - } - TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_7771bcc5c13a5ff83890 failed: " - << cudaGetErrorString(smem_status); - dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); - dim3 block(512u, 1u, 1u); - - // Extended direct-source launch for a linked CUDA kernel symbol. - cudaLaunchAttribute attrs[2]{}; - int n = 0; - attrs[n].id = cudaLaunchAttributeClusterDimension; - attrs[n].val.clusterDim.x = 2u; - attrs[n].val.clusterDim.y = 1u; - attrs[n].val.clusterDim.z = 1u; - ++n; - attrs[n].id = cudaLaunchAttributeClusterSchedulingPolicyPreference; - attrs[n].val.clusterSchedulingPolicyPreference = cudaClusterSchedulingPolicySpread; - ++n; - cudaLaunchConfig_t config{}; - config.gridDim = grid; - config.blockDim = block; - config.dynamicSmemBytes = 223744u; - config.stream = stream; - config.attrs = attrs; - config.numAttrs = n; - cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_7771bcc5c13a5ff83890), kargs); - TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_7771bcc5c13a5ff83890 failed: " - << cudaGetErrorString(launch_status); -} - -} // namespace cake_host_shim_392ec7fc32b7041f - -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_392ec7fc32b7041f::Run); diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_7771bcc5c13a5ff83890_kernel.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_7771bcc5c13a5ff83890_kernel.cu deleted file mode 100644 index 048415f9878..00000000000 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_7771bcc5c13a5ff83890_kernel.cu +++ /dev/null @@ -1,2803 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -typedef signed char int8_t; -typedef unsigned char uint8_t; -typedef unsigned short uint16_t; -typedef unsigned int uint32_t; -#if defined(__CUDACC_RTC__) -typedef unsigned long long uint64_t; -#else -typedef unsigned long uint64_t; -#endif -static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); -typedef signed int int32_t; -typedef short int int16_t; -struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; -struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; -static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); -static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); - -#if defined(__CUDACC_RTC__) -typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; -#else -#include -#endif - -static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); -static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); -#include -#include - -#define CAKE_INF CUDART_INF_F -#define TMEM_NCOLS 512 -#define TMEM_TMEM_SCRATCH_OFFSET 0 -#define NUM_K_PIPE_STAGES 3 -#define NUM_V_PIPE_STAGES 2 -#define NUM_KV_PIPE_STAGES 1 -#define NUM_INDEX_PIPE_STAGES 6 -#define SMEM_SMEM_Q_OFF 1024 -#define SMEM_SMEM_Q_STAGE_BYTES 8192 -#define SMEM_SMEM_Q_STRIDE 8192 -#define SMEM_SMEM_KV_OFF 33792 -#define SMEM_SMEM_KV_STAGE_BYTES 16384 -#define SMEM_SMEM_KV_STRIDE 16384 -#define SMEM_SMEM_V_OFF 33792 -#define SMEM_SMEM_V_STAGE_BYTES 16384 -#define SMEM_SMEM_V_STRIDE 16384 -#define SMEM_SMEM_Q_FULL_OFF 1024 -#define SMEM_SMEM_Q_FULL_STAGE_BYTES 32768 -#define SMEM_SMEM_Q_FULL_STRIDE 32768 -#define SMEM_SMEM_K_FULL_OFF 33792 -#define SMEM_SMEM_K_FULL_STAGE_BYTES 32768 -#define SMEM_SMEM_K_FULL_STRIDE 32768 -#define SMEM_SMEM_V_FULL_OFF 132096 -#define SMEM_SMEM_V_FULL_STAGE_BYTES 32768 -#define SMEM_SMEM_V_FULL_STRIDE 32768 -#define SMEM_SMEM_STATS_MAX_OFF 197632 -#define SMEM_SMEM_STATS_MAX_STAGE_BYTES 1024 -#define SMEM_SMEM_STATS_MAX_STRIDE 1024 -#define SMEM_SMEM_STATS_SUM_OFF 198656 -#define SMEM_SMEM_STATS_SUM_STAGE_BYTES 512 -#define SMEM_SMEM_STATS_SUM_STRIDE 512 -#define SMEM_SMEM_STATS_FINAL_MAX_OFF 199168 -#define SMEM_SMEM_STATS_FINAL_MAX_STAGE_BYTES 512 -#define SMEM_SMEM_STATS_FINAL_MAX_STRIDE 512 -#define SMEM_SMEM_SOFTMAX_WARP_PAIR_EXCHANGE_OFF 199680 -#define SMEM_SMEM_SOFTMAX_WARP_PAIR_EXCHANGE_STAGE_BYTES 1024 -#define SMEM_SMEM_SOFTMAX_WARP_PAIR_EXCHANGE_STRIDE 1024 -#define SMEM_SMEM_CORR_WARP_PAIR_EXCHANGE_OFF 200704 -#define SMEM_SMEM_CORR_WARP_PAIR_EXCHANGE_STAGE_BYTES 512 -#define SMEM_SMEM_CORR_WARP_PAIR_EXCHANGE_STRIDE 512 -#define SMEM_SMEM_SPARSE_INDICES_OFF 201216 -#define SMEM_SMEM_SPARSE_INDICES_STAGE_BYTES 1024 -#define SMEM_SMEM_SPARSE_INDICES_STRIDE 1024 -#define SMEM_SMEM_P_FP8_OFF 207360 -#define SMEM_SMEM_P_FP8_STAGE_BYTES 8192 -#define SMEM_SMEM_P_FP8_STRIDE 8192 -#define SMEM_TOTAL 223744 -#define THREADS 512 - -#include - -__device__ __forceinline__ uint32_t elect_sync() { - uint32_t pred = 0; - asm volatile( - "{\n\t" - ".reg .pred %%px;\n\t" - "elect.sync _|%%px, %1;\n\t" - "@%%px mov.s32 %0, 1;\n\t" - "}\n" - : "+r"(pred) - : "r"(0xFFFFFFFF)); - return pred; -} - - -__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { - asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" - :: "r"(mbar_addr), "r"(count) : "memory"); -} - -__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { - asm volatile("mbarrier.init.b64 [%0], %1;" - :: "l"(mbar_addr), "r"(count) : "memory"); -} - - -__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - - -// CTA-local pipelines have short, resident producer/consumer edges. Omitting -// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the -// explicit loop still makes this helper blocking until acquire succeeds. -__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE;\n\t" - "bra.uni LAB_WAIT;\n\t" - "DONE:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Source-faithful relaxed CTA wait used only by a typed protocol that does -// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. -__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, 10000000;\n\t" - "@P1 bra.uni DONE_RELAXED;\n\t" - "bra.uni LAB_WAIT_RELAXED;\n\t" - "DONE_RELAXED:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Exact source ports may request the PTX suspendTimeHint operand explicitly. -// The hint is expressed in nanoseconds and is kept separate from the canonical -// no-hint CTA helper so unrelated schedules retain their existing retry path. -__device__ __forceinline__ void mbarrier_wait_suspend( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_SUSPEND:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_SUSPEND;\n\t" - "bra.uni LAB_WAIT_SUSPEND;\n\t" - "DONE_SUSPEND:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE_CLUSTER;\n\t" - "bra.uni LAB_WAIT_CLUSTER;\n\t" - "DONE_CLUSTER:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - ".reg .u32 WAIT_ADDR;\n\t" - "mov.u32 WAIT_ADDR, %0;\n\t" - "LAB_WAIT_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [WAIT_ADDR], %1, %2;\n\t" - "@P1 bra.uni DONE_HINT;\n\t" - "bra.uni LAB_WAIT_HINT;\n\t" - "DONE_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -// Exact unqualified CTA wait used by source schedules whose PTX intentionally -// omits the acquire qualifier while retaining a typed suspendTimeHint operand. -__device__ __forceinline__ void mbarrier_wait_relaxed_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED_HINT:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra DONE_RELAXED_HINT;\n\t" - "bra LAB_WAIT_RELAXED_HINT;\n\t" - "DONE_RELAXED_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); -} - -__device__ __forceinline__ void mbarrier_wait_cluster_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" - "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" - "DONE_CLUSTER_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_suspend( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait_cluster(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); - } -} - - -__device__ __forceinline__ void tcgen05_mma_f8f6f4_cta2( - int taddr, uint64_t a_desc, uint64_t b_desc, - uint32_t i_desc, int enable_input_d) { - asm volatile( - "{\n\t" - ".reg .pred p;\n\t" - ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - "setp.ne.b32 p, %4, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" - "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "tcgen05.mma.cta_group::2.kind::f8f6f4 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" - "}\n" - :: "r"(taddr), "l"(a_desc), "l"(b_desc), - "r"(i_desc), "r"(enable_input_d) - : "memory"); -} - - -__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { - return (x & 0x3FFFFULL) >> 4ULL; -} - - -union MmaSmemDesc { - uint64_t u64; - uint32_t u32[2]; -}; - -__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { - MmaSmemDesc tmp; - tmp.u64 = smem_desc; - tmp.u32[0] += offset; - smem_desc = tmp.u64; -} - - -__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { - asm volatile( - "{\n\t" - ".reg .pred leader;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" - ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" - "}\n" - :: "r"(mbar_addr), "h"(cta_mask) : "memory"); -} - - -__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { - asm volatile( - "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" - :: "r"(mbar_addr), "r"(bytes) : "memory"); -} - - -__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7," - " %8, %9, %10, %11, %12, %13, %14, %15," - " %16, %17, %18, %19, %20, %21, %22, %23," - " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), - "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), - "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), - "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), - "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), - "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), - "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x16.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7," - " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), - "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), - "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ void tmem_ld_x4(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x4.b32" - " {%0, %1, %2, %3}, [%4];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { - asm volatile( - "tcgen05.st.sync.aligned.32x32b.x16.b32" - " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," - " %9, %10, %11, %12, %13, %14, %15, %16};" - :: "r"(tmem_addr), - "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), - "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), - "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), - "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); -} - - -__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { - asm volatile( - "tcgen05.st.sync.aligned.32x32b.x4.b32" - " [%0], {%1, %2, %3, %4};" - :: "r"(tmem_addr), - "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); -} - - -__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { - uint32_t addr; - asm("{\n\t" - ".reg .u64 u64addr;\n\t" - "cvta.to.shared.u64 u64addr, %1;\n\t" - "cvt.u32.u64 %0, u64addr;\n\t" - "}\n" : "=r"(addr) : "l"(ptr)); - return addr; -} - - -__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { - uint32_t remote; - asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" - : "=r"(remote) : "r"(local_addr), "r"(rank)); - return remote; -} - - -__device__ __forceinline__ float approx_exp2(float x) { - float y; - asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float approx_rcp(float x) { - float y; - asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float max_noftz(float a, float b) { - float c; - asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); - return c; -} - - -__device__ __forceinline__ float warp_reduce_max(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); - return val; -} - - -__device__ __forceinline__ float warp_reduce_sum(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val += __shfl_xor_sync(0xFFFFFFFF, val, offset); - return val; -} - - -__device__ __forceinline__ float row_max_reduce(float2 acc) { - return max_noftz(acc.x, acc.y); -} - - -__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { - #pragma unroll - for (int j = 0; j < 16; j++) { - if (j % 2 == 0) - acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); - else - acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); - } -} - - -__device__ __forceinline__ float2 ex2_emulation_f32x2_value(float2 value) { - const float c0 = 1.0f, c1 = 0.695146143436431884765625f; - const float c2 = 0.227564394474029541015625f, c3 = 0.077119089663028717041015625f; - const float magic = 12582912.0f; - float x0 = max_noftz(value.x, -127.0f), x1 = max_noftz(value.y, -127.0f); - float2 xc2 = make_float2(x0, x1), magic2 = make_float2(magic, magic); - float2 xr2; - asm("add.rm.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xr2) - : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&magic2)); - float2 c3_2 = make_float2(c3, c3), c2_2 = make_float2(c2, c2); - float2 c1_2 = make_float2(c1, c1), c0_2 = make_float2(c0, c0); - float2 xrb2, xfrac2; - asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xrb2) - : "l"(*(unsigned long long*)&xr2), "l"(*(unsigned long long*)&magic2)); - asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xfrac2) - : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&xrb2)); - float2 poly2; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&c3_2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c2_2)); - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c1_2)); - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c0_2)); - int x0r_i, x1r_i, p0_i, p1_i; - asm("mov.b64 {%0, %1}, %2;" : "=r"(x0r_i), "=r"(x1r_i) : "l"(*(unsigned long long*)&xr2)); - asm("mov.b64 {%0, %1}, %2;" : "=r"(p0_i), "=r"(p1_i) : "l"(*(unsigned long long*)&poly2)); - float r0, r1; - asm("mov.b32 %0, %1;" : "=f"(r0) : "r"((x0r_i << 23) + p0_i)); - asm("mov.b32 %0, %1;" : "=f"(r1) : "r"((x1r_i << 23) + p1_i)); - return make_float2(r0, r1); -} - -__device__ __forceinline__ void ex2_emulation_f32x2(float* x0_ptr, float* x1_ptr) { - float2 result = ex2_emulation_f32x2_value(make_float2(*x0_ptr, *x1_ptr)); - *x0_ptr = result.x; *x1_ptr = result.y; -} - -__device__ __forceinline__ void softmax_frag_exp2_cast( - float* sv, uint32_t* pv, int use_emu) -{ - #pragma unroll - for (int j = 0; j < 16; j++) { - if (use_emu && j >= 12) - ex2_emulation_f32x2(&sv[j*2], &sv[j*2+1]); - else { - sv[j*2] = approx_exp2(sv[j*2]); - sv[j*2+1] = approx_exp2(sv[j*2+1]); - } - } - #pragma unroll - for (int j = 0; j < 16; j++) { - __nv_bfloat162 bf = __float22bfloat162_rn({sv[j*2], sv[j*2+1]}); - pv[j] = reinterpret_cast(bf); - } -} - - - -__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { - const float2* sv2 = reinterpret_cast(sv); - #pragma unroll - for (int j = 0; j < 16; j++) { - asm("add.f32x2 %0, %1, %2;" - : "+l"(reinterpret_cast(*acc)) - : "l"(reinterpret_cast(*acc)), - "l"(reinterpret_cast(sv2[j]))); - } -} - - -__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { - asm("mul.rn.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { - asm("mul.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { - asm("add.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { - asm("add.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { - asm("sub.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { - asm("sub.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("add.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { - float2 r; - asm("sub.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("sub.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ void fma_scale_x32( - float* sv, const float2* scale2, const float2* neg_max2) -{ - float2* sv_2 = reinterpret_cast(sv); - #pragma unroll - for (int j = 0; j < 16; j++) - fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); -} - -__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("mul.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - - -// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) - -__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("add.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("add.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("add.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("add.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("add.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("add.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("add.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - - -__device__ __forceinline__ void fence_async_shared() { - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); -} - - -__device__ __forceinline__ uint64_t make_smem_desc(int addr) { - const int SBO = 1024; - return desc_encode(addr) - | (desc_encode(SBO) << 32ULL) - | (1ULL << 46ULL) - | (2ULL << 61ULL); -} - - -__device__ __forceinline__ void tma_2d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3}], [%4];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_4d_gmem2smem_cta2( - int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.4d.shared::cluster.global" - ".mbarrier::complete_tx::bytes.cta_group::2" - " [%0], [%1, {%2, %3, %4, %5}], [%6];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr) { - // Canonical .shared::cta form for non-multicast gather4, matching - // trtllm-gen / cuda_ptx and the PTX ISA qualifier order - // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, - // .shared::cluster is reserved for the multicast variant (ctaMask). - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem_cta2( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr) { - // Canonical .shared::cta form; see tma_gather4_gmem2smem above. - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" - ".mbarrier::complete_tx::bytes.cta_group::2" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem_mc( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr, unsigned short cta_mask) { - // Multicast variant: the PTX grammar ties the .shared::cluster - // destination to .multicast::cluster + ctaMask (cf. cuda_ptx / - // SM100_TMA_LOAD_MULTICAST_2D_GATHER4 in CUTLASS). - asm volatile( - "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" - ".mbarrier::complete_tx::bytes.multicast::cluster" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr), "h"(cta_mask) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem_mc_cta2( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr, unsigned short cta_mask) { - // Multicast + cta_group::2 variant; see tma_gather4_gmem2smem_mc. - asm volatile( - "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" - ".mbarrier::complete_tx::bytes.multicast::cluster.cta_group::2" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr), "h"(cta_mask) : "memory"); -} - - -__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { - asm volatile( - "{\n\t" - ".reg .b16 lo, hi;\n\t" - "mov.b32 {lo, hi}, %1;\n\t" - "tcgen05.commit.cta_group::2.mbarrier::arrive::one" - ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" - "}\n" - :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); -} - - -__device__ __forceinline__ void tmem_st_x8_u32(int addr, const uint32_t* src) { - asm volatile( - "tcgen05.st.sync.aligned.32x32b.x8.b32" - " [%0], {%1,%2,%3,%4,%5,%6,%7,%8};" - :: "r"(addr), - "r"(src[0]), "r"(src[1]), "r"(src[2]), "r"(src[3]), - "r"(src[4]), "r"(src[5]), "r"(src[6]), "r"(src[7])); -} - - -__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { - uint32_t result; - asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" - : "=r"(result) : "r"(val)); - return result; -} - -extern "C" { - -__global__ __launch_bounds__(512, 1) __cluster_dims__(2,1,1) void -kernel_cake_dsv4_7771bcc5c13a5ff83890(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items) -{ - const int tid = threadIdx.x; - const int warp = make_warp_uniform(tid / 32); - const int lane = tid % 32; - - extern __shared__ __align__(1024) char smem_raw[]; - int smem; - asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); - smem = make_warp_uniform(smem); - - const int mbar_base = smem; - #define q_full_addr (mbar_base + 0) - #define q_empty_addr (mbar_base + 8) - #define k_full_addr (mbar_base + 16) - #define k_empty_addr (mbar_base + 40) - #define v_full_addr (mbar_base + 64) - #define v_empty_addr (mbar_base + 80) - #define index_full_addr (mbar_base + 96) - #define index_empty_addr (mbar_base + 144) - #define s_full_addr (mbar_base + 192) - #define p_full_addr (mbar_base + 208) - #define corr_done_addr (mbar_base + 224) - #define stats_addr (mbar_base + 240) - #define sum_ready_addr (mbar_base + 256) - #define o_done_addr (mbar_base + 264) - #define pv_done_addr (mbar_base + 272) - #define s_seeded_addr (mbar_base + 280) - #define q_pair_ready_addr (mbar_base + 288) - #define kv_pair_ready_addr (mbar_base + 296) - #define pv_pair_ready_addr (mbar_base + 304) - #define tmem_dealloc_addr (mbar_base + 320) - #define tmem_dealloc_peer_addr (mbar_base + 328) - - const int bid = blockIdx.x; - const int num_bids = gridDim.x; - const unsigned int clusters_x = gridDim.x / 2; - const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; - const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; - - int cta_rank; - asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); - - // Kernel setup ops - uint8_t* smem_q = reinterpret_cast(smem_raw + 1024); - const int smem_q_addr = smem + 1024; - uint8_t* smem_kv = reinterpret_cast(smem_raw + 33792); - const int smem_kv_addr = smem + 33792; - uint8_t* smem_v = reinterpret_cast(smem_raw + 33792); - const int smem_v_addr = smem + 33792; - uint8_t* smem_q_full = reinterpret_cast(smem_raw + 1024); - const int smem_q_full_addr = smem + 1024; - uint8_t* smem_k_full = reinterpret_cast(smem_raw + 33792); - const int smem_k_full_addr = smem + 33792; - uint8_t* smem_v_full = reinterpret_cast(smem_raw + 132096); - const int smem_v_full_addr = smem + 132096; - float* smem_stats_max = reinterpret_cast(smem_raw + 197632); - const int smem_stats_max_addr = smem + 197632; - float* smem_stats_sum = reinterpret_cast(smem_raw + 198656); - const int smem_stats_sum_addr = smem + 198656; - float* smem_stats_final_max = reinterpret_cast(smem_raw + 199168); - const int smem_stats_final_max_addr = smem + 199168; - float* smem_softmax_warp_pair_exchange = reinterpret_cast(smem_raw + 199680); - const int smem_softmax_warp_pair_exchange_addr = smem + 199680; - float* smem_corr_warp_pair_exchange = reinterpret_cast(smem_raw + 200704); - const int smem_corr_warp_pair_exchange_addr = smem + 200704; - int* smem_sparse_indices = reinterpret_cast(smem_raw + 201216); - const int smem_sparse_indices_addr = smem + 201216; - uint8_t* smem_p_fp8 = reinterpret_cast(smem_raw + 207360); - const int smem_p_fp8_addr = smem + 207360; - asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); - asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); - - // Mbarrier init (21 pipeline groups, 0 ordered-sequence groups, 42 barriers) - // Mbarriers at smem_raw[0..336) - - if (warp == 12) { - uint32_t leader = elect_sync(); - if (leader) { - // q_full: 1 barriers, init_count=2 - mbarrier_init(smem + 0, 2); - // q_empty: 1 barriers, init_count=1 - mbarrier_init(smem + 8, 1); - // --- pipeline 'k_pipe' --- - // k_full: 3 barriers, init_count=2 - mbarrier_init(smem + 16, 2); - mbarrier_init(smem + 24, 2); - mbarrier_init(smem + 32, 2); - // k_empty: 3 barriers, init_count=1 - mbarrier_init(smem + 40, 1); - mbarrier_init(smem + 48, 1); - mbarrier_init(smem + 56, 1); - // --- pipeline 'v_pipe' --- - // v_full: 2 barriers, init_count=2 - mbarrier_init(smem + 64, 2); - mbarrier_init(smem + 72, 2); - // v_empty: 2 barriers, init_count=1 - mbarrier_init(smem + 80, 1); - mbarrier_init(smem + 88, 1); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - if (warp == 9) { - uint32_t leader = elect_sync(); - if (leader) { - // --- pipeline 'index_pipe' --- - // index_full: 6 barriers, init_count=32 - mbarrier_init(smem + 96, 32); - mbarrier_init(smem + 104, 32); - mbarrier_init(smem + 112, 32); - mbarrier_init(smem + 120, 32); - mbarrier_init(smem + 128, 32); - mbarrier_init(smem + 136, 32); - // index_empty: 6 barriers, init_count=256 - mbarrier_init(smem + 144, 256); - mbarrier_init(smem + 152, 256); - mbarrier_init(smem + 160, 256); - mbarrier_init(smem + 168, 256); - mbarrier_init(smem + 176, 256); - mbarrier_init(smem + 184, 256); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - if (warp == 0) { - uint32_t leader = elect_sync(); - if (leader) { - // s_full: stages (0,), init_count=1 - mbarrier_init(smem + 192, 1); - // o_done: 1 barriers, init_count=1 - mbarrier_init(smem + 264, 1); - // pv_done: 1 barriers, init_count=1 - mbarrier_init(smem + 272, 1); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - if (warp == 1) { - uint32_t leader = elect_sync(); - if (leader) { - // s_full: stages (1,), init_count=1 - mbarrier_init(smem + 200, 1); - // p_full: 2 barriers, init_count=128 - mbarrier_init(smem + 208, 128); - mbarrier_init(smem + 216, 128); - // corr_done: stages (0,), init_count=128 - mbarrier_init(smem + 224, 128); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - if (warp == 5) { - uint32_t leader = elect_sync(); - if (leader) { - // corr_done: stages (1,), init_count=128 - mbarrier_init(smem + 232, 128); - // stats: 2 barriers, init_count=128 - mbarrier_init(smem + 240, 128); - mbarrier_init(smem + 248, 128); - // sum_ready: 1 barriers, init_count=128 - mbarrier_init(smem + 256, 128); - // tmem_dealloc_peer: 1 barriers, init_count=32 - mbarrier_init(smem + 328, 32); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - if (warp == 2) { - uint32_t leader = elect_sync(); - if (leader) { - // s_seeded: 1 barriers, init_count=256 - mbarrier_init(smem + 280, 256); - // q_pair_ready: 1 barriers, init_count=64 - mbarrier_init(smem + 288, 64); - // kv_pair_ready: 1 barriers, init_count=64 - mbarrier_init(smem + 296, 64); - // pv_pair_ready: stages (0,), init_count=64 - mbarrier_init(smem + 304, 64); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - if (warp == 4) { - uint32_t leader = elect_sync(); - if (leader) { - // pv_pair_ready: stages (1,), init_count=64 - mbarrier_init(smem + 312, 64); - // tmem_dealloc: 1 barriers, init_count=416 - mbarrier_init(smem + 320, 416); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - - __syncwarp(); - - asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); - asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); - - // TMEM alloc (512 columns, 512 used) - volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 336); - if (warp == 0) { - int _tmem_hold = smem + 336; - asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); - __syncwarp(); - asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); - } - - __syncthreads(); - asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); - asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); - - const int taddr = tmem_addr_storage[0]; - - // Kernel post-init ops - const int tmem_tmem_scratch = taddr; - - // ---- Ordered hardware-WG register redistribution ---- - // Inc phase consumes the registers released above. - if (warp >= 8 && warp <= 11) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 136;"); - } - - // ---- Role: index_warp ---- - if (warp == 9) { - { // index_warp_main - const int index_dummy = 0; - unsigned int _phase_index_empty = 1; - unsigned int _phase_q_full_0 = 0; - { - int all_num_kv_tiles = (sparse_topk + 128 - 1) / 128; - unsigned int index_prod_stage = 0; - #pragma unroll - for (unsigned int static_work = 0; static_work < 1; static_work++) { - int split_idx = 0; - int query_idx = bid / 2; - int tiles_per_split = all_num_kv_tiles + 1 - 1; - int first_tile = split_idx * tiles_per_split; - int num_index_passes = (tiles_per_split + 1) / 2; - #pragma unroll 1 - for (int index_pass = 0; index_pass < num_index_passes; index_pass++) { - mbarrier_wait(index_empty_addr + (index_prod_stage) * 8, _phase_index_empty); - int index_stage_base = smem_sparse_indices_addr + index_prod_stage * 1024; - int stage_row_base = index_pass * 256; - #pragma unroll - for (int index_half = 0; index_half < 2; index_half++) { - int index_offset = index_half * 128 + lane * 4; - int tile_col_base = first_tile * 128 + stage_row_base + index_half * 128; - int* row_tile_ptr = ((tile_col_base < 128) ? swa_indices + (query_idx * swa_index_stride) : compressed_indices + (query_idx * compressed_index_stride + (tile_col_base - 128))); - int lane_col = tile_col_base + lane * 4; - if (lane_col + 4 <= sparse_topk) { - asm volatile("cp.async.cg.shared::cta.global.L2::128B [%0], [%1], 16;" - :: "r"(index_stage_base + index_offset * 4), "l"(row_tile_ptr + (lane * 4))); - } else { - int tail_rows[4]; - #pragma unroll - for (int row_i = 0; row_i < 4; row_i++) { - tail_rows[row_i] = -1; - } - #pragma unroll - for (int row_i_1 = 0; row_i_1 < 4; row_i_1++) { - if (lane_col + row_i_1 < sparse_topk) { - tail_rows[row_i_1] = row_tile_ptr[lane * 4 + row_i_1]; - } - } - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(smem_sparse_indices_addr + ((unsigned int)index_stage_base - smem_sparse_indices_addr + (unsigned int)(index_offset * 4))), "r"(tail_rows[0]), "r"(tail_rows[1]), "r"(tail_rows[2]), "r"(tail_rows[3]) : "memory"); - } - } - asm volatile("cp.async.commit_group;"); - asm volatile("cp.async.wait_group 0;"); - mbarrier_arrive(index_full_addr + (index_prod_stage) * 8); - index_prod_stage += 1; - if (index_prod_stage == 6) { index_prod_stage = 0; _phase_index_empty ^= 1; } - } - } - } - } - } - // ---- Role: load_warp ---- - if (warp >= 12 && warp <= 15) { - asm volatile("setmaxnreg.dec.sync.aligned.u32 72;"); - { // load_warp_main - const int wg2_dummy = 0; - const int load_warp_rank = warp - ((1) ? 12 : 9); - int all_num_kv_tiles_1 = (sparse_topk + 128 - 1) / 128; - unsigned int load_k_stage = 0; - unsigned int load_v_stage = 0; - unsigned int load_kv_stage = 0; - unsigned int load_k_index_stage = 0; - unsigned int load_v_index_stage = 0; - int k_cta_offset = bid % 2 * 64; - unsigned int _phase_q_empty_0 = 1; - unsigned int _phase_index_full = 0; - unsigned int _phase_k_empty = 1; - unsigned int _phase_k_empty_1 = 1; - unsigned int _phase_v_empty = 1; - #pragma unroll - for (unsigned int work_idx = 0; work_idx < 1; work_idx++) { - int split_idx_1 = 0; - int query_idx_1 = bid / 2 >> 1; - int v_chunk = bid / 2 & 1; - { - split_idx_1 = 0; - query_idx_1 = bid / 2; - v_chunk = bid % 2; - } - int tiles_per_split_1 = all_num_kv_tiles_1 + 1 - 1; - int first_tile_1 = split_idx_1 * tiles_per_split_1; - int num_kv_tiles = tiles_per_split_1; - int sparse_extent = num_kv_tiles * 128; - mbarrier_wait(q_empty_addr, _phase_q_empty_0); - _phase_q_empty_0 ^= 1; - if (load_warp_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); - #pragma unroll - for (int q_stage = 0; q_stage < 4; q_stage++) { - asm volatile( - "cp.async.bulk.tensor.4d.shared::cluster.global.mbarrier::complete_tx::bytes.multicast::cluster.cta_group::2" - " [%0], [%1, {%2, %3, %4, %5}], [%6], %7;" - :: "r"(smem_q_full_addr + (unsigned int)(q_stage * 8192)), "l"((&tmap_q)), "r"(0), "r"(bid % 2 * 64), "r"(q_stage), "r"(query_idx_1), - "r"(((q_full_addr) & 0xFEFFFFFF)), "h"((uint16_t)(1 << bid % 2)) : "memory"); - } - } - } - { - mbarrier_wait(index_full_addr + (load_k_index_stage) * 8, _phase_index_full); - asm volatile("tcgen05.fence::after_thread_sync;"); - } - int k_index_stage_base = smem_sparse_indices_addr + load_k_index_stage * 1024; - { - mbarrier_wait(k_empty_addr + (load_k_stage) * 8, _phase_k_empty); - if (load_warp_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); - } - } - int k_dst = smem_k_full_addr + load_k_stage * 32768 + (unsigned int)(load_warp_rank * 8192); - int group = lane; - if (group < 16) { - int group_offset = k_cta_offset + group * 4; - int raw_rows[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows[0])), "=r"(*reinterpret_cast(&raw_rows[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows[(0) + 3])) - : "r"(k_index_stage_base + group_offset * 4)); - int row0 = ((raw_rows[0] >= 0) ? raw_rows[0] : 0); - int row1 = ((raw_rows[1] >= 0) ? raw_rows[1] : 0); - int row2 = ((raw_rows[2] >= 0) ? raw_rows[2] : 0); - int row3 = ((raw_rows[3] >= 0) ? raw_rows[3] : 0); - if (first_tile_1 == 0) { - tma_gather4_gmem2smem_mc_cta2(k_dst + group * 512, (&tmap_swa_kv), load_warp_rank * 128, row0, row1, row2, row3, ((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } else { - tma_gather4_gmem2smem_mc_cta2(k_dst + group * 512, (&tmap_compressed_kv), load_warp_rank * 128, row0, row1, row2, row3, ((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } - } - load_k_stage += 1; - if (load_k_stage == 3) { load_k_stage = 0; _phase_k_empty ^= 1; } - } - #pragma unroll 1 - for (int tile = 1; tile < num_kv_tiles; tile++) { - if ((1 & (int)((tile & 1) == 0)) != 0) { - mbarrier_wait(index_full_addr + (load_k_index_stage) * 8, _phase_index_full); - asm volatile("tcgen05.fence::after_thread_sync;"); - } - k_index_stage_base = smem_sparse_indices_addr + load_k_index_stage * 1024; - { - mbarrier_wait(k_empty_addr + (load_k_stage) * 8, _phase_k_empty); - if (load_warp_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); - } - } - int k_dst_1 = smem_k_full_addr + load_k_stage * 32768 + (unsigned int)(load_warp_rank * 8192); - int group_1 = lane; - if (group_1 < 16) { - int group_offset_1 = (tile & 1) * 128 + k_cta_offset + group_1 * 4; - int raw_rows_1[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_1[0])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 3])) - : "r"(k_index_stage_base + group_offset_1 * 4)); - int row0_1 = ((raw_rows_1[0] >= 0) ? raw_rows_1[0] : 0); - int row1_1 = ((raw_rows_1[1] >= 0) ? raw_rows_1[1] : 0); - int row2_1 = ((raw_rows_1[2] >= 0) ? raw_rows_1[2] : 0); - int row3_1 = ((raw_rows_1[3] >= 0) ? raw_rows_1[3] : 0); - tma_gather4_gmem2smem_mc_cta2(k_dst_1 + group_1 * 512, (&tmap_compressed_kv), load_warp_rank * 128, row0_1, row1_1, row2_1, row3_1, ((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } - load_k_stage += 1; - if (load_k_stage == 3) { load_k_stage = 0; _phase_k_empty ^= 1; } - if ((tile & 1) != 0) { - load_k_index_stage += 1; - if (load_k_index_stage == 6) { load_k_index_stage = 0; _phase_index_full ^= 1; } - } - } - int prev_tile = tile - 1; - int global_prev_tile = first_tile_1 + prev_tile; - int v_index_stage_base = smem_sparse_indices_addr + load_v_index_stage * 1024; - { - mbarrier_wait(v_empty_addr + (load_v_stage) * 8, _phase_v_empty); - if (load_warp_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); - } - } - if (load_warp_rank < 2) { - int v_dst = smem_v_full_addr + load_v_stage * 32768 + (unsigned int)(load_warp_rank * 16384); - int v_col = v_chunk * 256 + load_warp_rank * 128; - int group_2 = lane; - int group_offset_2 = (prev_tile & 1) * 128 + group_2 * 4; - int raw_rows_2[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_2[0])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 3])) - : "r"(v_index_stage_base + group_offset_2 * 4)); - int row0_2 = ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : 0); - int row1_2 = ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : 0); - int row2_2 = ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : 0); - int row3_2 = ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : 0); - if (global_prev_tile == 0) { - tma_gather4_gmem2smem_mc_cta2(v_dst + group_2 * 512, (&tmap_swa_kv), v_col, row0_2, row1_2, row2_2, row3_2, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } else { - tma_gather4_gmem2smem_mc_cta2(v_dst + group_2 * 512, (&tmap_compressed_kv), v_col, row0_2, row1_2, row2_2, row3_2, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } - } - load_v_stage += 1; - if (load_v_stage == 2) { load_v_stage = 0; _phase_v_empty ^= 1; } - if ((prev_tile & 1) != 0) { - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(index_empty_addr + (load_v_index_stage) * 8); - load_v_index_stage = load_v_index_stage + 1; - if (load_v_index_stage == 6) { - load_v_index_stage = 0; - } - } - } - } - int last_tile = num_kv_tiles - 1; - int global_last_tile = first_tile_1 + last_tile; - { - mbarrier_wait(v_empty_addr + (load_v_stage) * 8, _phase_v_empty); - if (load_warp_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); - } - } - if (load_warp_rank < 2) { - int v_dst_1 = smem_v_full_addr + load_v_stage * 32768 + (unsigned int)(load_warp_rank * 16384); - int v_col_1 = v_chunk * 256 + load_warp_rank * 128; - int group_3 = lane; - int group_offset_3 = (last_tile & 1) * 128 + group_3 * 4; - int raw_rows_3[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_3[0])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 3])) - : "r"(smem_sparse_indices_addr + load_v_index_stage * 1024 + (unsigned int)(group_offset_3 * 4))); - int row0_3 = ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : 0); - int row1_3 = ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : 0); - int row2_3 = ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : 0); - int row3_3 = ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : 0); - if (global_last_tile == 0) { - tma_gather4_gmem2smem_mc_cta2(v_dst_1 + group_3 * 512, (&tmap_swa_kv), v_col_1, row0_3, row1_3, row2_3, row3_3, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } else { - tma_gather4_gmem2smem_mc_cta2(v_dst_1 + group_3 * 512, (&tmap_compressed_kv), v_col_1, row0_3, row1_3, row2_3, row3_3, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } - } - load_v_stage += 1; - if (load_v_stage == 2) { load_v_stage = 0; _phase_v_empty ^= 1; } - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(index_empty_addr + (load_v_index_stage) * 8); - load_v_index_stage = load_v_index_stage + 1; - if (load_v_index_stage == 6) { - load_v_index_stage = 0; - } - } - } - mbarrier_arrive(tmem_dealloc_addr); - } - } - // ---- Role: softmax_wg ---- - if (warp <= 3) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 152;"); - { // softmax_wg_main - float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; - const int wg_dummy_inc = 0; - int all_num_kv_tiles_2 = (sparse_topk + 128 - 1) / 128; - const int tmem_row_base = ((1) ? warp % 2 * 32 : warp % 4 * 32); - const int tmem_score_row_base = ((1) ? (int)(warp % 4 * 32) : tmem_row_base); - const int n_half = ((1) ? (int)(warp % 4 / 2) : 0); - const int my_row = tmem_row_base + lane; - const int stats_row = n_half * 64 + my_row; - int softmax_tile_cursor = 0; - unsigned int softmax_index_stage = 0; - unsigned int _phase_q_full_0_1 = 0; - #pragma unroll - for (unsigned int work_idx_1 = 0; work_idx_1 < 1; work_idx_1++) { - int split_idx_2 = 0; - int query_idx_2 = bid / 2 >> 1; - { - split_idx_2 = 0; - query_idx_2 = bid / 2; - } - int tiles_per_split_2 = all_num_kv_tiles_2 + 1 - 1; - int first_tile_2 = split_idx_2 * tiles_per_split_2; - int num_kv_tiles_1 = tiles_per_split_2; - int _max_0 = ((sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) : (0)); - int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); - int active_topk = _min_0; - float row_max_val = -CAKE_INF; - float row_sum_val = 0.0f; - int sink_head = ((1) ? bid % 2 * 64 + my_row : my_row); - if (has_sinks != 0 && sink_head < num_heads && split_idx_2 == 0) { - row_max_val = sinks[sink_head] * 1.4426950408889634f / softmax_scale_log2; - row_sum_val = 1.0f; - } - #pragma unroll 1 - for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { - int pipeline_tile = softmax_tile_cursor + tile_1; - int phase = pipeline_tile & 1; - int s_wait_phase = pipeline_tile >> 1 & 1; - mbarrier_wait(s_full_addr + (phase) * 8, s_wait_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - int s_off = ((phase != 0) ? 128 : 0); - int s_base = taddr + (unsigned int)s_off + (unsigned int)(tmem_score_row_base << 16); - float new_max = row_max_val; - int valid_sparse_cols = ((active_topk < sparse_topk) ? active_topk : sparse_topk); - valid_sparse_cols = valid_sparse_cols - (first_tile_2 + tile_1) * 128 - n_half * 64; - if (valid_sparse_cols < 0) { - valid_sparse_cols = 0; - } - if (valid_sparse_cols > 64) { - valid_sparse_cols = 64; - } - int index_tile_addr = smem_sparse_indices_addr + softmax_index_stage * 1024 + (unsigned int)((tile_1 & 1) * 512); - int index_lane_addr = index_tile_addr + (n_half * 64 + lane) * 4; - int staged_index[2]; - unsigned int invalid_cols[2]; - unsigned int any_invalid = 0; - #pragma unroll - for (int w = 0; w < 2; w++) { - asm volatile("ld.shared.b32 %0, [%1];" : "=r"(*reinterpret_cast(&staged_index[w])) : "r"(index_lane_addr + w * 128)); - unsigned int _vote_0 = __ballot_sync(0xFFFFFFFF, staged_index[w] < 0); - invalid_cols[w] = _vote_0; - any_invalid = any_invalid | invalid_cols[w]; - } - if ((tile_1 & 1) != 0 || tile_1 + 1 == num_kv_tiles_1) { - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(index_empty_addr + (softmax_index_stage) * 8); - softmax_index_stage = softmax_index_stage + 1; - if (softmax_index_stage == 6) { - softmax_index_stage = 0; - } - } - float _tmem_load_0[4]; - tmem_ld_x4(&_tmem_load_0[0], s_base); - float _tmem_load_1[64]; - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]) - : "r"(s_base)); - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]) - : "r"(s_base + 32)); - { - uint32_t _slice_lo_mask_0; - { - int _lim_0 = valid_sparse_cols; - if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } - else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); - } - } - if (!(_slice_lo_mask_0 & (1u << 0))) _tmem_load_1[0] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 1))) _tmem_load_1[1] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 2))) _tmem_load_1[2] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 3))) _tmem_load_1[3] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 4))) _tmem_load_1[4] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 5))) _tmem_load_1[5] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 6))) _tmem_load_1[6] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 7))) _tmem_load_1[7] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 8))) _tmem_load_1[8] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 9))) _tmem_load_1[9] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 10))) _tmem_load_1[10] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 11))) _tmem_load_1[11] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 12))) _tmem_load_1[12] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 13))) _tmem_load_1[13] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 14))) _tmem_load_1[14] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 15))) _tmem_load_1[15] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 16))) _tmem_load_1[16] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 17))) _tmem_load_1[17] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 18))) _tmem_load_1[18] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 19))) _tmem_load_1[19] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 20))) _tmem_load_1[20] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 21))) _tmem_load_1[21] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 22))) _tmem_load_1[22] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 23))) _tmem_load_1[23] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 24))) _tmem_load_1[24] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 25))) _tmem_load_1[25] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 26))) _tmem_load_1[26] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 27))) _tmem_load_1[27] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 28))) _tmem_load_1[28] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 29))) _tmem_load_1[29] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 30))) _tmem_load_1[30] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 31))) _tmem_load_1[31] = -CAKE_INF; - uint32_t _slice_lo_mask_1; - { - int _lim_1 = valid_sparse_cols - 32; - if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } - else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); - } - } - if (!(_slice_lo_mask_1 & (1u << 0))) _tmem_load_1[32] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 1))) _tmem_load_1[33] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 2))) _tmem_load_1[34] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 3))) _tmem_load_1[35] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 4))) _tmem_load_1[36] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 5))) _tmem_load_1[37] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 6))) _tmem_load_1[38] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 7))) _tmem_load_1[39] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 8))) _tmem_load_1[40] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 9))) _tmem_load_1[41] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 10))) _tmem_load_1[42] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 11))) _tmem_load_1[43] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 12))) _tmem_load_1[44] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 13))) _tmem_load_1[45] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 14))) _tmem_load_1[46] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 15))) _tmem_load_1[47] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 16))) _tmem_load_1[48] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 17))) _tmem_load_1[49] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 18))) _tmem_load_1[50] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 19))) _tmem_load_1[51] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 20))) _tmem_load_1[52] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 21))) _tmem_load_1[53] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 22))) _tmem_load_1[54] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 23))) _tmem_load_1[55] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 24))) _tmem_load_1[56] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 25))) _tmem_load_1[57] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 26))) _tmem_load_1[58] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 27))) _tmem_load_1[59] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 28))) _tmem_load_1[60] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 29))) _tmem_load_1[61] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 30))) _tmem_load_1[62] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 31))) _tmem_load_1[63] = -CAKE_INF; - if (any_invalid != 0) { - #pragma unroll - for (int w_1 = 0; w_1 < 2; w_1++) { - unsigned int mask_word = invalid_cols[w_1]; - _tmem_load_1[w_1 * 32] = (((mask_word & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32]); - _tmem_load_1[w_1 * 32 + 1] = (((mask_word >> 1 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 1]); - _tmem_load_1[w_1 * 32 + 2] = (((mask_word >> 2 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 2]); - _tmem_load_1[w_1 * 32 + 3] = (((mask_word >> 3 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 3]); - _tmem_load_1[w_1 * 32 + 4] = (((mask_word >> 4 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 4]); - _tmem_load_1[w_1 * 32 + 5] = (((mask_word >> 5 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 5]); - _tmem_load_1[w_1 * 32 + 6] = (((mask_word >> 6 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 6]); - _tmem_load_1[w_1 * 32 + 7] = (((mask_word >> 7 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 7]); - _tmem_load_1[w_1 * 32 + 8] = (((mask_word >> 8 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 8]); - _tmem_load_1[w_1 * 32 + 9] = (((mask_word >> 9 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 9]); - _tmem_load_1[w_1 * 32 + 10] = (((mask_word >> 10 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 10]); - _tmem_load_1[w_1 * 32 + 11] = (((mask_word >> 11 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 11]); - _tmem_load_1[w_1 * 32 + 12] = (((mask_word >> 12 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 12]); - _tmem_load_1[w_1 * 32 + 13] = (((mask_word >> 13 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 13]); - _tmem_load_1[w_1 * 32 + 14] = (((mask_word >> 14 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 14]); - _tmem_load_1[w_1 * 32 + 15] = (((mask_word >> 15 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 15]); - _tmem_load_1[w_1 * 32 + 16] = (((mask_word >> 16 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 16]); - _tmem_load_1[w_1 * 32 + 17] = (((mask_word >> 17 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 17]); - _tmem_load_1[w_1 * 32 + 18] = (((mask_word >> 18 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 18]); - _tmem_load_1[w_1 * 32 + 19] = (((mask_word >> 19 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 19]); - _tmem_load_1[w_1 * 32 + 20] = (((mask_word >> 20 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 20]); - _tmem_load_1[w_1 * 32 + 21] = (((mask_word >> 21 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 21]); - _tmem_load_1[w_1 * 32 + 22] = (((mask_word >> 22 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 22]); - _tmem_load_1[w_1 * 32 + 23] = (((mask_word >> 23 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 23]); - _tmem_load_1[w_1 * 32 + 24] = (((mask_word >> 24 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 24]); - _tmem_load_1[w_1 * 32 + 25] = (((mask_word >> 25 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 25]); - _tmem_load_1[w_1 * 32 + 26] = (((mask_word >> 26 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 26]); - _tmem_load_1[w_1 * 32 + 27] = (((mask_word >> 27 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 27]); - _tmem_load_1[w_1 * 32 + 28] = (((mask_word >> 28 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 28]); - _tmem_load_1[w_1 * 32 + 29] = (((mask_word >> 29 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 29]); - _tmem_load_1[w_1 * 32 + 30] = (((mask_word >> 30 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 30]); - _tmem_load_1[w_1 * 32 + 31] = (((mask_word >> 31 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 31]); - } - } - float2 _reg_reduce_max2_2 = {-CAKE_INF, -CAKE_INF}; - row_max_x32_accum(&_tmem_load_1[0], _reg_reduce_max2_2); - row_max_x32_accum(&_tmem_load_1[32], _reg_reduce_max2_2); - float _tmem_load_1_max = row_max_reduce(_reg_reduce_max2_2); - float _max_1 = max_noftz(new_max, _tmem_load_1_max); - new_max = _max_1; - smem_softmax_warp_pair_exchange[phase * 128 + stats_row] = new_max; - asm volatile("barrier.sync %0, 64;" :: "r"(2 + warp % 2) : "memory"); - float _max_2 = max_noftz(new_max, smem_softmax_warp_pair_exchange[phase * 128 + (stats_row ^ 64)]); - new_max = _max_2; - } - float _fma_0 = __fmaf_rn(row_max_val, softmax_scale_log2, (-new_max) * softmax_scale_log2); - float delta = _fma_0; - float _exp2_0 = approx_exp2(delta); - float exp_delta = _exp2_0; - float acc_scale = ((row_max_val > -CAKE_INF) ? exp_delta : 1.0f); - smem_stats_max[phase * 128 + stats_row] = acc_scale; - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(stats_addr + (phase) * 8); - row_max_val = new_max; - float safe_max = ((new_max == -CAKE_INF) ? 0.0f : new_max); - float max_scaled = safe_max * softmax_scale_log2; - float block_sum = 0.0f; - { - { - const float2 _fma_b2_3 = {softmax_scale_log2, softmax_scale_log2}; - const float2 _fma_c2_4 = {-max_scaled, -max_scaled}; - #pragma unroll - for (int _lf = 0; _lf < 32; _lf++) - fma_f32x2_inplace(&reinterpret_cast(_tmem_load_1)[_lf], _fma_b2_3, _fma_c2_4); - #pragma unroll - for (int _le = 0; _le < 64; _le++) { - _tmem_load_1[_le] = approx_exp2(_tmem_load_1[_le]); - } - uint32_t _fp8_0[16]; - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[0]), "f"(_tmem_load_1[1]), - "f"(_tmem_load_1[2]), "f"(_tmem_load_1[3])); - _fp8_0[0] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[4]), "f"(_tmem_load_1[5]), - "f"(_tmem_load_1[6]), "f"(_tmem_load_1[7])); - _fp8_0[1] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[8]), "f"(_tmem_load_1[9]), - "f"(_tmem_load_1[10]), "f"(_tmem_load_1[11])); - _fp8_0[2] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[12]), "f"(_tmem_load_1[13]), - "f"(_tmem_load_1[14]), "f"(_tmem_load_1[15])); - _fp8_0[3] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[16]), "f"(_tmem_load_1[17]), - "f"(_tmem_load_1[18]), "f"(_tmem_load_1[19])); - _fp8_0[4] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[20]), "f"(_tmem_load_1[21]), - "f"(_tmem_load_1[22]), "f"(_tmem_load_1[23])); - _fp8_0[5] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[24]), "f"(_tmem_load_1[25]), - "f"(_tmem_load_1[26]), "f"(_tmem_load_1[27])); - _fp8_0[6] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[28]), "f"(_tmem_load_1[29]), - "f"(_tmem_load_1[30]), "f"(_tmem_load_1[31])); - _fp8_0[7] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[32]), "f"(_tmem_load_1[33]), - "f"(_tmem_load_1[34]), "f"(_tmem_load_1[35])); - _fp8_0[8] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[36]), "f"(_tmem_load_1[37]), - "f"(_tmem_load_1[38]), "f"(_tmem_load_1[39])); - _fp8_0[9] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[40]), "f"(_tmem_load_1[41]), - "f"(_tmem_load_1[42]), "f"(_tmem_load_1[43])); - _fp8_0[10] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[44]), "f"(_tmem_load_1[45]), - "f"(_tmem_load_1[46]), "f"(_tmem_load_1[47])); - _fp8_0[11] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[48]), "f"(_tmem_load_1[49]), - "f"(_tmem_load_1[50]), "f"(_tmem_load_1[51])); - _fp8_0[12] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[52]), "f"(_tmem_load_1[53]), - "f"(_tmem_load_1[54]), "f"(_tmem_load_1[55])); - _fp8_0[13] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[56]), "f"(_tmem_load_1[57]), - "f"(_tmem_load_1[58]), "f"(_tmem_load_1[59])); - _fp8_0[14] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[60]), "f"(_tmem_load_1[61]), - "f"(_tmem_load_1[62]), "f"(_tmem_load_1[63])); - _fp8_0[15] = _packed; - } - int p_stage_base = smem_p_fp8_addr + (unsigned int)(phase * 8192); - #pragma unroll - for (int p_vec = 0; p_vec < 4; p_vec++) { - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_stage_base + (my_row * 128 + (n_half * 4 + p_vec) * 16 ^ (my_row * 128 + (n_half * 4 + p_vec) * 16 >> 7 & 7) << 4))), "r"(_fp8_0[p_vec * 4]), "r"(_fp8_0[p_vec * 4 + 1]), "r"(_fp8_0[p_vec * 4 + 2]), "r"(_fp8_0[p_vec * 4 + 3]) : "memory"); - } - } - } - { - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - } - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - { - asm volatile("" ::: "memory"); - } - mbarrier_arrive(p_full_addr + (phase) * 8); - { - { - asm volatile("" ::: "memory"); - const float2* _reg_reduce_src2_5 = reinterpret_cast(&_tmem_load_1[0]); - float2 _reg_reduce_sum2_5_0 = make_float2(0.0f, 0.0f); - float2 _reg_reduce_sum2_5_1 = make_float2(0.0f, 0.0f); - float2 _reg_reduce_sum2_5_2 = make_float2(0.0f, 0.0f); - float2 _reg_reduce_sum2_5_3 = make_float2(0.0f, 0.0f); - #pragma unroll - for (int _rr = 0; _rr < 8; _rr++) { - add_f32x2_inplace(&_reg_reduce_sum2_5_0, _reg_reduce_src2_5[_rr * 4 + 0]); - add_f32x2_inplace(&_reg_reduce_sum2_5_1, _reg_reduce_src2_5[_rr * 4 + 1]); - add_f32x2_inplace(&_reg_reduce_sum2_5_2, _reg_reduce_src2_5[_rr * 4 + 2]); - add_f32x2_inplace(&_reg_reduce_sum2_5_3, _reg_reduce_src2_5[_rr * 4 + 3]); - } - add_f32x2_inplace(&_reg_reduce_sum2_5_0, _reg_reduce_sum2_5_1); - add_f32x2_inplace(&_reg_reduce_sum2_5_2, _reg_reduce_sum2_5_3); - add_f32x2_inplace(&_reg_reduce_sum2_5_0, _reg_reduce_sum2_5_2); - float _tmem_load_1_sum = _reg_reduce_sum2_5_0.x + _reg_reduce_sum2_5_0.y; - block_sum = _tmem_load_1_sum; - float _fma_2 = __fmaf_rn(row_sum_val, acc_scale, block_sum); - row_sum_val = _fma_2; - } - } - } - smem_stats_sum[stats_row] = row_sum_val; - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(sum_ready_addr); - softmax_tile_cursor = softmax_tile_cursor + num_kv_tiles_1; - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(tmem_dealloc_addr); - } - } - // ---- Role: correction_wg ---- - if (warp >= 4 && warp <= 7) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 144;"); - { // correction_wg_main - float softmax_scale_log2_1 = bmm1_scale[0] * 1.4426950408889634f; - float output_scale = bmm2_scale[0]; - const int wg_dummy_inc_1 = 0; - int all_num_kv_tiles_3 = (sparse_topk + 128 - 1) / 128; - const int tmem_row_base_1 = ((1) ? warp % 2 * 32 : warp % 4 * 32); - const int n_half_1 = ((1) ? (int)(warp % 4 / 2) : 0); - const int my_row_1 = tmem_row_base_1 + lane; - const int stats_row_1 = n_half_1 * 64 + my_row_1; - const int corr_row = tmem_row_base_1 << 16; - int correction_tile_cursor = 0; - unsigned int _phase_q_full_0_2 = 0; - unsigned int _phase_pv_done_0 = 0; - unsigned int _phase_o_done_0 = 0; - unsigned int _phase_sum_ready_0 = 0; - #pragma unroll - for (unsigned int work_idx_2 = 0; work_idx_2 < 1; work_idx_2++) { - int split_idx_3 = 0; - int query_idx_3 = bid / 2 >> 1; - int v_chunk_1 = bid / 2 & 1; - { - split_idx_3 = 0; - query_idx_3 = bid / 2; - v_chunk_1 = bid % 2; - } - int tiles_per_split_3 = all_num_kv_tiles_3 + 1 - 1; - int first_tile_3 = split_idx_3 * tiles_per_split_3; - int num_kv_tiles_2 = tiles_per_split_3; - #pragma unroll 1 - for (int tile_2 = 0; tile_2 < num_kv_tiles_2; tile_2++) { - int pipeline_tile_1 = correction_tile_cursor + tile_2; - int phase_1 = pipeline_tile_1 & 1; - int stats_wait_phase = pipeline_tile_1 >> 1 & 1; - mbarrier_wait(stats_addr + (phase_1) * 8, stats_wait_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - float acc_scale_1 = smem_stats_max[phase_1 * 128 + stats_row_1]; - if (tile_2 > 0) { - mbarrier_wait(pv_done_addr, _phase_pv_done_0); - _phase_pv_done_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - } - if (tile_2 > 0) { - int _vote_1 = __any_sync(0xFFFFFFFF, acc_scale_1 < 1.0f); - int any_rescale = _vote_1; - if (any_rescale != 0) { - #pragma unroll - for (int vs = 0; vs < 2; vs++) { - int o_base = taddr + 256 + (unsigned int)(vs * 128) + (unsigned int)corr_row; - #pragma unroll - for (int c = 0; c < 128; c += 16) { - float _tmem_load_4[16]; - tmem_ld_x16(&_tmem_load_4[0], o_base + c); - const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; - #pragma unroll - for (int _ls = 0; _ls < 8; _ls++) - mul_f32x2_inplace(&reinterpret_cast(_tmem_load_4)[_ls], _scale2_0); - tmem_st_x16_f32(o_base + c, _tmem_load_4); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - } - } - } - } - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(corr_done_addr + (phase_1) * 8); - } - mbarrier_wait(o_done_addr, _phase_o_done_0); - _phase_o_done_0 ^= 1; - mbarrier_wait(sum_ready_addr, _phase_sum_ready_0); - _phase_sum_ready_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - float total_sum = smem_stats_sum[stats_row_1]; - float final_max = smem_stats_final_max[stats_row_1]; - { - smem_corr_warp_pair_exchange[stats_row_1] = total_sum; - asm volatile("barrier.sync %0, 64;" :: "r"(4 + warp % 2) : "memory"); - total_sum = total_sum + smem_corr_warp_pair_exchange[stats_row_1 ^ 64]; - } - float _rcp_0 = approx_rcp(total_sum); - float inv_sum = ((total_sum > 0.0f) ? _rcp_0 : 0.0f); - int head_idx = my_row_1; - { - head_idx = bid % 2 * 64 + my_row_1; - } - int direct_o_offset = (query_idx_3 * num_heads + head_idx) * 512 + v_chunk_1 * 256; - int partial_o_offset = (query_idx_3 * num_heads + head_idx + split_idx_3) * 512 + ((0) ? v_chunk_1 : n_half_1) * 256; - int o_offset = partial_o_offset; - #pragma unroll - for (int vs_1 = 0; vs_1 < 2; vs_1++) { - int o_base_epi = taddr + 256 + (unsigned int)(vs_1 * 128) + (unsigned int)corr_row; - #pragma unroll - for (int c_1 = 0; c_1 < 128; c_1 += 64) { - float _tmem_load_5[64]; - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x64.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, [%64];" - : "=f"(_tmem_load_5[0]), "=f"(_tmem_load_5[1]), "=f"(_tmem_load_5[2]), "=f"(_tmem_load_5[3]), "=f"(_tmem_load_5[4]), "=f"(_tmem_load_5[5]), "=f"(_tmem_load_5[6]), "=f"(_tmem_load_5[7]), "=f"(_tmem_load_5[8]), "=f"(_tmem_load_5[9]), "=f"(_tmem_load_5[10]), "=f"(_tmem_load_5[11]), "=f"(_tmem_load_5[12]), "=f"(_tmem_load_5[13]), "=f"(_tmem_load_5[14]), "=f"(_tmem_load_5[15]), "=f"(_tmem_load_5[16]), "=f"(_tmem_load_5[17]), "=f"(_tmem_load_5[18]), "=f"(_tmem_load_5[19]), "=f"(_tmem_load_5[20]), "=f"(_tmem_load_5[21]), "=f"(_tmem_load_5[22]), "=f"(_tmem_load_5[23]), "=f"(_tmem_load_5[24]), "=f"(_tmem_load_5[25]), "=f"(_tmem_load_5[26]), "=f"(_tmem_load_5[27]), "=f"(_tmem_load_5[28]), "=f"(_tmem_load_5[29]), "=f"(_tmem_load_5[30]), "=f"(_tmem_load_5[31]), "=f"(_tmem_load_5[32]), "=f"(_tmem_load_5[33]), "=f"(_tmem_load_5[34]), "=f"(_tmem_load_5[35]), "=f"(_tmem_load_5[36]), "=f"(_tmem_load_5[37]), "=f"(_tmem_load_5[38]), "=f"(_tmem_load_5[39]), "=f"(_tmem_load_5[40]), "=f"(_tmem_load_5[41]), "=f"(_tmem_load_5[42]), "=f"(_tmem_load_5[43]), "=f"(_tmem_load_5[44]), "=f"(_tmem_load_5[45]), "=f"(_tmem_load_5[46]), "=f"(_tmem_load_5[47]), "=f"(_tmem_load_5[48]), "=f"(_tmem_load_5[49]), "=f"(_tmem_load_5[50]), "=f"(_tmem_load_5[51]), "=f"(_tmem_load_5[52]), "=f"(_tmem_load_5[53]), "=f"(_tmem_load_5[54]), "=f"(_tmem_load_5[55]), "=f"(_tmem_load_5[56]), "=f"(_tmem_load_5[57]), "=f"(_tmem_load_5[58]), "=f"(_tmem_load_5[59]), "=f"(_tmem_load_5[60]), "=f"(_tmem_load_5[61]), "=f"(_tmem_load_5[62]), "=f"(_tmem_load_5[63]) - : "r"(o_base_epi + c_1)); - int gmem_base = o_offset + vs_1 * 128 + c_1; - if (bid % 2 == 0) { - #pragma unroll - for (int j = 0; j < 64; j += 16) { - { - const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; - #if __CUDA_ARCH__ >= 1000 - #pragma unroll - for (int _ps = 0; _ps < 8; _ps++) - mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_5[j])[_ps], _prescale2_1); - #else - #pragma unroll - for (int _ps = 0; _ps < 16; _ps++) - _tmem_load_5[j + _ps] *= inv_sum * output_scale; - #endif - { - __nv_bfloat162 _pk0 = __floats2bfloat162_rn(_tmem_load_5[j + 0], _tmem_load_5[j + 1]); - unsigned _pk_u0 = *reinterpret_cast(&_pk0); - __nv_bfloat162 _pk1 = __floats2bfloat162_rn(_tmem_load_5[j + 2], _tmem_load_5[j + 3]); - unsigned _pk_u1 = *reinterpret_cast(&_pk1); - __nv_bfloat162 _pk2 = __floats2bfloat162_rn(_tmem_load_5[j + 4], _tmem_load_5[j + 5]); - unsigned _pk_u2 = *reinterpret_cast(&_pk2); - __nv_bfloat162 _pk3 = __floats2bfloat162_rn(_tmem_load_5[j + 6], _tmem_load_5[j + 7]); - unsigned _pk_u3 = *reinterpret_cast(&_pk3); - __nv_bfloat162 _pk4 = __floats2bfloat162_rn(_tmem_load_5[j + 8], _tmem_load_5[j + 9]); - unsigned _pk_u4 = *reinterpret_cast(&_pk4); - __nv_bfloat162 _pk5 = __floats2bfloat162_rn(_tmem_load_5[j + 10], _tmem_load_5[j + 11]); - unsigned _pk_u5 = *reinterpret_cast(&_pk5); - __nv_bfloat162 _pk6 = __floats2bfloat162_rn(_tmem_load_5[j + 12], _tmem_load_5[j + 13]); - unsigned _pk_u6 = *reinterpret_cast(&_pk6); - __nv_bfloat162 _pk7 = __floats2bfloat162_rn(_tmem_load_5[j + 14], _tmem_load_5[j + 15]); - unsigned _pk_u7 = *reinterpret_cast(&_pk7); - asm volatile( - "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" - :: "l"((void*)(&((__nv_bfloat16*)(O + (gmem_base + j)))[0])), "r"(_pk_u0), "r"(_pk_u1), "r"(_pk_u2), "r"(_pk_u3), "r"(_pk_u4), "r"(_pk_u5), "r"(_pk_u6), "r"(_pk_u7) : "memory"); - } - } - } - } - } - } - correction_tile_cursor = correction_tile_cursor + num_kv_tiles_2; - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(tmem_dealloc_addr); - } - } - // ---- Role: mma_warp ---- - if (warp == 8) { - { // mma_warp_main - const int wg2_dummy_1 = 0; - int all_num_kv_tiles_4 = (sparse_topk + 128 - 1) / 128; - unsigned int mma_k_stage = 0; - unsigned int mma_v_stage = 0; - unsigned int mma_kv_stage = 0; - int mma_tile_cursor = 0; - unsigned int _phase_s_seeded_0 = 0; - unsigned int _phase_q_full_0_3 = 0; - unsigned int _phase_q_pair_ready_0 = 0; - unsigned int _phase_k_full = 0; - unsigned int _phase_k_full_1 = 0; - unsigned int _phase_kv_pair_ready = 0; - unsigned int _phase_v_full = 0; - #pragma unroll - for (unsigned int work_idx_3 = 0; work_idx_3 < 1; work_idx_3++) { - int split_idx_4 = 0; - { - split_idx_4 = 0; - } - int tiles_per_split_4 = all_num_kv_tiles_4 + 1 - 1; - int first_tile_4 = split_idx_4 * tiles_per_split_4; - int num_kv_tiles_3 = tiles_per_split_4; - { - if (cta_rank == 0) { - mbarrier_wait(q_full_addr, _phase_q_full_0_3); - _phase_q_full_0_3 ^= 1; - } - } - int first_pv = 1; - #pragma unroll 1 - for (int tile_3 = 0; tile_3 < num_kv_tiles_3; tile_3++) { - int pipeline_tile_2 = mma_tile_cursor + tile_3; - int phase_2 = pipeline_tile_2 & 1; - int score_col = ((phase_2 != 0) ? 128 : 0); - { - if (cta_rank == 0) { - mbarrier_wait(k_full_addr + (mma_k_stage) * 8, _phase_k_full); - int _mma_a_lo_0 = (((smem_q_full_addr) >> 4) & 0x3FFF) + (0) * 2048; - int _mma_b_lo_0 = (((smem_k_full_addr) >> 4) & 0x3FFF) + (mma_k_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 136314896;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (score_col))), "r"(0)); - int _mma_a_lo_1 = (((smem_q_full_addr + 8192) >> 4) & 0x3FFF) + (0) * 2048; - int _mma_b_lo_1 = (((smem_k_full_addr + 8192) >> 4) & 0x3FFF) + (mma_k_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 136314896;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); - int _mma_a_lo_2 = (((smem_q_full_addr + 16384) >> 4) & 0x3FFF) + (0) * 2048; - int _mma_b_lo_2 = (((smem_k_full_addr + 16384) >> 4) & 0x3FFF) + (mma_k_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 136314896;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); - int _mma_a_lo_3 = (((smem_q_full_addr + 24576) >> 4) & 0x3FFF) + (0) * 2048; - int _mma_b_lo_3 = (((smem_k_full_addr + 24576) >> 4) & 0x3FFF) + (mma_k_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 136314896;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); - if (cta_rank == 0) { - elect_commit_cg2_multicast(s_full_addr + (phase_2) * 8, (uint16_t)(3)); - } - if (cta_rank == 0) { - elect_commit_cg2_multicast(k_empty_addr + (mma_k_stage) * 8, (uint16_t)(3)); - } - } - mma_k_stage += 1; - if (mma_k_stage == 3) { mma_k_stage = 0; _phase_k_full ^= 1; } - } - if (tile_3 > 0) { - int prev_pipeline_tile = pipeline_tile_2 - 1; - int prev_phase = prev_pipeline_tile & 1; - int pv_wait_phase = prev_pipeline_tile >> 1 & 1; - mbarrier_wait(p_full_addr + (prev_phase) * 8, pv_wait_phase); - mbarrier_wait(corr_done_addr + (prev_phase) * 8, pv_wait_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - { - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cluster.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(pv_pair_ready_addr + (unsigned int)(prev_phase * 8)), "r"(0) : "memory"); - if (cta_rank == 0) { - mbarrier_wait_cluster_hint(pv_pair_ready_addr + (prev_phase) * 8, pv_wait_phase, 10000000); - } - } - int p_col = ((prev_phase != 0) ? 224 : 96); - { - if (cta_rank == 0) { - mbarrier_wait(v_full_addr + (mma_v_stage) * 8, _phase_v_full); - int _mma_a_lo_6 = (((smem_p_fp8_addr) >> 4) & 0x3FFF) + (prev_phase) * 512; - int _mma_b_lo_6 = ((((smem_v_full_addr) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 138477584;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_6), "r"(_mma_b_lo_6), "r"((tmem_tmem_scratch + (256))), "r"(((first_pv) ? 0 : 1))); - int _mma_a_lo_7 = (((smem_p_fp8_addr) >> 4) & 0x3FFF) + (prev_phase) * 512; - int _mma_b_lo_7 = ((((smem_v_full_addr + 16384) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 138477584;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_7), "r"(_mma_b_lo_7), "r"((tmem_tmem_scratch + (384))), "r"(((first_pv) ? 0 : 1))); - if (cta_rank == 0) { - elect_commit_cg2_multicast(v_empty_addr + (mma_v_stage) * 8, (uint16_t)(3)); - } - } - mma_v_stage += 1; - if (mma_v_stage == 2) { mma_v_stage = 0; _phase_v_full ^= 1; } - } - first_pv = 0; - if (cta_rank == 0) { - elect_commit_cg2_multicast(pv_done_addr, (uint16_t)(3)); - } - } - } - int last_pipeline_tile = mma_tile_cursor + num_kv_tiles_3 - 1; - int last_phase = last_pipeline_tile & 1; - int drain_wait_phase = last_pipeline_tile >> 1 & 1; - mbarrier_wait(p_full_addr + (last_phase) * 8, drain_wait_phase); - mbarrier_wait(corr_done_addr + (last_phase) * 8, drain_wait_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - { - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cluster.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(pv_pair_ready_addr + (unsigned int)(last_phase * 8)), "r"(0) : "memory"); - if (cta_rank == 0) { - mbarrier_wait_cluster_hint(pv_pair_ready_addr + (last_phase) * 8, drain_wait_phase, 10000000); - } - } - int p_col_last = ((last_phase != 0) ? 224 : 96); - { - if (cta_rank == 0) { - mbarrier_wait(v_full_addr + (mma_v_stage) * 8, _phase_v_full); - int _mma_a_lo_10 = (((smem_p_fp8_addr) >> 4) & 0x3FFF) + (last_phase) * 512; - int _mma_b_lo_10 = ((((smem_v_full_addr) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 138477584;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_10), "r"(_mma_b_lo_10), "r"((tmem_tmem_scratch + (256))), "r"(((first_pv) ? 0 : 1))); - int _mma_a_lo_11 = (((smem_p_fp8_addr) >> 4) & 0x3FFF) + (last_phase) * 512; - int _mma_b_lo_11 = ((((smem_v_full_addr + 16384) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 138477584;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_11), "r"(_mma_b_lo_11), "r"((tmem_tmem_scratch + (384))), "r"(((first_pv) ? 0 : 1))); - if (cta_rank == 0) { - elect_commit_cg2_multicast(v_empty_addr + (mma_v_stage) * 8, (uint16_t)(3)); - } - } - mma_v_stage += 1; - if (mma_v_stage == 2) { mma_v_stage = 0; _phase_v_full ^= 1; } - } - if (cta_rank == 0) { - elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); - } - if (cta_rank == 0) { - elect_commit_cg2_multicast(o_done_addr, (uint16_t)(3)); - } - mma_tile_cursor = mma_tile_cursor + num_kv_tiles_3; - } - mbarrier_arrive(tmem_dealloc_addr); - unsigned int _phase_tmem_dealloc_0 = 0; - mbarrier_wait(tmem_dealloc_addr, _phase_tmem_dealloc_0); - _phase_tmem_dealloc_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - { - int tmem_dealloc_peer_rank = bid % 2 ^ 1; - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cluster.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(tmem_dealloc_peer_addr), "r"(tmem_dealloc_peer_rank) : "memory"); - mbarrier_wait(tmem_dealloc_peer_addr, 0); - int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); - asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); - } - } - } - // ---- Role: empty1 ---- - if (warp >= 10 && warp <= 11) { - { // empty1_main - const int wg2_dummy_2 = 0; - unsigned int _phase_q_full_0_4 = 0; - } - } - - // Cleanup -} - -} // extern "C" diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_9526f80c89f2b5ef4c6c_binding.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_9526f80c89f2b5ef4c6c_binding.cu deleted file mode 100644 index 931e515eb10..00000000000 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_9526f80c89f2b5ef4c6c_binding.cu +++ /dev/null @@ -1,436 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -// Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_9526f80c89f2b5ef4c6c'. -#include -#include -#include -#include - -#include "tvm_ffi_utils.h" - -#include -#include -#include -#include -#include -#include -#include - -extern "C" __global__ void kernel_cake_dsv4_9526f80c89f2b5ef4c6c(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int num_query_tokens, int has_sinks); - - -namespace cake_host_shim_638caaf740c22a7d { - -using tvm::ffi::Optional; -using tvm::ffi::TensorView; - -class ScopedCudaDevice { - public: - explicit ScopedCudaDevice(int device_id) { - cudaError_t error = cudaGetDevice(&previous_device_); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed before host-shim launch: cudaError=" - << static_cast(error); - if (previous_device_ != device_id) { - error = cudaSetDevice(device_id); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaSetDevice failed before host-shim launch for cuda:" - << device_id << ": cudaError=" << static_cast(error); - restore_ = true; - } - } - - ScopedCudaDevice(const ScopedCudaDevice&) = delete; - ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; - - ~ScopedCudaDevice() noexcept { - if (restore_) { - (void)cudaSetDevice(previous_device_); - } - } - - private: - int previous_device_ = -1; - bool restore_ = false; -}; - -inline int64_t CakeDeviceMultiprocessorCount(int device_id) { - constexpr int kMaxCachedCudaDevices = 64; - TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) - << "physical-SM-count cache does not cover cuda:" << device_id; - static std::atomic count_by_device[kMaxCachedCudaDevices]{}; - int cached = count_by_device[device_id].load(std::memory_order_acquire); - if (cached > 0) return cached; - - int count = 0; - cudaError_t error = cudaDeviceGetAttribute( - &count, cudaDevAttrMultiProcessorCount, device_id); - TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) - << "querying multiProcessorCount failed for cuda:" << device_id - << ": cudaError=" << static_cast(error) << ", count=" << count; - // Concurrent first launches may repeat the immutable device query, but all - // publication is atomic and every later hot-path lookup is one acquire load. - count_by_device[device_id].store(count, std::memory_order_release); - return count; -} - -inline void CheckCudaTensor(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) - << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; -} - -inline void CheckSameCudaDevice( - const TensorView& t, - const TensorView& reference, - const char* name, - const char* reference_name) { - TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) - << name << " must be on the same CUDA device as " << reference_name - << ": got cuda:" << t.device().device_id - << " versus cuda:" << reference.device().device_id; -} - -inline void CheckCurrentCudaDevice( - const TensorView& reference, - const char* reference_name) { - int current_device = -1; - cudaError_t error = cudaGetDevice(¤t_device); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed while validating " << reference_name - << ": cudaError=" << static_cast(error); - TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) - << "current CUDA device must match " << reference_name - << ": current=cuda:" << current_device - << ", tensor=cuda:" << reference.device().device_id; -} - -inline void CheckContiguous(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; -} - -inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { - DLDataType d = t.dtype(); - TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) - << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits - << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits - << ", lanes=" << (int)d.lanes << ")"; -} - -// A logical axis.outer(trailing) folds every source dim above the trailing -// dimensions. Shape products are independent of physical strides, so verify -// the leading dimensions form one dense row-major chain instead of inventing -// a "folded stride". The descriptor reads its exact adjacent physical step -// separately through stride[-(trailing + 1)]. -inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { - TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) - << name << " cannot fold leading dimensions above " << trailing - << " trailing dims from ndim=" << t.ndim(); - int outer_last = t.ndim() - trailing - 1; - if (outer_last <= 0) { - return; - } - int64_t step = t.stride(outer_last); - TVM_FFI_CHECK(step > 0, ValueError) - << name << " physical strides must be positive"; - int64_t expected = step; - for (int axis = outer_last - 1; axis >= 0; --axis) { - expected *= t.size(axis + 1); - if (t.size(axis) > 1) { - TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) - << name << " leading dims are not physically foldable above " << trailing - << " trailing dims: stride(" << axis << ")=" << t.stride(axis) - << ", expected " << expected; - } - } -} - -// 4D TMA descriptor for buffer 'tmap_q' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 3, ValueError) - << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - int64_t d2 = t.size(t.ndim() - 2); - int64_t d3 = t.size(t.ndim() - 3); - TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) - << "TMA source 'tmap_q' trailing dims must be positive"; - TVM_FFI_CHECK(d1 % 64 == 0, ValueError) - << "TMA source 'tmap_q' extent " << d1 - << " must divide exactly by " << 64; - uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(d3)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) - << "TMA box (64, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; - int64_t carrier_stride_0 = d1; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; - int64_t carrier_stride_1 = 64; - TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; - TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; - TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; - int64_t carrier_stride_2 = (d2 * d1); - TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; - TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; - TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; - uint64_t global_strides[3] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - (uint64_t)((carrier_stride_1 * 16) / 8), - (uint64_t)((carrier_stride_2 * 16) / 8), - }; - uint32_t box_dim[4] = {64u, 64u, 1u, 1u}; - uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_swa_kv' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_swa_kv' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_kv'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_compressed_kv' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_compressed_kv(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_compressed_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_compressed_kv' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_compressed_kv' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_compressed_kv"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_compressed_kv' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_kv'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_kv') failed: CUresult=" << (int)r; - return tm; -} - -void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_compressed_kv, TensorView arg_O, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_num_query_tokens, int64_t arg_has_sinks, int64_t grid_x, int64_t grid_y, int64_t grid_z) { - DLDevice dev = arg_tmap_q.device(); - ScopedCudaDevice device_guard(dev.device_id); - CheckCudaTensor(arg_tmap_q, "tmap_q"); - CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); - CheckContiguous(arg_tmap_q, "tmap_q"); - CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); - CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 4, 16, 1); - CheckCudaTensor(arg_tmap_compressed_kv, "tmap_compressed_kv"); - CheckDtype(arg_tmap_compressed_kv, "tmap_compressed_kv", 4, 16, 1); - CheckCudaTensor(arg_O, "O"); - CheckDtype(arg_O, "O", 4, 16, 1); - CheckContiguous(arg_O, "O"); - CheckCudaTensor(arg_swa_indices, "swa_indices"); - CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); - CheckContiguous(arg_swa_indices, "swa_indices"); - CheckCudaTensor(arg_compressed_indices, "compressed_indices"); - CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); - CheckContiguous(arg_compressed_indices, "compressed_indices"); - CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); - CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckCudaTensor(arg_sinks, "sinks"); - CheckDtype(arg_sinks, "sinks", 2, 32, 1); - CheckContiguous(arg_sinks, "sinks"); - CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); - CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); - CheckContiguous(arg_bmm1_scale, "bmm1_scale"); - CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); - CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); - CheckContiguous(arg_bmm2_scale, "bmm2_scale"); - TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) - << "scalar 'num_heads' value " << arg_num_heads - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) - << "scalar 'swa_index_stride' value " << arg_swa_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) - << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) - << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) - << "scalar 'sparse_topk' value " << arg_sparse_topk - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_num_query_tokens >= -2147483648LL && arg_num_query_tokens <= 2147483647LL, ValueError) - << "scalar 'num_query_tokens' value " << arg_num_query_tokens - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) - << "scalar 'has_sinks' value " << arg_has_sinks - << " is outside i32 range [-2147483648, 2147483647]"; - CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); - CheckSameCudaDevice(arg_tmap_compressed_kv, arg_tmap_q, "tmap_compressed_kv", "tmap_q"); - CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); - CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); - CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); - CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); - CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); - CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); - CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); - TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) - << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y - << ", " << grid_z << ")"; - - cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); - CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); - CUtensorMap p_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); - CUtensorMap p_tmap_compressed_kv = EncodeTma_tmap_compressed_kv(arg_tmap_compressed_kv); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); - int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); - int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); - int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); - float* p_sinks = static_cast(arg_sinks.data_ptr()); - float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); - float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); - int32_t v_num_heads = (int32_t)arg_num_heads; - int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; - int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; - int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; - int32_t v_sparse_topk = (int32_t)arg_sparse_topk; - int32_t v_num_query_tokens = (int32_t)arg_num_query_tokens; - int32_t v_has_sinks = (int32_t)arg_has_sinks; - void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_tmap_compressed_kv, &p_O, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_num_query_tokens, &v_has_sinks}; - - static std::mutex smem_mu; - static auto* smem_status_by_device = - new std::unordered_map(); - cudaError_t smem_status = cudaSuccess; - { - std::lock_guard lock(smem_mu); - auto it = smem_status_by_device->find(dev.device_id); - if (it == smem_status_by_device->end()) { - int current_device = -1; - smem_status = cudaGetDevice(¤t_device); - if (smem_status == cudaSuccess && current_device != dev.device_id) { - smem_status = cudaErrorInvalidDevice; - } - if (smem_status == cudaSuccess) { - smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_9526f80c89f2b5ef4c6c), - cudaFuncAttributeMaxDynamicSharedMemorySize, - 203904); - } - smem_status_by_device->emplace(dev.device_id, smem_status); - } else { - smem_status = it->second; - } - } - TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_9526f80c89f2b5ef4c6c failed: " - << cudaGetErrorString(smem_status); - dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); - dim3 block(512u, 1u, 1u); - - cudaLaunchConfig_t config{}; - config.gridDim = grid; - config.blockDim = block; - config.dynamicSmemBytes = 203904u; - config.stream = stream; - cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_9526f80c89f2b5ef4c6c), kargs); - TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_9526f80c89f2b5ef4c6c failed: " - << cudaGetErrorString(launch_status); -} - -} // namespace cake_host_shim_638caaf740c22a7d - -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_638caaf740c22a7d::Run); diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_9526f80c89f2b5ef4c6c_kernel.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_9526f80c89f2b5ef4c6c_kernel.cu deleted file mode 100644 index b17e9aa149f..00000000000 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_9526f80c89f2b5ef4c6c_kernel.cu +++ /dev/null @@ -1,2542 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -typedef signed char int8_t; -typedef unsigned char uint8_t; -typedef unsigned short uint16_t; -typedef unsigned int uint32_t; -#if defined(__CUDACC_RTC__) -typedef unsigned long long uint64_t; -#else -typedef unsigned long uint64_t; -#endif -static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); -typedef signed int int32_t; -typedef short int int16_t; -struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; -struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; -static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); -static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); - -#if defined(__CUDACC_RTC__) -typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; -#else -#include -#endif - -static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); -static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); -#include -#include - -#define CAKE_INF CUDART_INF_F -#define TMEM_NCOLS 512 -#define TMEM_TMEM_OFFSET 0 -#define NUM_KV_PIPE_STAGES 4 -#define NUM_PAGE_PIPE_STAGES 6 -#define NUM_WORK_PIPE_STAGES 2 -#define NUM_THROTTLE_PIPE_STAGES 2 -#define SMEM_SMEM_Q_OFF 1024 -#define SMEM_SMEM_Q_STAGE_BYTES 16384 -#define SMEM_SMEM_Q_STRIDE 16384 -#define SMEM_SMEM_KV_OFF 66560 -#define SMEM_SMEM_KV_STAGE_BYTES 32768 -#define SMEM_SMEM_KV_STRIDE 32768 -#define SMEM_SMEM_V_OFF 66560 -#define SMEM_SMEM_V_STAGE_BYTES 32768 -#define SMEM_SMEM_V_STRIDE 32768 -#define SMEM_SMEM_PAGE_OFFSETS_OFF 197632 -#define SMEM_SMEM_PAGE_OFFSETS_STAGE_BYTES 1024 -#define SMEM_SMEM_PAGE_OFFSETS_STRIDE 1024 -#define SMEM_WORK_RESPONSE_OFF 203776 -#define SMEM_WORK_RESPONSE_STAGE_BYTES 16 -#define SMEM_WORK_RESPONSE_STRIDE 16 -#define SMEM_TOTAL 203904 -#define THREADS 512 - -#include - -__device__ __forceinline__ uint32_t elect_sync() { - uint32_t pred = 0; - asm volatile( - "{\n\t" - ".reg .pred %%px;\n\t" - "elect.sync _|%%px, %1;\n\t" - "@%%px mov.s32 %0, 1;\n\t" - "}\n" - : "+r"(pred) - : "r"(0xFFFFFFFF)); - return pred; -} - - -__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { - asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" - :: "r"(mbar_addr), "r"(count) : "memory"); -} - -__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { - asm volatile("mbarrier.init.b64 [%0], %1;" - :: "l"(mbar_addr), "r"(count) : "memory"); -} - - -__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - - -// CTA-local pipelines have short, resident producer/consumer edges. Omitting -// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the -// explicit loop still makes this helper blocking until acquire succeeds. -__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE;\n\t" - "bra.uni LAB_WAIT;\n\t" - "DONE:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Source-faithful relaxed CTA wait used only by a typed protocol that does -// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. -__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, 10000000;\n\t" - "@P1 bra.uni DONE_RELAXED;\n\t" - "bra.uni LAB_WAIT_RELAXED;\n\t" - "DONE_RELAXED:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Exact source ports may request the PTX suspendTimeHint operand explicitly. -// The hint is expressed in nanoseconds and is kept separate from the canonical -// no-hint CTA helper so unrelated schedules retain their existing retry path. -__device__ __forceinline__ void mbarrier_wait_suspend( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_SUSPEND:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_SUSPEND;\n\t" - "bra.uni LAB_WAIT_SUSPEND;\n\t" - "DONE_SUSPEND:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE_CLUSTER;\n\t" - "bra.uni LAB_WAIT_CLUSTER;\n\t" - "DONE_CLUSTER:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - ".reg .u32 WAIT_ADDR;\n\t" - "mov.u32 WAIT_ADDR, %0;\n\t" - "LAB_WAIT_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [WAIT_ADDR], %1, %2;\n\t" - "@P1 bra.uni DONE_HINT;\n\t" - "bra.uni LAB_WAIT_HINT;\n\t" - "DONE_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -// Exact unqualified CTA wait used by source schedules whose PTX intentionally -// omits the acquire qualifier while retaining a typed suspendTimeHint operand. -__device__ __forceinline__ void mbarrier_wait_relaxed_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED_HINT:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra DONE_RELAXED_HINT;\n\t" - "bra LAB_WAIT_RELAXED_HINT;\n\t" - "DONE_RELAXED_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); -} - -__device__ __forceinline__ void mbarrier_wait_cluster_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" - "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" - "DONE_CLUSTER_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_suspend( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait_cluster(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); - } -} - - -__device__ __forceinline__ void tcgen05_mma_f16( - int taddr, uint64_t a_desc, uint64_t b_desc, - uint32_t i_desc, int enable_input_d) { - asm volatile( - "{\n\t" - ".reg .pred p;\n\t" - "setp.ne.b32 p, %4, 0;\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" - "}\n" - :: "r"(taddr), "l"(a_desc), "l"(b_desc), - "r"(i_desc), "r"(enable_input_d) - : "memory"); -} - - -__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { - return (x & 0x3FFFFULL) >> 4ULL; -} - - -union MmaSmemDesc { - uint64_t u64; - uint32_t u32[2]; -}; - -__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { - MmaSmemDesc tmp; - tmp.u64 = smem_desc; - tmp.u32[0] += offset; - smem_desc = tmp.u64; -} - - -__device__ __forceinline__ void elect_commit(int mbar_addr) { - asm volatile( - "{\n\t" - ".reg .pred leader;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" - ".shared::cluster.b64 [%0];\n\t" - "}\n" - :: "r"(mbar_addr)); -} - - -__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { - asm volatile( - "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" - :: "r"(mbar_addr), "r"(bytes) : "memory"); -} - - -__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7," - " %8, %9, %10, %11, %12, %13, %14, %15," - " %16, %17, %18, %19, %20, %21, %22, %23," - " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), - "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), - "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), - "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), - "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), - "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), - "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x16.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7," - " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), - "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), - "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ float approx_exp2(float x) { - float y; - asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float approx_rcp(float x) { - float y; - asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float max_noftz(float a, float b) { - float c; - asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); - return c; -} - - -__device__ __forceinline__ float warp_reduce_max(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); - return val; -} - - -__device__ __forceinline__ float warp_reduce_sum(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val += __shfl_xor_sync(0xFFFFFFFF, val, offset); - return val; -} - - -__device__ __forceinline__ float row_max_reduce(float2 acc) { - return max_noftz(acc.x, acc.y); -} - - -__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { - #pragma unroll - for (int j = 0; j < 16; j++) { - if (j % 2 == 0) - acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); - else - acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); - } -} - - -__device__ __forceinline__ float2 ex2_emulation_f32x2_value(float2 value) { - const float c0 = 1.0f, c1 = 0.695146143436431884765625f; - const float c2 = 0.227564394474029541015625f, c3 = 0.077119089663028717041015625f; - const float magic = 12582912.0f; - float x0 = max_noftz(value.x, -127.0f), x1 = max_noftz(value.y, -127.0f); - float2 xc2 = make_float2(x0, x1), magic2 = make_float2(magic, magic); - float2 xr2; - asm("add.rm.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xr2) - : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&magic2)); - float2 c3_2 = make_float2(c3, c3), c2_2 = make_float2(c2, c2); - float2 c1_2 = make_float2(c1, c1), c0_2 = make_float2(c0, c0); - float2 xrb2, xfrac2; - asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xrb2) - : "l"(*(unsigned long long*)&xr2), "l"(*(unsigned long long*)&magic2)); - asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xfrac2) - : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&xrb2)); - float2 poly2; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&c3_2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c2_2)); - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c1_2)); - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c0_2)); - int x0r_i, x1r_i, p0_i, p1_i; - asm("mov.b64 {%0, %1}, %2;" : "=r"(x0r_i), "=r"(x1r_i) : "l"(*(unsigned long long*)&xr2)); - asm("mov.b64 {%0, %1}, %2;" : "=r"(p0_i), "=r"(p1_i) : "l"(*(unsigned long long*)&poly2)); - float r0, r1; - asm("mov.b32 %0, %1;" : "=f"(r0) : "r"((x0r_i << 23) + p0_i)); - asm("mov.b32 %0, %1;" : "=f"(r1) : "r"((x1r_i << 23) + p1_i)); - return make_float2(r0, r1); -} - -__device__ __forceinline__ void ex2_emulation_f32x2(float* x0_ptr, float* x1_ptr) { - float2 result = ex2_emulation_f32x2_value(make_float2(*x0_ptr, *x1_ptr)); - *x0_ptr = result.x; *x1_ptr = result.y; -} - -__device__ __forceinline__ void softmax_frag_exp2_cast( - float* sv, uint32_t* pv, int use_emu) -{ - #pragma unroll - for (int j = 0; j < 16; j++) { - if (use_emu && j >= 12) - ex2_emulation_f32x2(&sv[j*2], &sv[j*2+1]); - else { - sv[j*2] = approx_exp2(sv[j*2]); - sv[j*2+1] = approx_exp2(sv[j*2+1]); - } - } - #pragma unroll - for (int j = 0; j < 16; j++) { - __nv_bfloat162 bf = __float22bfloat162_rn({sv[j*2], sv[j*2+1]}); - pv[j] = reinterpret_cast(bf); - } -} - - - -__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { - const float2* sv2 = reinterpret_cast(sv); - #pragma unroll - for (int j = 0; j < 16; j++) { - asm("add.f32x2 %0, %1, %2;" - : "+l"(reinterpret_cast(*acc)) - : "l"(reinterpret_cast(*acc)), - "l"(reinterpret_cast(sv2[j]))); - } -} - - -__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { - asm("mul.rn.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { - asm("mul.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { - asm("add.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { - asm("add.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { - asm("sub.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { - asm("sub.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("add.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { - float2 r; - asm("sub.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("sub.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ void fma_scale_x32( - float* sv, const float2* scale2, const float2* neg_max2) -{ - float2* sv_2 = reinterpret_cast(sv); - #pragma unroll - for (int j = 0; j < 16; j++) - fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); -} - -__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("mul.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - - -// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) - -__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("add.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("add.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("add.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("add.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("add.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("add.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("add.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - - -__device__ __forceinline__ void fence_async_shared() { - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); -} - - -__device__ __forceinline__ uint64_t make_smem_desc(int addr) { - const int SBO = 1024; - return desc_encode(addr) - | (desc_encode(SBO) << 32ULL) - | (1ULL << 46ULL) - | (2ULL << 61ULL); -} - - -__device__ __forceinline__ void tma_2d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3}], [%4];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_4d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.4d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5}], [%6];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr) { - // Canonical .shared::cta form for non-multicast gather4, matching - // trtllm-gen / cuda_ptx and the PTX ISA qualifier order - // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, - // .shared::cluster is reserved for the multicast variant (ctaMask). - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { - asm volatile( - "tcgen05.commit.cta_group::1.mbarrier::arrive::one" - ".shared::cluster.b64 [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { - uint32_t result; - asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" - : "=r"(result) : "r"(val)); - return result; -} - -extern "C" { - -__global__ __launch_bounds__(512, 1) void -kernel_cake_dsv4_9526f80c89f2b5ef4c6c(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int num_query_tokens, int has_sinks) -{ - const int tid = threadIdx.x; - const int warp = make_warp_uniform(tid / 32); - const int lane = tid % 32; - - extern __shared__ __align__(1024) char smem_raw[]; - int smem; - asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); - smem = make_warp_uniform(smem); - - const int mbar_base = smem; - #define q_full_addr (mbar_base + 0) - #define q_empty_addr (mbar_base + 8) - #define kv_full_addr (mbar_base + 16) - #define kv_empty_addr (mbar_base + 48) - #define s_full_addr (mbar_base + 80) - #define p_full_addr (mbar_base + 96) - #define stats_addr (mbar_base + 112) - #define stats_empty_addr (mbar_base + 128) - #define o_empty_addr (mbar_base + 144) - #define o_full_addr (mbar_base + 152) - #define page_full_addr (mbar_base + 160) - #define page_empty_addr (mbar_base + 208) - #define work_full_addr (mbar_base + 256) - #define work_empty_addr (mbar_base + 272) - #define throttle_full_addr (mbar_base + 288) - #define throttle_empty_addr (mbar_base + 304) - - const int bid = blockIdx.x; - const int num_bids = gridDim.x; - - const int cta_rank = 0; - - // Kernel setup ops - __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); - const int smem_q_addr = smem + 1024; - __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); - const int smem_kv_addr = smem + 66560; - __nv_bfloat16* smem_v = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); - const int smem_v_addr = smem + 66560; - int* smem_page_offsets = reinterpret_cast(smem_raw + 197632); - const int smem_page_offsets_addr = smem + 197632; - unsigned int* work_response = reinterpret_cast(smem_raw + 203776); - const int work_response_addr = smem + 203776; - - // Mbarrier init (16 pipeline groups, 0 ordered-sequence groups, 40 barriers) - // Mbarriers at smem_raw[0..320) - - if (warp == 0) { - uint32_t leader = elect_sync(); - if (leader) { - // q_full: 1 barriers, init_count=1 - mbarrier_init(smem + 0, 1); - // q_empty: 1 barriers, init_count=1 - mbarrier_init(smem + 8, 1); - // --- pipeline 'kv_pipe' --- - // kv_full: 4 barriers, init_count=1 - mbarrier_init(smem + 16, 1); - mbarrier_init(smem + 24, 1); - mbarrier_init(smem + 32, 1); - mbarrier_init(smem + 40, 1); - // kv_empty: 4 barriers, init_count=1 - mbarrier_init(smem + 48, 1); - mbarrier_init(smem + 56, 1); - mbarrier_init(smem + 64, 1); - mbarrier_init(smem + 72, 1); - // s_full: 2 barriers, init_count=1 - mbarrier_init(smem + 80, 1); - mbarrier_init(smem + 88, 1); - // p_full: 2 barriers, init_count=128 - mbarrier_init(smem + 96, 128); - mbarrier_init(smem + 104, 128); - // stats: 2 barriers, init_count=128 - mbarrier_init(smem + 112, 128); - mbarrier_init(smem + 120, 128); - // stats_empty: 2 barriers, init_count=128 - mbarrier_init(smem + 128, 128); - mbarrier_init(smem + 136, 128); - // o_empty: 1 barriers, init_count=128 - mbarrier_init(smem + 144, 128); - // o_full: 1 barriers, init_count=1 - mbarrier_init(smem + 152, 1); - // --- pipeline 'page_pipe' --- - // page_full: 6 barriers, init_count=32 - mbarrier_init(smem + 160, 32); - mbarrier_init(smem + 168, 32); - mbarrier_init(smem + 176, 32); - mbarrier_init(smem + 184, 32); - mbarrier_init(smem + 192, 32); - mbarrier_init(smem + 200, 32); - // page_empty: 6 barriers, init_count=128 - mbarrier_init(smem + 208, 128); - mbarrier_init(smem + 216, 128); - mbarrier_init(smem + 224, 128); - mbarrier_init(smem + 232, 128); - mbarrier_init(smem + 240, 128); - mbarrier_init(smem + 248, 128); - // --- pipeline 'work_pipe' --- - // work_full: 2 barriers, init_count=1 - mbarrier_init(smem + 256, 1); - mbarrier_init(smem + 264, 1); - // work_empty: 2 barriers, init_count=512 - mbarrier_init(smem + 272, 512); - mbarrier_init(smem + 280, 512); - // --- pipeline 'throttle_pipe' --- - // throttle_full: 2 barriers, init_count=128 - mbarrier_init(smem + 288, 128); - mbarrier_init(smem + 296, 128); - // throttle_empty: 2 barriers, init_count=32 - mbarrier_init(smem + 304, 32); - mbarrier_init(smem + 312, 32); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - - __syncwarp(); - - // TMEM alloc (512 columns, 512 used) - volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 320); - if (warp == 0) { - int _tmem_hold = smem + 320; - asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); - __syncwarp(); - asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); - } - - __syncthreads(); - asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); - - const int taddr = tmem_addr_storage[0]; - - // Kernel post-init ops - const int tmem_tmem = taddr; - - // ---- Ordered hardware-WG register redistribution ---- - // Inc phase consumes the registers released above. - if (warp >= 8 && warp <= 11) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 136;"); - } - - // ---- Role: softmax ---- - if (warp <= 3) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 152;"); - { // softmax_main - asm volatile("griddepcontrol.wait;" ::: "memory"); - float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; - unsigned int work_stage = 0; - unsigned int query_idx = blockIdx.x; - unsigned int global_tile = 0; - unsigned int stats_cursor = 0; - const int warp_in_compute = warp; - const int tmem_row_origin = warp_in_compute * 32; - const int logical_row_origin = warp_in_compute * 16; - const int my_row = logical_row_origin + lane % 16; - const int col_half = lane / 16; - unsigned int _phase_work_full = 0; - #pragma unroll 1 - for (unsigned int work_iter = 0; work_iter < num_query_tokens; work_iter++) { - if (query_idx >= (unsigned int)num_query_tokens) { - break; - } - int _max_4 = ((sparse_topk_lens[query_idx] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx] + sparse_topk_lens_offset) : (0)); - int _min_2 = ((_max_4) < (sparse_topk) ? (_max_4) : (sparse_topk)); - int active_topk = _min_2; - int _max_5 = (((active_topk + 128 - 1) / 128) > (1) ? ((active_topk + 128 - 1) / 128) : (1)); - int num_loop_steps = _max_5; - float row_max = -CAKE_INF; - float row_sum = 0.0f; - if (has_sinks != 0 && my_row < num_heads) { - row_max = sinks[my_row] * 1.4426950408889634f / softmax_scale_log2; - row_sum = 1.0f; - } - if (num_loop_steps > 0) { - unsigned int first_stats_stage = stats_cursor & 1; - unsigned int first_stats_empty_phase = stats_cursor >> 1 & 1 ^ 1; - mbarrier_wait(stats_empty_addr + (first_stats_stage) * 8, first_stats_empty_phase); - } - #pragma unroll 1 - for (int tile = 0; tile < num_loop_steps; tile++) { - unsigned int sync_stage = global_tile & 1; - unsigned int sync_phase = global_tile >> 1 & 1; - mbarrier_wait(s_full_addr + (sync_stage) * 8, sync_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - int score_col = ((sync_stage != 0) ? 128 : 0); - int score_addr = taddr + (unsigned int)score_col + (unsigned int)(tmem_row_origin << 16); - float _tmem_load_0[64]; - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_0[0])), "=r"(*reinterpret_cast(&_tmem_load_0[1])), "=r"(*reinterpret_cast(&_tmem_load_0[2])), "=r"(*reinterpret_cast(&_tmem_load_0[3])), "=r"(*reinterpret_cast(&_tmem_load_0[4])), "=r"(*reinterpret_cast(&_tmem_load_0[5])), "=r"(*reinterpret_cast(&_tmem_load_0[6])), "=r"(*reinterpret_cast(&_tmem_load_0[7])), "=r"(*reinterpret_cast(&_tmem_load_0[8])), "=r"(*reinterpret_cast(&_tmem_load_0[9])), "=r"(*reinterpret_cast(&_tmem_load_0[10])), "=r"(*reinterpret_cast(&_tmem_load_0[11])), "=r"(*reinterpret_cast(&_tmem_load_0[12])), "=r"(*reinterpret_cast(&_tmem_load_0[13])), "=r"(*reinterpret_cast(&_tmem_load_0[14])), "=r"(*reinterpret_cast(&_tmem_load_0[15])), "=r"(*reinterpret_cast(&_tmem_load_0[16])), "=r"(*reinterpret_cast(&_tmem_load_0[17])), "=r"(*reinterpret_cast(&_tmem_load_0[18])), "=r"(*reinterpret_cast(&_tmem_load_0[19])), "=r"(*reinterpret_cast(&_tmem_load_0[20])), "=r"(*reinterpret_cast(&_tmem_load_0[21])), "=r"(*reinterpret_cast(&_tmem_load_0[22])), "=r"(*reinterpret_cast(&_tmem_load_0[23])), "=r"(*reinterpret_cast(&_tmem_load_0[24])), "=r"(*reinterpret_cast(&_tmem_load_0[25])), "=r"(*reinterpret_cast(&_tmem_load_0[26])), "=r"(*reinterpret_cast(&_tmem_load_0[27])), "=r"(*reinterpret_cast(&_tmem_load_0[28])), "=r"(*reinterpret_cast(&_tmem_load_0[29])), "=r"(*reinterpret_cast(&_tmem_load_0[30])), "=r"(*reinterpret_cast(&_tmem_load_0[31])) - : "r"(score_addr)); - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_0[32])), "=r"(*reinterpret_cast(&_tmem_load_0[33])), "=r"(*reinterpret_cast(&_tmem_load_0[34])), "=r"(*reinterpret_cast(&_tmem_load_0[35])), "=r"(*reinterpret_cast(&_tmem_load_0[36])), "=r"(*reinterpret_cast(&_tmem_load_0[37])), "=r"(*reinterpret_cast(&_tmem_load_0[38])), "=r"(*reinterpret_cast(&_tmem_load_0[39])), "=r"(*reinterpret_cast(&_tmem_load_0[40])), "=r"(*reinterpret_cast(&_tmem_load_0[41])), "=r"(*reinterpret_cast(&_tmem_load_0[42])), "=r"(*reinterpret_cast(&_tmem_load_0[43])), "=r"(*reinterpret_cast(&_tmem_load_0[44])), "=r"(*reinterpret_cast(&_tmem_load_0[45])), "=r"(*reinterpret_cast(&_tmem_load_0[46])), "=r"(*reinterpret_cast(&_tmem_load_0[47])), "=r"(*reinterpret_cast(&_tmem_load_0[48])), "=r"(*reinterpret_cast(&_tmem_load_0[49])), "=r"(*reinterpret_cast(&_tmem_load_0[50])), "=r"(*reinterpret_cast(&_tmem_load_0[51])), "=r"(*reinterpret_cast(&_tmem_load_0[52])), "=r"(*reinterpret_cast(&_tmem_load_0[53])), "=r"(*reinterpret_cast(&_tmem_load_0[54])), "=r"(*reinterpret_cast(&_tmem_load_0[55])), "=r"(*reinterpret_cast(&_tmem_load_0[56])), "=r"(*reinterpret_cast(&_tmem_load_0[57])), "=r"(*reinterpret_cast(&_tmem_load_0[58])), "=r"(*reinterpret_cast(&_tmem_load_0[59])), "=r"(*reinterpret_cast(&_tmem_load_0[60])), "=r"(*reinterpret_cast(&_tmem_load_0[61])), "=r"(*reinterpret_cast(&_tmem_load_0[62])), "=r"(*reinterpret_cast(&_tmem_load_0[63])) - : "r"(score_addr + 32)); - int valid_cols = active_topk - tile * 128 - col_half * 64; - if (valid_cols < 0) { - valid_cols = 0; - } - if (valid_cols > 64) { - valid_cols = 64; - } - uint32_t _slice_lo_mask_0; - { - int _lim_0 = valid_cols; - if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } - else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); - } - } - if (!(_slice_lo_mask_0 & (1u << 0))) _tmem_load_0[0] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 1))) _tmem_load_0[1] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 2))) _tmem_load_0[2] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 3))) _tmem_load_0[3] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 4))) _tmem_load_0[4] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 5))) _tmem_load_0[5] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 6))) _tmem_load_0[6] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 7))) _tmem_load_0[7] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 8))) _tmem_load_0[8] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 9))) _tmem_load_0[9] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 10))) _tmem_load_0[10] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 11))) _tmem_load_0[11] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 12))) _tmem_load_0[12] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 13))) _tmem_load_0[13] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 14))) _tmem_load_0[14] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 15))) _tmem_load_0[15] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 16))) _tmem_load_0[16] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 17))) _tmem_load_0[17] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 18))) _tmem_load_0[18] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 19))) _tmem_load_0[19] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 20))) _tmem_load_0[20] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 21))) _tmem_load_0[21] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 22))) _tmem_load_0[22] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 23))) _tmem_load_0[23] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 24))) _tmem_load_0[24] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 25))) _tmem_load_0[25] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 26))) _tmem_load_0[26] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 27))) _tmem_load_0[27] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 28))) _tmem_load_0[28] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 29))) _tmem_load_0[29] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 30))) _tmem_load_0[30] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 31))) _tmem_load_0[31] = -CAKE_INF; - uint32_t _slice_lo_mask_1; - { - int _lim_1 = valid_cols - 32; - if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } - else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); - } - } - if (!(_slice_lo_mask_1 & (1u << 0))) _tmem_load_0[32] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 1))) _tmem_load_0[33] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 2))) _tmem_load_0[34] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 3))) _tmem_load_0[35] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 4))) _tmem_load_0[36] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 5))) _tmem_load_0[37] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 6))) _tmem_load_0[38] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 7))) _tmem_load_0[39] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 8))) _tmem_load_0[40] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 9))) _tmem_load_0[41] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 10))) _tmem_load_0[42] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 11))) _tmem_load_0[43] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 12))) _tmem_load_0[44] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 13))) _tmem_load_0[45] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 14))) _tmem_load_0[46] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 15))) _tmem_load_0[47] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 16))) _tmem_load_0[48] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 17))) _tmem_load_0[49] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 18))) _tmem_load_0[50] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 19))) _tmem_load_0[51] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 20))) _tmem_load_0[52] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 21))) _tmem_load_0[53] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 22))) _tmem_load_0[54] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 23))) _tmem_load_0[55] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 24))) _tmem_load_0[56] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 25))) _tmem_load_0[57] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 26))) _tmem_load_0[58] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 27))) _tmem_load_0[59] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 28))) _tmem_load_0[60] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 29))) _tmem_load_0[61] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 30))) _tmem_load_0[62] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 31))) _tmem_load_0[63] = -CAKE_INF; - float2 _reg_reduce_max2_2 = {-CAKE_INF, -CAKE_INF}; - row_max_x32_accum(&_tmem_load_0[0], _reg_reduce_max2_2); - row_max_x32_accum(&_tmem_load_0[32], _reg_reduce_max2_2); - float _tmem_load_0_max = row_max_reduce(_reg_reduce_max2_2); - float tile_max = _tmem_load_0_max; - float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, tile_max, 16); - float _max_6 = max_noftz(tile_max, _shfl_xor_0); - tile_max = _max_6; - float old_max = row_max; - float _max_7 = max_noftz(row_max, tile_max); - float new_max = _max_7; - float _fma_0 = __fmaf_rn(old_max, softmax_scale_log2, (-new_max) * softmax_scale_log2); - float delta = _fma_0; - float _exp2_0 = approx_exp2(delta); - float acc_scale = ((old_max > -CAKE_INF) ? _exp2_0 : 1.0f); - unsigned int stats_stage = stats_cursor & 1; - int stats_col = ((stats_stage != 0) ? 128 : 0); - float max_stats[2]; - max_stats[0] = old_max; - max_stats[1] = new_max; - asm volatile( - "tcgen05.st.sync.aligned.32x32b.x2.b32" - " [%0], {%1, %2};" - :: "r"(taddr + (unsigned int)stats_col), "f"(max_stats[0]), "f"(max_stats[1])); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - mbarrier_arrive(stats_addr + (stats_stage) * 8); - stats_cursor = stats_cursor + 1; - unsigned int next_stats_stage = stats_cursor & 1; - unsigned int next_stats_empty_phase = stats_cursor >> 1 & 1 ^ 1; - mbarrier_wait(stats_empty_addr + (next_stats_stage) * 8, next_stats_empty_phase); - bool empty_row = new_max == -CAKE_INF; - float safe_max = ((empty_row) ? 0.0f : new_max); - float max_scaled = safe_max * softmax_scale_log2; - const float2 _fma_b2_3 = {softmax_scale_log2, softmax_scale_log2}; - const float2 _fma_c2_4 = {-max_scaled, -max_scaled}; - #pragma unroll - for (int _lf = 0; _lf < 32; _lf++) - fma_f32x2_inplace(&reinterpret_cast(_tmem_load_0)[_lf], _fma_b2_3, _fma_c2_4); - #pragma unroll - for (int _le = 0; _le < 64; _le++) { - _tmem_load_0[_le] = approx_exp2(_tmem_load_0[_le]); - } - float2 _reg_reduce_sum2_5 = make_float2(0.0f, 0.0f); - softmax_block_sum(&_tmem_load_0[0], &_reg_reduce_sum2_5); - softmax_block_sum(&_tmem_load_0[32], &_reg_reduce_sum2_5); - float _tmem_load_0_sum = _reg_reduce_sum2_5.x + _reg_reduce_sum2_5.y; - float block_sum = _tmem_load_0_sum; - float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, block_sum, 16); - block_sum = block_sum + _shfl_xor_1; - float _fma_1 = __fmaf_rn(row_sum, acc_scale, block_sum); - row_sum = _fma_1; - row_max = new_max; - if (tile == num_loop_steps - 1) { - float final_stats[2]; - float _max_8 = max_noftz(row_sum, 1.1754943508222875e-38f); - final_stats[0] = _max_8; - final_stats[1] = row_max; - int final_stats_col = ((next_stats_stage != 0) ? 128 : 0); - asm volatile( - "tcgen05.st.sync.aligned.32x32b.x2.b32" - " [%0], {%1, %2};" - :: "r"(taddr + (unsigned int)final_stats_col), "f"(final_stats[0]), "f"(final_stats[1])); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - mbarrier_arrive(stats_addr + (next_stats_stage) * 8); - stats_cursor = stats_cursor + 1; - unsigned int tail_stats_stage = stats_cursor & 1; - unsigned int tail_stats_empty_phase = stats_cursor >> 1 & 1 ^ 1; - mbarrier_wait(stats_empty_addr + (tail_stats_stage) * 8, tail_stats_empty_phase); - mbarrier_arrive(stats_addr + (tail_stats_stage) * 8); - stats_cursor = stats_cursor + 1; - } - unsigned int packed_p[32]; - #pragma unroll - for (int _lp = 0; _lp < 32; _lp++) { - __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_0[_lp*2 + 0], _tmem_load_0[_lp*2+1 + 0])); - packed_p[_lp] = *(uint32_t*)&_bf2; - } - int p_addr = taddr + (unsigned int)score_col + 32 + (unsigned int)(tmem_row_origin << 16); - asm volatile( - "tcgen05.st.sync.aligned.16x32bx2.x16.b32" - " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" - :: "r"(p_addr), "r"(*reinterpret_cast(&packed_p[0])), "r"(*reinterpret_cast(&packed_p[1])), "r"(*reinterpret_cast(&packed_p[2])), "r"(*reinterpret_cast(&packed_p[3])), "r"(*reinterpret_cast(&packed_p[4])), "r"(*reinterpret_cast(&packed_p[5])), "r"(*reinterpret_cast(&packed_p[6])), "r"(*reinterpret_cast(&packed_p[7])), "r"(*reinterpret_cast(&packed_p[8])), "r"(*reinterpret_cast(&packed_p[9])), "r"(*reinterpret_cast(&packed_p[10])), "r"(*reinterpret_cast(&packed_p[11])), "r"(*reinterpret_cast(&packed_p[12])), "r"(*reinterpret_cast(&packed_p[13])), "r"(*reinterpret_cast(&packed_p[14])), "r"(*reinterpret_cast(&packed_p[15]))); - asm volatile( - "tcgen05.st.sync.aligned.16x32bx2.x16.b32" - " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" - :: "r"(p_addr + 16), "r"(*reinterpret_cast(&(packed_p + 16)[0])), "r"(*reinterpret_cast(&(packed_p + 16)[1])), "r"(*reinterpret_cast(&(packed_p + 16)[2])), "r"(*reinterpret_cast(&(packed_p + 16)[3])), "r"(*reinterpret_cast(&(packed_p + 16)[4])), "r"(*reinterpret_cast(&(packed_p + 16)[5])), "r"(*reinterpret_cast(&(packed_p + 16)[6])), "r"(*reinterpret_cast(&(packed_p + 16)[7])), "r"(*reinterpret_cast(&(packed_p + 16)[8])), "r"(*reinterpret_cast(&(packed_p + 16)[9])), "r"(*reinterpret_cast(&(packed_p + 16)[10])), "r"(*reinterpret_cast(&(packed_p + 16)[11])), "r"(*reinterpret_cast(&(packed_p + 16)[12])), "r"(*reinterpret_cast(&(packed_p + 16)[13])), "r"(*reinterpret_cast(&(packed_p + 16)[14])), "r"(*reinterpret_cast(&(packed_p + 16)[15]))); - int p_addr_high = p_addr + 1048576; - asm volatile( - "tcgen05.st.sync.aligned.16x32bx2.x16.b32" - " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" - :: "r"(p_addr_high), "r"(*reinterpret_cast(&packed_p[0])), "r"(*reinterpret_cast(&packed_p[1])), "r"(*reinterpret_cast(&packed_p[2])), "r"(*reinterpret_cast(&packed_p[3])), "r"(*reinterpret_cast(&packed_p[4])), "r"(*reinterpret_cast(&packed_p[5])), "r"(*reinterpret_cast(&packed_p[6])), "r"(*reinterpret_cast(&packed_p[7])), "r"(*reinterpret_cast(&packed_p[8])), "r"(*reinterpret_cast(&packed_p[9])), "r"(*reinterpret_cast(&packed_p[10])), "r"(*reinterpret_cast(&packed_p[11])), "r"(*reinterpret_cast(&packed_p[12])), "r"(*reinterpret_cast(&packed_p[13])), "r"(*reinterpret_cast(&packed_p[14])), "r"(*reinterpret_cast(&packed_p[15]))); - asm volatile( - "tcgen05.st.sync.aligned.16x32bx2.x16.b32" - " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" - :: "r"(p_addr_high + 16), "r"(*reinterpret_cast(&(packed_p + 16)[0])), "r"(*reinterpret_cast(&(packed_p + 16)[1])), "r"(*reinterpret_cast(&(packed_p + 16)[2])), "r"(*reinterpret_cast(&(packed_p + 16)[3])), "r"(*reinterpret_cast(&(packed_p + 16)[4])), "r"(*reinterpret_cast(&(packed_p + 16)[5])), "r"(*reinterpret_cast(&(packed_p + 16)[6])), "r"(*reinterpret_cast(&(packed_p + 16)[7])), "r"(*reinterpret_cast(&(packed_p + 16)[8])), "r"(*reinterpret_cast(&(packed_p + 16)[9])), "r"(*reinterpret_cast(&(packed_p + 16)[10])), "r"(*reinterpret_cast(&(packed_p + 16)[11])), "r"(*reinterpret_cast(&(packed_p + 16)[12])), "r"(*reinterpret_cast(&(packed_p + 16)[13])), "r"(*reinterpret_cast(&(packed_p + 16)[14])), "r"(*reinterpret_cast(&(packed_p + 16)[15]))); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(p_full_addr + (sync_stage) * 8); - global_tile = global_tile + 1; - } - mbarrier_wait(work_full_addr + (work_stage) * 8, _phase_work_full); - uint32_t _clc_valid_2 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_2) - : "r"(work_response_addr + work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_2 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_2) - : "r"(work_response_addr + work_stage * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(work_empty_addr + (work_stage) * 8); - work_stage += 1; - if (work_stage == 2) { work_stage = 0; _phase_work_full ^= 1; } - unsigned int valid = _clc_valid_2; - query_idx = _clc_ctaid_2; - if (valid == 0) { - break; - } - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - } - } - // ---- Role: correction ---- - if (warp >= 4 && warp <= 7) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 144;"); - { // correction_main - asm volatile("griddepcontrol.wait;" ::: "memory"); - float softmax_scale_log2_1 = bmm1_scale[0] * 1.4426950408889634f; - float output_scale = bmm2_scale[0]; - unsigned int work_stage_1 = 0; - unsigned int query_idx_1 = blockIdx.x; - unsigned int global_tile_1 = 0; - unsigned int stats_cursor_1 = 0; - const int warp_in_role = warp - 4; - const int tmem_row_origin_1 = warp_in_role * 32; - const int logical_row_origin_1 = warp_in_role * 16; - const int my_row_1 = logical_row_origin_1 + lane % 16; - const int col_half_1 = lane / 16; - unsigned int _phase_work_full_1 = 0; - #pragma unroll 1 - for (unsigned int work_iter_1 = 0; work_iter_1 < num_query_tokens; work_iter_1++) { - if (query_idx_1 >= (unsigned int)num_query_tokens) { - break; - } - int _max_9 = ((sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) : (0)); - int _min_3 = ((_max_9) < (sparse_topk) ? (_max_9) : (sparse_topk)); - int active_topk_1 = _min_3; - int _max_10 = (((active_topk_1 + 128 - 1) / 128) > (1) ? ((active_topk_1 + 128 - 1) / 128) : (1)); - int num_loop_steps_1 = _max_10; - unsigned int first_stats_stage_1 = stats_cursor_1 & 1; - unsigned int first_stats_phase = stats_cursor_1 >> 1 & 1; - mbarrier_wait(stats_addr + (first_stats_stage_1) * 8, first_stats_phase); - mbarrier_arrive(stats_empty_addr + (first_stats_stage_1) * 8); - stats_cursor_1 = stats_cursor_1 + 1; - global_tile_1 = global_tile_1 + 1; - #pragma unroll 1 - for (int tile_1 = 0; tile_1 < num_loop_steps_1 - 1; tile_1++) { - unsigned int stats_stage_1 = stats_cursor_1 & 1; - unsigned int stats_phase = stats_cursor_1 >> 1 & 1; - mbarrier_wait(stats_addr + (stats_stage_1) * 8, stats_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - int stats_col_1 = ((stats_stage_1 != 0) ? 128 : 0); - float _tmem_load_1[2]; - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x2.b32" - " {%0, %1}, [%2];" - : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]) - : "r"(taddr + (unsigned int)stats_col_1)); - mbarrier_arrive(stats_empty_addr + (stats_stage_1) * 8); - stats_cursor_1 = stats_cursor_1 + 1; - float max_diff = _tmem_load_1[0] - _tmem_load_1[1]; - float _exp2_1 = approx_exp2(softmax_scale_log2_1 * max_diff); - float acc_scale_1 = ((max_diff != 0.0f) ? _exp2_1 : 1.0f); - mbarrier_wait(o_full_addr, global_tile_1 - 1 & 1); - asm volatile("tcgen05.fence::after_thread_sync;"); - int _vote_0 = __any_sync(0xFFFFFFFF, acc_scale_1 < 1.0f); - int any_rescale = _vote_0; - if (any_rescale != 0) { - #pragma unroll - for (int v_stage = 0; v_stage < 4; v_stage++) { - int output_addr = taddr + 256 + (unsigned int)((v_stage & 1) * 128) + (unsigned int)((v_stage >> 1) * 16 << 16) + (unsigned int)(tmem_row_origin_1 << 16); - float _tmem_load_2[64]; - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_2[0])), "=r"(*reinterpret_cast(&_tmem_load_2[1])), "=r"(*reinterpret_cast(&_tmem_load_2[2])), "=r"(*reinterpret_cast(&_tmem_load_2[3])), "=r"(*reinterpret_cast(&_tmem_load_2[4])), "=r"(*reinterpret_cast(&_tmem_load_2[5])), "=r"(*reinterpret_cast(&_tmem_load_2[6])), "=r"(*reinterpret_cast(&_tmem_load_2[7])), "=r"(*reinterpret_cast(&_tmem_load_2[8])), "=r"(*reinterpret_cast(&_tmem_load_2[9])), "=r"(*reinterpret_cast(&_tmem_load_2[10])), "=r"(*reinterpret_cast(&_tmem_load_2[11])), "=r"(*reinterpret_cast(&_tmem_load_2[12])), "=r"(*reinterpret_cast(&_tmem_load_2[13])), "=r"(*reinterpret_cast(&_tmem_load_2[14])), "=r"(*reinterpret_cast(&_tmem_load_2[15])), "=r"(*reinterpret_cast(&_tmem_load_2[16])), "=r"(*reinterpret_cast(&_tmem_load_2[17])), "=r"(*reinterpret_cast(&_tmem_load_2[18])), "=r"(*reinterpret_cast(&_tmem_load_2[19])), "=r"(*reinterpret_cast(&_tmem_load_2[20])), "=r"(*reinterpret_cast(&_tmem_load_2[21])), "=r"(*reinterpret_cast(&_tmem_load_2[22])), "=r"(*reinterpret_cast(&_tmem_load_2[23])), "=r"(*reinterpret_cast(&_tmem_load_2[24])), "=r"(*reinterpret_cast(&_tmem_load_2[25])), "=r"(*reinterpret_cast(&_tmem_load_2[26])), "=r"(*reinterpret_cast(&_tmem_load_2[27])), "=r"(*reinterpret_cast(&_tmem_load_2[28])), "=r"(*reinterpret_cast(&_tmem_load_2[29])), "=r"(*reinterpret_cast(&_tmem_load_2[30])), "=r"(*reinterpret_cast(&_tmem_load_2[31])) - : "r"(output_addr)); - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_2[32])), "=r"(*reinterpret_cast(&_tmem_load_2[33])), "=r"(*reinterpret_cast(&_tmem_load_2[34])), "=r"(*reinterpret_cast(&_tmem_load_2[35])), "=r"(*reinterpret_cast(&_tmem_load_2[36])), "=r"(*reinterpret_cast(&_tmem_load_2[37])), "=r"(*reinterpret_cast(&_tmem_load_2[38])), "=r"(*reinterpret_cast(&_tmem_load_2[39])), "=r"(*reinterpret_cast(&_tmem_load_2[40])), "=r"(*reinterpret_cast(&_tmem_load_2[41])), "=r"(*reinterpret_cast(&_tmem_load_2[42])), "=r"(*reinterpret_cast(&_tmem_load_2[43])), "=r"(*reinterpret_cast(&_tmem_load_2[44])), "=r"(*reinterpret_cast(&_tmem_load_2[45])), "=r"(*reinterpret_cast(&_tmem_load_2[46])), "=r"(*reinterpret_cast(&_tmem_load_2[47])), "=r"(*reinterpret_cast(&_tmem_load_2[48])), "=r"(*reinterpret_cast(&_tmem_load_2[49])), "=r"(*reinterpret_cast(&_tmem_load_2[50])), "=r"(*reinterpret_cast(&_tmem_load_2[51])), "=r"(*reinterpret_cast(&_tmem_load_2[52])), "=r"(*reinterpret_cast(&_tmem_load_2[53])), "=r"(*reinterpret_cast(&_tmem_load_2[54])), "=r"(*reinterpret_cast(&_tmem_load_2[55])), "=r"(*reinterpret_cast(&_tmem_load_2[56])), "=r"(*reinterpret_cast(&_tmem_load_2[57])), "=r"(*reinterpret_cast(&_tmem_load_2[58])), "=r"(*reinterpret_cast(&_tmem_load_2[59])), "=r"(*reinterpret_cast(&_tmem_load_2[60])), "=r"(*reinterpret_cast(&_tmem_load_2[61])), "=r"(*reinterpret_cast(&_tmem_load_2[62])), "=r"(*reinterpret_cast(&_tmem_load_2[63])) - : "r"(output_addr + 32)); - const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; - #pragma unroll - for (int _ls = 0; _ls < 32; _ls++) - mul_f32x2_inplace(&reinterpret_cast(_tmem_load_2)[_ls], _scale2_0); - asm volatile( - "tcgen05.st.sync.aligned.16x32bx2.x64.b32" - " [%0], 64, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64};" - :: "r"(output_addr), "r"(*reinterpret_cast(&_tmem_load_2[0])), "r"(*reinterpret_cast(&_tmem_load_2[1])), "r"(*reinterpret_cast(&_tmem_load_2[2])), "r"(*reinterpret_cast(&_tmem_load_2[3])), "r"(*reinterpret_cast(&_tmem_load_2[4])), "r"(*reinterpret_cast(&_tmem_load_2[5])), "r"(*reinterpret_cast(&_tmem_load_2[6])), "r"(*reinterpret_cast(&_tmem_load_2[7])), "r"(*reinterpret_cast(&_tmem_load_2[8])), "r"(*reinterpret_cast(&_tmem_load_2[9])), "r"(*reinterpret_cast(&_tmem_load_2[10])), "r"(*reinterpret_cast(&_tmem_load_2[11])), "r"(*reinterpret_cast(&_tmem_load_2[12])), "r"(*reinterpret_cast(&_tmem_load_2[13])), "r"(*reinterpret_cast(&_tmem_load_2[14])), "r"(*reinterpret_cast(&_tmem_load_2[15])), "r"(*reinterpret_cast(&_tmem_load_2[16])), "r"(*reinterpret_cast(&_tmem_load_2[17])), "r"(*reinterpret_cast(&_tmem_load_2[18])), "r"(*reinterpret_cast(&_tmem_load_2[19])), "r"(*reinterpret_cast(&_tmem_load_2[20])), "r"(*reinterpret_cast(&_tmem_load_2[21])), "r"(*reinterpret_cast(&_tmem_load_2[22])), "r"(*reinterpret_cast(&_tmem_load_2[23])), "r"(*reinterpret_cast(&_tmem_load_2[24])), "r"(*reinterpret_cast(&_tmem_load_2[25])), "r"(*reinterpret_cast(&_tmem_load_2[26])), "r"(*reinterpret_cast(&_tmem_load_2[27])), "r"(*reinterpret_cast(&_tmem_load_2[28])), "r"(*reinterpret_cast(&_tmem_load_2[29])), "r"(*reinterpret_cast(&_tmem_load_2[30])), "r"(*reinterpret_cast(&_tmem_load_2[31])), "r"(*reinterpret_cast(&_tmem_load_2[32])), "r"(*reinterpret_cast(&_tmem_load_2[33])), "r"(*reinterpret_cast(&_tmem_load_2[34])), "r"(*reinterpret_cast(&_tmem_load_2[35])), "r"(*reinterpret_cast(&_tmem_load_2[36])), "r"(*reinterpret_cast(&_tmem_load_2[37])), "r"(*reinterpret_cast(&_tmem_load_2[38])), "r"(*reinterpret_cast(&_tmem_load_2[39])), "r"(*reinterpret_cast(&_tmem_load_2[40])), "r"(*reinterpret_cast(&_tmem_load_2[41])), "r"(*reinterpret_cast(&_tmem_load_2[42])), "r"(*reinterpret_cast(&_tmem_load_2[43])), "r"(*reinterpret_cast(&_tmem_load_2[44])), "r"(*reinterpret_cast(&_tmem_load_2[45])), "r"(*reinterpret_cast(&_tmem_load_2[46])), "r"(*reinterpret_cast(&_tmem_load_2[47])), "r"(*reinterpret_cast(&_tmem_load_2[48])), "r"(*reinterpret_cast(&_tmem_load_2[49])), "r"(*reinterpret_cast(&_tmem_load_2[50])), "r"(*reinterpret_cast(&_tmem_load_2[51])), "r"(*reinterpret_cast(&_tmem_load_2[52])), "r"(*reinterpret_cast(&_tmem_load_2[53])), "r"(*reinterpret_cast(&_tmem_load_2[54])), "r"(*reinterpret_cast(&_tmem_load_2[55])), "r"(*reinterpret_cast(&_tmem_load_2[56])), "r"(*reinterpret_cast(&_tmem_load_2[57])), "r"(*reinterpret_cast(&_tmem_load_2[58])), "r"(*reinterpret_cast(&_tmem_load_2[59])), "r"(*reinterpret_cast(&_tmem_load_2[60])), "r"(*reinterpret_cast(&_tmem_load_2[61])), "r"(*reinterpret_cast(&_tmem_load_2[62])), "r"(*reinterpret_cast(&_tmem_load_2[63]))); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - } - } - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(o_empty_addr); - global_tile_1 = global_tile_1 + 1; - } - unsigned int final_stats_stage = stats_cursor_1 & 1; - unsigned int final_stats_phase = stats_cursor_1 >> 1 & 1; - mbarrier_wait(stats_addr + (final_stats_stage) * 8, final_stats_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - int final_stats_col_1 = ((final_stats_stage != 0) ? 128 : 0); - float _tmem_load_3[2]; - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x2.b32" - " {%0, %1}, [%2];" - : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]) - : "r"(taddr + (unsigned int)final_stats_col_1)); - mbarrier_arrive(stats_empty_addr + (final_stats_stage) * 8); - stats_cursor_1 = stats_cursor_1 + 1; - unsigned int tail_stats_stage_1 = stats_cursor_1 & 1; - unsigned int tail_stats_phase = stats_cursor_1 >> 1 & 1; - mbarrier_wait(stats_addr + (tail_stats_stage_1) * 8, tail_stats_phase); - mbarrier_arrive(stats_empty_addr + (tail_stats_stage_1) * 8); - stats_cursor_1 = stats_cursor_1 + 1; - mbarrier_wait(o_full_addr, global_tile_1 - 1 & 1); - asm volatile("tcgen05.fence::after_thread_sync;"); - float final_sum = _tmem_load_3[0]; - float _rcp_0 = approx_rcp(final_sum); - float inv_sum = _rcp_0; - int output_base = (query_idx_1 * (unsigned int)num_heads + (unsigned int)my_row_1) * 512; - #pragma unroll - for (int v_stage_1 = 0; v_stage_1 < 4; v_stage_1++) { - int output_addr_1 = taddr + 256 + (unsigned int)((v_stage_1 & 1) * 128) + (unsigned int)((v_stage_1 >> 1) * 16 << 16) + (unsigned int)(tmem_row_origin_1 << 16); - float _tmem_load_4[64]; - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_4[0])), "=r"(*reinterpret_cast(&_tmem_load_4[1])), "=r"(*reinterpret_cast(&_tmem_load_4[2])), "=r"(*reinterpret_cast(&_tmem_load_4[3])), "=r"(*reinterpret_cast(&_tmem_load_4[4])), "=r"(*reinterpret_cast(&_tmem_load_4[5])), "=r"(*reinterpret_cast(&_tmem_load_4[6])), "=r"(*reinterpret_cast(&_tmem_load_4[7])), "=r"(*reinterpret_cast(&_tmem_load_4[8])), "=r"(*reinterpret_cast(&_tmem_load_4[9])), "=r"(*reinterpret_cast(&_tmem_load_4[10])), "=r"(*reinterpret_cast(&_tmem_load_4[11])), "=r"(*reinterpret_cast(&_tmem_load_4[12])), "=r"(*reinterpret_cast(&_tmem_load_4[13])), "=r"(*reinterpret_cast(&_tmem_load_4[14])), "=r"(*reinterpret_cast(&_tmem_load_4[15])), "=r"(*reinterpret_cast(&_tmem_load_4[16])), "=r"(*reinterpret_cast(&_tmem_load_4[17])), "=r"(*reinterpret_cast(&_tmem_load_4[18])), "=r"(*reinterpret_cast(&_tmem_load_4[19])), "=r"(*reinterpret_cast(&_tmem_load_4[20])), "=r"(*reinterpret_cast(&_tmem_load_4[21])), "=r"(*reinterpret_cast(&_tmem_load_4[22])), "=r"(*reinterpret_cast(&_tmem_load_4[23])), "=r"(*reinterpret_cast(&_tmem_load_4[24])), "=r"(*reinterpret_cast(&_tmem_load_4[25])), "=r"(*reinterpret_cast(&_tmem_load_4[26])), "=r"(*reinterpret_cast(&_tmem_load_4[27])), "=r"(*reinterpret_cast(&_tmem_load_4[28])), "=r"(*reinterpret_cast(&_tmem_load_4[29])), "=r"(*reinterpret_cast(&_tmem_load_4[30])), "=r"(*reinterpret_cast(&_tmem_load_4[31])) - : "r"(output_addr_1)); - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_4[32])), "=r"(*reinterpret_cast(&_tmem_load_4[33])), "=r"(*reinterpret_cast(&_tmem_load_4[34])), "=r"(*reinterpret_cast(&_tmem_load_4[35])), "=r"(*reinterpret_cast(&_tmem_load_4[36])), "=r"(*reinterpret_cast(&_tmem_load_4[37])), "=r"(*reinterpret_cast(&_tmem_load_4[38])), "=r"(*reinterpret_cast(&_tmem_load_4[39])), "=r"(*reinterpret_cast(&_tmem_load_4[40])), "=r"(*reinterpret_cast(&_tmem_load_4[41])), "=r"(*reinterpret_cast(&_tmem_load_4[42])), "=r"(*reinterpret_cast(&_tmem_load_4[43])), "=r"(*reinterpret_cast(&_tmem_load_4[44])), "=r"(*reinterpret_cast(&_tmem_load_4[45])), "=r"(*reinterpret_cast(&_tmem_load_4[46])), "=r"(*reinterpret_cast(&_tmem_load_4[47])), "=r"(*reinterpret_cast(&_tmem_load_4[48])), "=r"(*reinterpret_cast(&_tmem_load_4[49])), "=r"(*reinterpret_cast(&_tmem_load_4[50])), "=r"(*reinterpret_cast(&_tmem_load_4[51])), "=r"(*reinterpret_cast(&_tmem_load_4[52])), "=r"(*reinterpret_cast(&_tmem_load_4[53])), "=r"(*reinterpret_cast(&_tmem_load_4[54])), "=r"(*reinterpret_cast(&_tmem_load_4[55])), "=r"(*reinterpret_cast(&_tmem_load_4[56])), "=r"(*reinterpret_cast(&_tmem_load_4[57])), "=r"(*reinterpret_cast(&_tmem_load_4[58])), "=r"(*reinterpret_cast(&_tmem_load_4[59])), "=r"(*reinterpret_cast(&_tmem_load_4[60])), "=r"(*reinterpret_cast(&_tmem_load_4[61])), "=r"(*reinterpret_cast(&_tmem_load_4[62])), "=r"(*reinterpret_cast(&_tmem_load_4[63])) - : "r"(output_addr_1 + 32)); - int col_base = v_stage_1 * 128 + col_half_1 * 64; - if (my_row_1 < num_heads) { - #pragma unroll - for (int offset = 0; offset < 64; offset += 8) { - { - const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; - #if __CUDA_ARCH__ >= 1000 - #pragma unroll - for (int _ps = 0; _ps < 4; _ps++) - mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_4[offset])[_ps], _prescale2_1); - #else - #pragma unroll - for (int _ps = 0; _ps < 8; _ps++) - _tmem_load_4[offset + _ps] *= inv_sum * output_scale; - #endif - __nv_bfloat162 _pk[4]; - _pk[0] = __floats2bfloat162_rn(_tmem_load_4[offset + 0], _tmem_load_4[offset + 1]); - _pk[1] = __floats2bfloat162_rn(_tmem_load_4[offset + 2], _tmem_load_4[offset + 3]); - _pk[2] = __floats2bfloat162_rn(_tmem_load_4[offset + 4], _tmem_load_4[offset + 5]); - _pk[3] = __floats2bfloat162_rn(_tmem_load_4[offset + 6], _tmem_load_4[offset + 7]); - *reinterpret_cast(&((__nv_bfloat16*)(O + (output_base + col_base + offset)))[0]) = *reinterpret_cast(&_pk[0]); - } - } - } - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(o_empty_addr); - mbarrier_wait(work_full_addr + (work_stage_1) * 8, _phase_work_full_1); - uint32_t _clc_valid_3 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_3) - : "r"(work_response_addr + work_stage_1 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_3 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_3) - : "r"(work_response_addr + work_stage_1 * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(work_empty_addr + (work_stage_1) * 8); - work_stage_1 += 1; - if (work_stage_1 == 2) { work_stage_1 = 0; _phase_work_full_1 ^= 1; } - unsigned int valid_1 = _clc_valid_3; - query_idx_1 = _clc_ctaid_3; - if (valid_1 == 0) { - break; - } - } - asm volatile("barrier.sync 7, 128;" ::: "memory"); - if (warp == 4) { - int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); - asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); - } - } - } - // ---- Role: mma_warp ---- - if (warp == 8) { - { // mma_warp_main - unsigned int work_stage_2 = 0; - unsigned int query_idx_2 = blockIdx.x; - unsigned int global_tile_2 = 0; - unsigned int _phase_work_full_2 = 0; - #pragma unroll 1 - for (unsigned int work_iter_2 = 0; work_iter_2 < num_query_tokens; work_iter_2++) { - if (query_idx_2 >= (unsigned int)num_query_tokens) { - break; - } - int _max_2 = ((sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) : (0)); - int _min_1 = ((_max_2) < (sparse_topk) ? (_max_2) : (sparse_topk)); - int active_topk_2 = _min_1; - int _max_3 = (((active_topk_2 + 128 - 1) / 128) > (1) ? ((active_topk_2 + 128 - 1) / 128) : (1)); - int num_loop_steps_2 = _max_3; - unsigned int query_phase = work_iter_2 & 1; - mbarrier_wait(q_full_addr, query_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - if (num_loop_steps_2 > 0) { - unsigned int score_stage = global_tile_2 & 1; - int score_col_1 = ((score_stage != 0) ? 128 : 0); - #pragma unroll - for (int k_stage = 0; k_stage < 4; k_stage++) { - mbarrier_wait(kv_full_addr + (k_stage) * 8, 0); - asm volatile("tcgen05.fence::after_thread_sync;"); - int _mma_a_lo_0 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (k_stage) * 1024); - int _mma_b_lo_0 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (k_stage) * 2048); - { - uint64_t _mma_ss_a_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_0); - uint64_t _mma_ss_b_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_0); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, ((k_stage == 0) ? 0 : 1)); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 506U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 1018U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); - } - } - if (k_stage == 3) { - elect_commit(s_full_addr + (score_stage) * 8); - } - elect_commit(kv_empty_addr + (k_stage) * 8); - } - #pragma unroll 1 - for (int tile_2 = 0; tile_2 < num_loop_steps_2 - 1; tile_2++) { - unsigned int next_global_tile = global_tile_2 + 1; - unsigned int next_score_stage = next_global_tile & 1; - int next_score_col = ((next_score_stage != 0) ? 128 : 0); - #pragma unroll - for (int k_stage_1 = 0; k_stage_1 < 4; k_stage_1++) { - mbarrier_wait(kv_full_addr + (k_stage_1) * 8, 1); - asm volatile("tcgen05.fence::after_thread_sync;"); - int _mma_a_lo_1 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (k_stage_1) * 1024); - int _mma_b_lo_1 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (k_stage_1) * 2048); - { - uint64_t _mma_ss_a_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_1); - uint64_t _mma_ss_b_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_1); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (next_score_col)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, ((k_stage_1 == 0) ? 0 : 1)); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (next_score_col)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (next_score_col)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (next_score_col)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 506U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 1018U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (next_score_col)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (next_score_col)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (next_score_col)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (next_score_col)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); - } - } - if (k_stage_1 == 3) { - elect_commit(s_full_addr + (next_score_stage) * 8); - } - elect_commit(kv_empty_addr + (k_stage_1) * 8); - } - unsigned int sync_stage_1 = global_tile_2 & 1; - unsigned int sync_phase_1 = global_tile_2 >> 1 & 1; - int current_score_col = ((sync_stage_1 != 0) ? 128 : 0); - mbarrier_wait(p_full_addr + (sync_stage_1) * 8, sync_phase_1); - mbarrier_wait(o_empty_addr, 1 - (global_tile_2 & 1)); - asm volatile("tcgen05.fence::after_thread_sync;"); - #pragma unroll - for (int v_stage_2 = 0; v_stage_2 < 4; v_stage_2++) { - mbarrier_wait(kv_full_addr + (v_stage_2) * 8, 0); - asm volatile("tcgen05.fence::after_thread_sync;"); - int output_addr_2 = 256 + (v_stage_2 & 1) * 128 + ((v_stage_2 >> 1) * 16 << 16); - int _mma_b_lo_2 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (v_stage_2) * 2048); - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 dhi, blo, ta, id;\n\t" - ".reg .b64 db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "" - "mov.b32 dhi, 0x40004040;\n\t" - "mov.b32 id, 69272720;\n\t" - "mov.b32 ta, %2;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p0;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "}\n" - :: "r"((tmem_tmem + (output_addr_2))), "r"(_mma_b_lo_2), "r"(tmem_tmem + (current_score_col + 32 + ((v_stage_2 >> 1) * 16 << 16))), "r"(((tile_2 == 0) ? 0 : 1))); - elect_commit(kv_empty_addr + (v_stage_2) * 8); - } - elect_commit(o_full_addr); - global_tile_2 = next_global_tile; - } - elect_commit(q_empty_addr); - int last_tile = num_loop_steps_2 - 1; - unsigned int last_sync_stage = global_tile_2 & 1; - unsigned int last_sync_phase = global_tile_2 >> 1 & 1; - int last_score_col = ((last_sync_stage != 0) ? 128 : 0); - mbarrier_wait(p_full_addr + (last_sync_stage) * 8, last_sync_phase); - mbarrier_wait(o_empty_addr, 1 - (global_tile_2 & 1)); - asm volatile("tcgen05.fence::after_thread_sync;"); - #pragma unroll - for (int v_stage_3 = 0; v_stage_3 < 4; v_stage_3++) { - mbarrier_wait(kv_full_addr + (v_stage_3) * 8, 1); - asm volatile("tcgen05.fence::after_thread_sync;"); - int output_addr_3 = 256 + (v_stage_3 & 1) * 128 + ((v_stage_3 >> 1) * 16 << 16); - int _mma_b_lo_3 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (v_stage_3) * 2048); - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 dhi, blo, ta, id;\n\t" - ".reg .b64 db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "" - "mov.b32 dhi, 0x40004040;\n\t" - "mov.b32 id, 69272720;\n\t" - "mov.b32 ta, %2;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p0;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "}\n" - :: "r"((tmem_tmem + (output_addr_3))), "r"(_mma_b_lo_3), "r"(tmem_tmem + (last_score_col + 32 + ((v_stage_3 >> 1) * 16 << 16))), "r"(((last_tile == 0) ? 0 : 1))); - elect_commit(kv_empty_addr + (v_stage_3) * 8); - } - elect_commit(o_full_addr); - global_tile_2 = global_tile_2 + 1; - } else { - elect_commit(q_empty_addr); - } - mbarrier_wait(work_full_addr + (work_stage_2) * 8, _phase_work_full_2); - uint32_t _clc_valid_1 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_1) - : "r"(work_response_addr + work_stage_2 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_1 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_1) - : "r"(work_response_addr + work_stage_2 * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(work_empty_addr + (work_stage_2) * 8); - work_stage_2 += 1; - if (work_stage_2 == 2) { work_stage_2 = 0; _phase_work_full_2 ^= 1; } - unsigned int valid_2 = _clc_valid_1; - query_idx_2 = _clc_ctaid_1; - if (valid_2 == 0) { - break; - } - } - } - } - // ---- Role: page_warp ---- - if (warp == 9) { - { // page_warp_main - unsigned int page_stage = 0; - unsigned int work_stage_3 = 0; - unsigned int query_idx_3 = blockIdx.x; - unsigned int _phase_page_empty = 1; - unsigned int _phase_work_full_3 = 0; - #pragma unroll 1 - for (unsigned int work_iter_3 = 0; work_iter_3 < num_query_tokens; work_iter_3++) { - if (query_idx_3 >= (unsigned int)num_query_tokens) { - break; - } - int* swa_row = swa_indices + (query_idx_3 * (unsigned int)swa_index_stride); - int* compressed_row = compressed_indices + (query_idx_3 * (unsigned int)compressed_index_stride); - int _max_11 = ((sparse_topk_lens[query_idx_3] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_3] + sparse_topk_lens_offset) : (0)); - int _min_4 = ((_max_11) < (sparse_topk) ? (_max_11) : (sparse_topk)); - int active_topk_3 = _min_4; - int _max_12 = (((active_topk_3 + 128 - 1) / 128) > (1) ? ((active_topk_3 + 128 - 1) / 128) : (1)); - int num_loop_steps_3 = _max_12; - int num_page_pairs = (num_loop_steps_3 + 1) / 2; - int _max_13 = ((active_topk_3 - 1 & -4) > (0) ? (active_topk_3 - 1 & -4) : (0)); - int last_aligned = _max_13; - #pragma unroll 1 - for (int pair = 0; pair < num_page_pairs; pair++) { - int pair_base = pair * 256; - #pragma unroll - for (int duplicate = 0; duplicate < 2; duplicate++) { - mbarrier_wait(page_empty_addr + (page_stage) * 8, _phase_page_empty); - int lane_base = lane * 4; - int page0 = pair_base + lane_base; - int page1 = page0 + 128; - int safe0 = ((page0 <= last_aligned) ? page0 : last_aligned); - int safe1 = ((page1 <= last_aligned) ? page1 : last_aligned); - int* src0 = ((safe0 < 128) ? swa_row + safe0 : compressed_row + (safe0 - 128)); - int* src1 = ((safe1 < 128) ? swa_row + safe1 : compressed_row + (safe1 - 128)); - asm volatile("cp.async.ca.shared::cta.global [%0], [%1], 16;" - :: "r"(smem_page_offsets_addr + page_stage * 1024 + (unsigned int)(lane_base * 4)), "l"(src0)); - asm volatile("cp.async.ca.shared::cta.global [%0], [%1], 16;" - :: "r"(smem_page_offsets_addr + page_stage * 1024 + (unsigned int)((128 + lane_base) * 4)), "l"(src1)); - asm volatile( - "{\n\t" - "cp.async.mbarrier.arrive.noinc.shared::cta.b64 [%0];\n\t" - "}" - :: "r"(page_full_addr + (page_stage) * 8) : "memory"); - page_stage += 1; - if (page_stage == 6) { page_stage = 0; _phase_page_empty ^= 1; } - } - } - mbarrier_wait(work_full_addr + (work_stage_3) * 8, _phase_work_full_3); - uint32_t _clc_valid_4 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_4) - : "r"(work_response_addr + work_stage_3 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_4 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_4) - : "r"(work_response_addr + work_stage_3 * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(work_empty_addr + (work_stage_3) * 8); - work_stage_3 += 1; - if (work_stage_3 == 2) { work_stage_3 = 0; _phase_work_full_3 ^= 1; } - unsigned int valid_3 = _clc_valid_4; - query_idx_3 = _clc_ctaid_4; - if (valid_3 == 0) { - break; - } - } - } - } - // ---- Role: scheduler_warp ---- - if (warp == 10) { - { // scheduler_warp_main - unsigned int work_stage_4 = 0; - unsigned int throttle_stage = 0; - unsigned int query_idx_4 = blockIdx.x; - unsigned int _phase_throttle_full = 0; - unsigned int _phase_work_empty = 1; - unsigned int _phase_work_full_4 = 0; - #pragma unroll 1 - for (unsigned int work_iter_4 = 0; work_iter_4 < num_query_tokens; work_iter_4++) { - if (query_idx_4 >= (unsigned int)num_query_tokens) { - break; - } - mbarrier_wait(throttle_full_addr + (throttle_stage) * 8, _phase_throttle_full); - mbarrier_arrive(throttle_empty_addr + (throttle_stage) * 8); - throttle_stage += 1; - if (throttle_stage == 2) { throttle_stage = 0; _phase_throttle_full ^= 1; } - if (elect_sync()) { - mbarrier_wait(work_empty_addr + (work_stage_4) * 8, _phase_work_empty); - mbarrier_arrive_expect_tx(work_full_addr + (work_stage_4) * 8, 16); - asm volatile( - "fence.proxy.async.shared::cta;\n\t" - "clusterlaunchcontrol.try_cancel.async.shared::cta" - ".mbarrier::complete_tx::bytes.b128" - " [%0], [%1];" - :: "r"(work_response_addr + work_stage_4 * 16 + 0 * 16), "r"(work_full_addr + work_stage_4 * 8) - : "memory"); - } - mbarrier_wait(work_full_addr + (work_stage_4) * 8, _phase_work_full_4); - uint32_t _clc_valid_5 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_5) - : "r"(work_response_addr + work_stage_4 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_5 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_5) - : "r"(work_response_addr + work_stage_4 * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(work_empty_addr + (work_stage_4) * 8); - work_stage_4 += 1; - if (work_stage_4 == 2) { work_stage_4 = 0; _phase_work_empty ^= 1; _phase_work_full_4 ^= 1; } - unsigned int valid_4 = _clc_valid_5; - query_idx_4 = _clc_ctaid_5; - if (valid_4 == 0) { - break; - } - } - } - } - // ---- Role: padding_warp ---- - if (warp == 11) { - { // padding_warp_main - unsigned int work_stage_5 = 0; - unsigned int query_idx_5 = blockIdx.x; - unsigned int _phase_work_full_5 = 0; - #pragma unroll 1 - for (unsigned int work_iter_5 = 0; work_iter_5 < num_query_tokens; work_iter_5++) { - if (query_idx_5 >= (unsigned int)num_query_tokens) { - break; - } - mbarrier_wait(work_full_addr + (work_stage_5) * 8, _phase_work_full_5); - uint32_t _clc_valid_6 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_6) - : "r"(work_response_addr + work_stage_5 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_6 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_6) - : "r"(work_response_addr + work_stage_5 * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(work_empty_addr + (work_stage_5) * 8); - work_stage_5 += 1; - if (work_stage_5 == 2) { work_stage_5 = 0; _phase_work_full_5 ^= 1; } - unsigned int valid_5 = _clc_valid_6; - query_idx_5 = _clc_ctaid_6; - if (valid_5 == 0) { - break; - } - } - } - } - // ---- Role: load_warp ---- - if (warp >= 12 && warp <= 15) { - asm volatile("setmaxnreg.dec.sync.aligned.u32 72;"); - { // load_warp_main - asm volatile("griddepcontrol.wait;" ::: "memory"); - const int load_warp_rank = warp - 12; - unsigned int load_page_stage = 0; - unsigned int page_k_stage = 0; - unsigned int page_v_stage = 0; - int page_k_addr = 0; - int page_v_addr = 0; - unsigned int work_stage_6 = 0; - unsigned int throttle_stage_1 = 0; - unsigned int query_idx_6 = blockIdx.x; - unsigned int _phase_throttle_empty = 1; - unsigned int _phase_page_full = 0; - unsigned int _phase_work_full_6 = 0; - #pragma unroll 1 - for (unsigned int work_iter_6 = 0; work_iter_6 < num_query_tokens; work_iter_6++) { - if (query_idx_6 >= (unsigned int)num_query_tokens) { - break; - } - int _max_0 = ((sparse_topk_lens[query_idx_6] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_6] + sparse_topk_lens_offset) : (0)); - int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); - int active_topk_4 = _min_0; - int _max_1 = (((active_topk_4 + 128 - 1) / 128) > (1) ? ((active_topk_4 + 128 - 1) / 128) : (1)); - int num_loop_steps_4 = _max_1; - mbarrier_wait(throttle_empty_addr + (throttle_stage_1) * 8, _phase_throttle_empty); - mbarrier_arrive(throttle_full_addr + (throttle_stage_1) * 8); - throttle_stage_1 += 1; - if (throttle_stage_1 == 2) { throttle_stage_1 = 0; _phase_throttle_empty ^= 1; } - unsigned int query_phase_1 = work_iter_6 & 1; - mbarrier_wait(q_empty_addr, 1 - query_phase_1); - if (warp == 12) { - if (elect_sync()) { - mbarrier_arrive_expect_tx(q_full_addr, 65536); - #pragma unroll - for (int q_subload = 0; q_subload < 8; q_subload++) { - tma_4d_gmem2smem(smem_q_addr + (unsigned int)(q_subload * 8192), (&tmap_q), 0, 0, q_subload, query_idx_6, q_full_addr); - } - } - } - if (num_loop_steps_4 > 0) { - mbarrier_wait(page_full_addr + (load_page_stage) * 8, _phase_page_full); - page_k_stage = load_page_stage; - page_k_addr = smem_page_offsets_addr + page_k_stage * 1024; - load_page_stage += 1; - if (load_page_stage == 6) { load_page_stage = 0; _phase_page_full ^= 1; } - #pragma unroll - for (int kv_stage = 0; kv_stage < 4; kv_stage++) { - mbarrier_wait(kv_empty_addr + (kv_stage) * 8, 1); - if (warp == 12) { - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage) * 8, 32768); - } - } - #pragma unroll - for (int group = 0; group < 8; group++) { - int row_offset = load_warp_rank * 4 + group * 16; - int raw_rows[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows[0])), "=r"(*reinterpret_cast(&raw_rows[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows[(0) + 3])) - : "r"(page_k_addr + row_offset * 4)); - int row0 = raw_rows[0]; - int row1 = raw_rows[1]; - int row2 = raw_rows[2]; - int row3 = raw_rows[3]; - int dst = smem_kv_addr + (unsigned int)(kv_stage * 32768) + (unsigned int)(row_offset * 64 * 2); - if (elect_sync()) { - { - tma_gather4_gmem2smem(dst, (&tmap_swa_kv), kv_stage * 128, row0, row1, row2, row3, kv_full_addr + (kv_stage) * 8); - tma_gather4_gmem2smem(dst + 16384, (&tmap_swa_kv), kv_stage * 128 + 64, row0, row1, row2, row3, kv_full_addr + (kv_stage) * 8); - } - } - } - } - if (num_loop_steps_4 > 1) { - #pragma unroll - for (int kv_stage_1 = 0; kv_stage_1 < 4; kv_stage_1++) { - mbarrier_wait(kv_empty_addr + (kv_stage_1) * 8, 0); - if (warp == 12) { - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_1) * 8, 32768); - } - } - #pragma unroll - for (int group_1 = 0; group_1 < 8; group_1++) { - int row_offset_1 = load_warp_rank * 4 + group_1 * 16; - int raw_rows_1[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_1[0])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 3])) - : "r"(page_k_addr + 512 + row_offset_1 * 4)); - int row0_1 = raw_rows_1[0]; - int row1_1 = raw_rows_1[1]; - int row2_1 = raw_rows_1[2]; - int row3_1 = raw_rows_1[3]; - int dst_1 = smem_kv_addr + (unsigned int)(kv_stage_1 * 32768) + (unsigned int)(row_offset_1 * 64 * 2); - if (elect_sync()) { - { - tma_gather4_gmem2smem(dst_1, (&tmap_compressed_kv), kv_stage_1 * 128, row0_1, row1_1, row2_1, row3_1, kv_full_addr + (kv_stage_1) * 8); - tma_gather4_gmem2smem(dst_1 + 16384, (&tmap_compressed_kv), kv_stage_1 * 128 + 64, row0_1, row1_1, row2_1, row3_1, kv_full_addr + (kv_stage_1) * 8); - } - } - } - } - if (num_loop_steps_4 > 2) { - mbarrier_arrive(page_empty_addr + (page_k_stage) * 8); - } - mbarrier_wait(page_full_addr + (load_page_stage) * 8, _phase_page_full); - page_v_stage = load_page_stage; - page_v_addr = smem_page_offsets_addr + page_v_stage * 1024; - load_page_stage += 1; - if (load_page_stage == 6) { load_page_stage = 0; _phase_page_full ^= 1; } - #pragma unroll - for (int kv_stage_2 = 0; kv_stage_2 < 4; kv_stage_2++) { - mbarrier_wait(kv_empty_addr + (kv_stage_2) * 8, 1); - if (warp == 12) { - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_2) * 8, 32768); - } - } - #pragma unroll - for (int group_2 = 0; group_2 < 8; group_2++) { - int row_offset_2 = load_warp_rank * 4 + group_2 * 16; - int raw_rows_2[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_2[0])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 3])) - : "r"(page_v_addr + row_offset_2 * 4)); - int row0_2 = raw_rows_2[0]; - int row1_2 = raw_rows_2[1]; - int row2_2 = raw_rows_2[2]; - int row3_2 = raw_rows_2[3]; - int dst_2 = smem_v_addr + (unsigned int)(kv_stage_2 * 32768) + (unsigned int)(row_offset_2 * 64 * 2); - if (elect_sync()) { - { - tma_gather4_gmem2smem(dst_2, (&tmap_swa_kv), kv_stage_2 * 128, row0_2, row1_2, row2_2, row3_2, kv_full_addr + (kv_stage_2) * 8); - tma_gather4_gmem2smem(dst_2 + 16384, (&tmap_swa_kv), kv_stage_2 * 128 + 64, row0_2, row1_2, row2_2, row3_2, kv_full_addr + (kv_stage_2) * 8); - } - } - } - } - } - #pragma unroll 1 - for (int tile_3 = 1; tile_3 < num_loop_steps_4 - 1; tile_3++) { - int next_tile = tile_3 + 1; - if ((next_tile & 1) == 0) { - mbarrier_wait(page_full_addr + (load_page_stage) * 8, _phase_page_full); - page_k_stage = load_page_stage; - page_k_addr = smem_page_offsets_addr + page_k_stage * 1024; - load_page_stage += 1; - if (load_page_stage == 6) { load_page_stage = 0; _phase_page_full ^= 1; } - } - #pragma unroll - for (int kv_stage_3 = 0; kv_stage_3 < 4; kv_stage_3++) { - mbarrier_wait(kv_empty_addr + (kv_stage_3) * 8, 0); - if (warp == 12) { - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_3) * 8, 32768); - } - } - #pragma unroll - for (int group_3 = 0; group_3 < 8; group_3++) { - int row_offset_3 = load_warp_rank * 4 + group_3 * 16; - int raw_rows_3[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_3[0])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 3])) - : "r"(page_k_addr + (next_tile & 1) * 128 * 4 + row_offset_3 * 4)); - int row0_3 = raw_rows_3[0]; - int row1_3 = raw_rows_3[1]; - int row2_3 = raw_rows_3[2]; - int row3_3 = raw_rows_3[3]; - int dst_3 = smem_kv_addr + (unsigned int)(kv_stage_3 * 32768) + (unsigned int)(row_offset_3 * 64 * 2); - if (elect_sync()) { - { - tma_gather4_gmem2smem(dst_3, (&tmap_compressed_kv), kv_stage_3 * 128, row0_3, row1_3, row2_3, row3_3, kv_full_addr + (kv_stage_3) * 8); - tma_gather4_gmem2smem(dst_3 + 16384, (&tmap_compressed_kv), kv_stage_3 * 128 + 64, row0_3, row1_3, row2_3, row3_3, kv_full_addr + (kv_stage_3) * 8); - } - } - } - } - if ((next_tile & 1) != 0 && tile_3 < num_loop_steps_4 - 2) { - mbarrier_arrive(page_empty_addr + (page_k_stage) * 8); - } - if ((tile_3 & 1) == 0) { - mbarrier_wait(page_full_addr + (load_page_stage) * 8, _phase_page_full); - page_v_stage = load_page_stage; - page_v_addr = smem_page_offsets_addr + page_v_stage * 1024; - load_page_stage += 1; - if (load_page_stage == 6) { load_page_stage = 0; _phase_page_full ^= 1; } - } - #pragma unroll - for (int kv_stage_4 = 0; kv_stage_4 < 4; kv_stage_4++) { - mbarrier_wait(kv_empty_addr + (kv_stage_4) * 8, 1); - if (warp == 12) { - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_4) * 8, 32768); - } - } - #pragma unroll - for (int group_4 = 0; group_4 < 8; group_4++) { - int row_offset_4 = load_warp_rank * 4 + group_4 * 16; - int raw_rows_4[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_4[0])), "=r"(*reinterpret_cast(&raw_rows_4[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_4[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_4[(0) + 3])) - : "r"(page_v_addr + (tile_3 & 1) * 128 * 4 + row_offset_4 * 4)); - int row0_4 = raw_rows_4[0]; - int row1_4 = raw_rows_4[1]; - int row2_4 = raw_rows_4[2]; - int row3_4 = raw_rows_4[3]; - int dst_4 = smem_v_addr + (unsigned int)(kv_stage_4 * 32768) + (unsigned int)(row_offset_4 * 64 * 2); - if (elect_sync()) { - { - tma_gather4_gmem2smem(dst_4, (&tmap_compressed_kv), kv_stage_4 * 128, row0_4, row1_4, row2_4, row3_4, kv_full_addr + (kv_stage_4) * 8); - tma_gather4_gmem2smem(dst_4 + 16384, (&tmap_compressed_kv), kv_stage_4 * 128 + 64, row0_4, row1_4, row2_4, row3_4, kv_full_addr + (kv_stage_4) * 8); - } - } - } - } - if ((tile_3 & 1) != 0) { - mbarrier_arrive(page_empty_addr + (page_v_stage) * 8); - } - } - int last_tile_1 = num_loop_steps_4 - 1; - mbarrier_arrive(page_empty_addr + (page_k_stage) * 8); - if ((last_tile_1 & 1) == 0) { - mbarrier_wait(page_full_addr + (load_page_stage) * 8, _phase_page_full); - page_v_stage = load_page_stage; - page_v_addr = smem_page_offsets_addr + page_v_stage * 1024; - load_page_stage += 1; - if (load_page_stage == 6) { load_page_stage = 0; _phase_page_full ^= 1; } - } - if (num_loop_steps_4 == 1) { - #pragma unroll - for (int kv_stage_5 = 0; kv_stage_5 < 4; kv_stage_5++) { - mbarrier_wait(kv_empty_addr + (kv_stage_5) * 8, 0); - if (warp == 12) { - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_5) * 8, 32768); - } - } - #pragma unroll - for (int group_5 = 0; group_5 < 8; group_5++) { - int row_offset_5 = load_warp_rank * 4 + group_5 * 16; - int raw_rows_5[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_5[0])), "=r"(*reinterpret_cast(&raw_rows_5[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_5[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_5[(0) + 3])) - : "r"(page_v_addr + row_offset_5 * 4)); - int row0_5 = raw_rows_5[0]; - int row1_5 = raw_rows_5[1]; - int row2_5 = raw_rows_5[2]; - int row3_5 = raw_rows_5[3]; - int dst_5 = smem_v_addr + (unsigned int)(kv_stage_5 * 32768) + (unsigned int)(row_offset_5 * 64 * 2); - if (elect_sync()) { - { - tma_gather4_gmem2smem(dst_5, (&tmap_swa_kv), kv_stage_5 * 128, row0_5, row1_5, row2_5, row3_5, kv_full_addr + (kv_stage_5) * 8); - tma_gather4_gmem2smem(dst_5 + 16384, (&tmap_swa_kv), kv_stage_5 * 128 + 64, row0_5, row1_5, row2_5, row3_5, kv_full_addr + (kv_stage_5) * 8); - } - } - } - } - } else { - #pragma unroll - for (int kv_stage_6 = 0; kv_stage_6 < 4; kv_stage_6++) { - mbarrier_wait(kv_empty_addr + (kv_stage_6) * 8, 0); - if (warp == 12) { - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_6) * 8, 32768); - } - } - #pragma unroll - for (int group_6 = 0; group_6 < 8; group_6++) { - int row_offset_6 = load_warp_rank * 4 + group_6 * 16; - int raw_rows_6[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_6[0])), "=r"(*reinterpret_cast(&raw_rows_6[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_6[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_6[(0) + 3])) - : "r"(page_v_addr + (last_tile_1 & 1) * 128 * 4 + row_offset_6 * 4)); - int row0_6 = raw_rows_6[0]; - int row1_6 = raw_rows_6[1]; - int row2_6 = raw_rows_6[2]; - int row3_6 = raw_rows_6[3]; - int dst_6 = smem_v_addr + (unsigned int)(kv_stage_6 * 32768) + (unsigned int)(row_offset_6 * 64 * 2); - if (elect_sync()) { - { - tma_gather4_gmem2smem(dst_6, (&tmap_compressed_kv), kv_stage_6 * 128, row0_6, row1_6, row2_6, row3_6, kv_full_addr + (kv_stage_6) * 8); - tma_gather4_gmem2smem(dst_6 + 16384, (&tmap_compressed_kv), kv_stage_6 * 128 + 64, row0_6, row1_6, row2_6, row3_6, kv_full_addr + (kv_stage_6) * 8); - } - } - } - } - } - mbarrier_arrive(page_empty_addr + (page_v_stage) * 8); - } - mbarrier_wait(work_full_addr + (work_stage_6) * 8, _phase_work_full_6); - uint32_t _clc_valid_0 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_0) - : "r"(work_response_addr + work_stage_6 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_0 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_0) - : "r"(work_response_addr + work_stage_6 * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(work_empty_addr + (work_stage_6) * 8); - work_stage_6 += 1; - if (work_stage_6 == 2) { work_stage_6 = 0; _phase_work_full_6 ^= 1; } - unsigned int valid_6 = _clc_valid_0; - query_idx_6 = _clc_ctaid_0; - if (valid_6 == 0) { - break; - } - } - } - } - - // Cleanup -} - -} // extern "C" diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_ba9ed6a6233761682cc3_binding.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_ba9ed6a6233761682cc3_binding.cu deleted file mode 100644 index 5883041ded4..00000000000 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_ba9ed6a6233761682cc3_binding.cu +++ /dev/null @@ -1,555 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -// Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_ba9ed6a6233761682cc3'. -#include -#include -#include -#include - -#include "tvm_ffi_utils.h" - -#include -#include -#include -#include -#include -#include -#include - -extern "C" __global__ void kernel_cake_dsv4_ba9ed6a6233761682cc3(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); - - -namespace cake_host_shim_37348cf6921b17a8 { - -using tvm::ffi::Optional; -using tvm::ffi::TensorView; - -class ScopedCudaDevice { - public: - explicit ScopedCudaDevice(int device_id) { - cudaError_t error = cudaGetDevice(&previous_device_); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed before host-shim launch: cudaError=" - << static_cast(error); - if (previous_device_ != device_id) { - error = cudaSetDevice(device_id); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaSetDevice failed before host-shim launch for cuda:" - << device_id << ": cudaError=" << static_cast(error); - restore_ = true; - } - } - - ScopedCudaDevice(const ScopedCudaDevice&) = delete; - ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; - - ~ScopedCudaDevice() noexcept { - if (restore_) { - (void)cudaSetDevice(previous_device_); - } - } - - private: - int previous_device_ = -1; - bool restore_ = false; -}; - -inline int64_t CakeDeviceMultiprocessorCount(int device_id) { - constexpr int kMaxCachedCudaDevices = 64; - TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) - << "physical-SM-count cache does not cover cuda:" << device_id; - static std::atomic count_by_device[kMaxCachedCudaDevices]{}; - int cached = count_by_device[device_id].load(std::memory_order_acquire); - if (cached > 0) return cached; - - int count = 0; - cudaError_t error = cudaDeviceGetAttribute( - &count, cudaDevAttrMultiProcessorCount, device_id); - TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) - << "querying multiProcessorCount failed for cuda:" << device_id - << ": cudaError=" << static_cast(error) << ", count=" << count; - // Concurrent first launches may repeat the immutable device query, but all - // publication is atomic and every later hot-path lookup is one acquire load. - count_by_device[device_id].store(count, std::memory_order_release); - return count; -} - -inline void CheckCudaTensor(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) - << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; -} - -inline void CheckSameCudaDevice( - const TensorView& t, - const TensorView& reference, - const char* name, - const char* reference_name) { - TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) - << name << " must be on the same CUDA device as " << reference_name - << ": got cuda:" << t.device().device_id - << " versus cuda:" << reference.device().device_id; -} - -inline void CheckCurrentCudaDevice( - const TensorView& reference, - const char* reference_name) { - int current_device = -1; - cudaError_t error = cudaGetDevice(¤t_device); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed while validating " << reference_name - << ": cudaError=" << static_cast(error); - TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) - << "current CUDA device must match " << reference_name - << ": current=cuda:" << current_device - << ", tensor=cuda:" << reference.device().device_id; -} - -inline void CheckContiguous(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; -} - -inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { - DLDataType d = t.dtype(); - TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) - << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits - << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits - << ", lanes=" << (int)d.lanes << ")"; -} - -// A logical axis.outer(trailing) folds every source dim above the trailing -// dimensions. Shape products are independent of physical strides, so verify -// the leading dimensions form one dense row-major chain instead of inventing -// a "folded stride". The descriptor reads its exact adjacent physical step -// separately through stride[-(trailing + 1)]. -inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { - TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) - << name << " cannot fold leading dimensions above " << trailing - << " trailing dims from ndim=" << t.ndim(); - int outer_last = t.ndim() - trailing - 1; - if (outer_last <= 0) { - return; - } - int64_t step = t.stride(outer_last); - TVM_FFI_CHECK(step > 0, ValueError) - << name << " physical strides must be positive"; - int64_t expected = step; - for (int axis = outer_last - 1; axis >= 0; --axis) { - expected *= t.size(axis + 1); - if (t.size(axis) > 1) { - TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) - << name << " leading dims are not physically foldable above " << trailing - << " trailing dims: stride(" << axis << ")=" << t.stride(axis) - << ", expected " << expected; - } - } -} - -// 4D TMA descriptor for buffer 'tmap_q' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 3, ValueError) - << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - int64_t d2 = t.size(t.ndim() - 2); - int64_t d3 = t.size(t.ndim() - 3); - TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) - << "TMA source 'tmap_q' trailing dims must be positive"; - TVM_FFI_CHECK(d1 % 64 == 0, ValueError) - << "TMA source 'tmap_q' extent " << d1 - << " must divide exactly by " << 64; - uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(d3)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) - << "TMA box (64, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; - int64_t carrier_stride_0 = d1; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; - int64_t carrier_stride_1 = 64; - TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; - TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; - TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; - int64_t carrier_stride_2 = (d2 * d1); - TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; - TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; - TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; - uint64_t global_strides[3] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - (uint64_t)((carrier_stride_1 * 16) / 8), - (uint64_t)((carrier_stride_2 * 16) / 8), - }; - uint32_t box_dim[4] = {64u, 64u, 1u, 1u}; - uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_swa_k' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_swa_k(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_swa_k' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_swa_k' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_swa_k' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_swa_k"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_swa_k' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_swa_k' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_k'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_swa_k') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_compressed_k' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_compressed_k(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_compressed_k' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_compressed_k' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_compressed_k' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_compressed_k"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_compressed_k' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_compressed_k' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_k'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_k') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_swa_v' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_swa_v(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_swa_v' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_swa_v' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_swa_v' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_swa_v"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_swa_v' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_swa_v' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_v'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_swa_v') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_compressed_v' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_compressed_v(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_compressed_v' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_compressed_v' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_compressed_v' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_compressed_v"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_compressed_v' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_compressed_v' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_v'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_v') failed: CUresult=" << (int)r; - return tm; -} - -void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_compressed_k, TensorView arg_tmap_swa_v, TensorView arg_tmap_compressed_v, TensorView arg_O, TensorView arg_partial_lse, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_num_query_tokens, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, int64_t grid_x, int64_t grid_y, int64_t grid_z) { - DLDevice dev = arg_tmap_q.device(); - ScopedCudaDevice device_guard(dev.device_id); - CheckCudaTensor(arg_tmap_q, "tmap_q"); - CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); - CheckContiguous(arg_tmap_q, "tmap_q"); - CheckCudaTensor(arg_tmap_swa_k, "tmap_swa_k"); - CheckDtype(arg_tmap_swa_k, "tmap_swa_k", 4, 16, 1); - CheckCudaTensor(arg_tmap_compressed_k, "tmap_compressed_k"); - CheckDtype(arg_tmap_compressed_k, "tmap_compressed_k", 4, 16, 1); - CheckCudaTensor(arg_tmap_swa_v, "tmap_swa_v"); - CheckDtype(arg_tmap_swa_v, "tmap_swa_v", 4, 16, 1); - CheckCudaTensor(arg_tmap_compressed_v, "tmap_compressed_v"); - CheckDtype(arg_tmap_compressed_v, "tmap_compressed_v", 4, 16, 1); - CheckCudaTensor(arg_O, "O"); - CheckDtype(arg_O, "O", 4, 16, 1); - CheckContiguous(arg_O, "O"); - CheckCudaTensor(arg_partial_lse, "partial_lse"); - CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); - CheckContiguous(arg_partial_lse, "partial_lse"); - CheckCudaTensor(arg_swa_indices, "swa_indices"); - CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); - CheckContiguous(arg_swa_indices, "swa_indices"); - CheckCudaTensor(arg_compressed_indices, "compressed_indices"); - CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); - CheckContiguous(arg_compressed_indices, "compressed_indices"); - CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); - CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckCudaTensor(arg_sinks, "sinks"); - CheckDtype(arg_sinks, "sinks", 2, 32, 1); - CheckContiguous(arg_sinks, "sinks"); - CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); - CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); - CheckContiguous(arg_bmm1_scale, "bmm1_scale"); - CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); - CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); - CheckContiguous(arg_bmm2_scale, "bmm2_scale"); - TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) - << "scalar 'num_heads' value " << arg_num_heads - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_num_query_tokens >= -2147483648LL && arg_num_query_tokens <= 2147483647LL, ValueError) - << "scalar 'num_query_tokens' value " << arg_num_query_tokens - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) - << "scalar 'swa_index_stride' value " << arg_swa_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) - << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) - << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) - << "scalar 'sparse_topk' value " << arg_sparse_topk - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) - << "scalar 'has_sinks' value " << arg_has_sinks - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) - << "scalar 'total_work_items' value " << arg_total_work_items - << " is outside i32 range [-2147483648, 2147483647]"; - CheckSameCudaDevice(arg_tmap_swa_k, arg_tmap_q, "tmap_swa_k", "tmap_q"); - CheckSameCudaDevice(arg_tmap_compressed_k, arg_tmap_q, "tmap_compressed_k", "tmap_q"); - CheckSameCudaDevice(arg_tmap_swa_v, arg_tmap_q, "tmap_swa_v", "tmap_q"); - CheckSameCudaDevice(arg_tmap_compressed_v, arg_tmap_q, "tmap_compressed_v", "tmap_q"); - CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); - CheckSameCudaDevice(arg_partial_lse, arg_tmap_q, "partial_lse", "tmap_q"); - CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); - CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); - CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); - CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); - CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); - CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); - TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) - << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y - << ", " << grid_z << ")"; - TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) - << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z - << ") must be divisible by cluster dims (2, 1, 1)"; - - cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); - CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); - CUtensorMap p_tmap_swa_k = EncodeTma_tmap_swa_k(arg_tmap_swa_k); - CUtensorMap p_tmap_compressed_k = EncodeTma_tmap_compressed_k(arg_tmap_compressed_k); - CUtensorMap p_tmap_swa_v = EncodeTma_tmap_swa_v(arg_tmap_swa_v); - CUtensorMap p_tmap_compressed_v = EncodeTma_tmap_compressed_v(arg_tmap_compressed_v); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); - float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); - int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); - int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); - int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); - float* p_sinks = static_cast(arg_sinks.data_ptr()); - float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); - float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); - int32_t v_num_heads = (int32_t)arg_num_heads; - int32_t v_num_query_tokens = (int32_t)arg_num_query_tokens; - int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; - int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; - int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; - int32_t v_sparse_topk = (int32_t)arg_sparse_topk; - int32_t v_has_sinks = (int32_t)arg_has_sinks; - int32_t v_total_work_items = (int32_t)arg_total_work_items; - void* kargs[] = {&p_tmap_q, &p_tmap_swa_k, &p_tmap_compressed_k, &p_tmap_swa_v, &p_tmap_compressed_v, &p_O, &p_partial_lse, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_num_query_tokens, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_has_sinks, &v_total_work_items}; - - static std::mutex smem_mu; - static auto* smem_status_by_device = - new std::unordered_map(); - cudaError_t smem_status = cudaSuccess; - { - std::lock_guard lock(smem_mu); - auto it = smem_status_by_device->find(dev.device_id); - if (it == smem_status_by_device->end()) { - int current_device = -1; - smem_status = cudaGetDevice(¤t_device); - if (smem_status == cudaSuccess && current_device != dev.device_id) { - smem_status = cudaErrorInvalidDevice; - } - if (smem_status == cudaSuccess) { - smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_ba9ed6a6233761682cc3), - cudaFuncAttributeMaxDynamicSharedMemorySize, - 231808); - } - smem_status_by_device->emplace(dev.device_id, smem_status); - } else { - smem_status = it->second; - } - } - TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_ba9ed6a6233761682cc3 failed: " - << cudaGetErrorString(smem_status); - dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); - dim3 block(512u, 1u, 1u); - - // Extended direct-source launch for a linked CUDA kernel symbol. - cudaLaunchAttribute attrs[2]{}; - int n = 0; - attrs[n].id = cudaLaunchAttributeClusterDimension; - attrs[n].val.clusterDim.x = 2u; - attrs[n].val.clusterDim.y = 1u; - attrs[n].val.clusterDim.z = 1u; - ++n; - attrs[n].id = cudaLaunchAttributeClusterSchedulingPolicyPreference; - attrs[n].val.clusterSchedulingPolicyPreference = cudaClusterSchedulingPolicySpread; - ++n; - cudaLaunchConfig_t config{}; - config.gridDim = grid; - config.blockDim = block; - config.dynamicSmemBytes = 231808u; - config.stream = stream; - config.attrs = attrs; - config.numAttrs = n; - cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_ba9ed6a6233761682cc3), kargs); - TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_ba9ed6a6233761682cc3 failed: " - << cudaGetErrorString(launch_status); -} - -} // namespace cake_host_shim_37348cf6921b17a8 - -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_37348cf6921b17a8::Run); diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_ba9ed6a6233761682cc3_kernel.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_ba9ed6a6233761682cc3_kernel.cu deleted file mode 100644 index 5677c4f50cd..00000000000 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_ba9ed6a6233761682cc3_kernel.cu +++ /dev/null @@ -1,2367 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -typedef signed char int8_t; -typedef unsigned char uint8_t; -typedef unsigned short uint16_t; -typedef unsigned int uint32_t; -#if defined(__CUDACC_RTC__) -typedef unsigned long long uint64_t; -#else -typedef unsigned long uint64_t; -#endif -static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); -typedef signed int int32_t; -typedef short int int16_t; -struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; -struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; -static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); -static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); - -#if defined(__CUDACC_RTC__) -typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; -#else -#include -#endif - -static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); -static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); -#include -#include - -__device__ __forceinline__ int make_warp_uniform(int x) { - int result; - asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" - : "=r"(result) : "r"(x)); - return result; -} - -#define CAKE_INF CUDART_INF_F -#define TMEM_NCOLS 512 -#define TMEM_TMEM_SCRATCH_OFFSET 0 -#define NUM_K_PIPE_STAGES 4 -#define NUM_V_PIPE_STAGES 2 -#define SMEM_SMEM_Q_OFF 0 -#define SMEM_SMEM_Q_STAGE_BYTES 8192 -#define SMEM_SMEM_Q_STRIDE 8192 -#define SMEM_SMEM_K_OFF 65536 -#define SMEM_SMEM_K_STAGE_BYTES 16384 -#define SMEM_SMEM_K_STRIDE 16384 -#define SMEM_SMEM_V_LO_OFF 131072 -#define SMEM_SMEM_V_LO_STAGE_BYTES 16384 -#define SMEM_SMEM_V_LO_STRIDE 32768 -#define SMEM_SMEM_V_HI_OFF 147456 -#define SMEM_SMEM_V_HI_STAGE_BYTES 16384 -#define SMEM_SMEM_V_HI_STRIDE 32768 -#define SMEM_SMEM_P_OFF 196608 -#define SMEM_SMEM_P_STAGE_BYTES 8192 -#define SMEM_SMEM_P_STRIDE 8192 -#define SMEM_SMEM_INDICES_OFF 229376 -#define SMEM_SMEM_INDICES_STAGE_BYTES 1024 -#define SMEM_SMEM_INDICES_STRIDE 1024 -#define SMEM_SMEM_SOFTMAX_EXCHANGE_OFF 230400 -#define SMEM_SMEM_SOFTMAX_EXCHANGE_STAGE_BYTES 512 -#define SMEM_SMEM_SOFTMAX_EXCHANGE_STRIDE 512 -#define SMEM_SMEM_EPILOGUE_EXCHANGE_OFF 230912 -#define SMEM_SMEM_EPILOGUE_EXCHANGE_STAGE_BYTES 512 -#define SMEM_SMEM_EPILOGUE_EXCHANGE_STRIDE 512 -#define SMEM_SMEM_INDEX_VALID_OFF 231424 -#define SMEM_SMEM_INDEX_VALID_STAGE_BYTES 64 -#define SMEM_SMEM_INDEX_VALID_STRIDE 64 -#define SMEM_TOTAL 231808 -#define THREADS 512 - -#include - -__device__ __forceinline__ uint32_t elect_sync() { - uint32_t pred = 0; - asm volatile( - "{\n\t" - ".reg .pred %%px;\n\t" - "elect.sync _|%%px, %1;\n\t" - "@%%px mov.s32 %0, 1;\n\t" - "}\n" - : "+r"(pred) - : "r"(0xFFFFFFFF)); - return pred; -} - - -__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { - asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" - :: "r"(mbar_addr), "r"(count) : "memory"); -} - -__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { - asm volatile("mbarrier.init.b64 [%0], %1;" - :: "l"(mbar_addr), "r"(count) : "memory"); -} - - -__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - - -// CTA-local pipelines have short, resident producer/consumer edges. Omitting -// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the -// explicit loop still makes this helper blocking until acquire succeeds. -__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE;\n\t" - "bra.uni LAB_WAIT;\n\t" - "DONE:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Source-faithful relaxed CTA wait used only by a typed protocol that does -// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. -__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, 10000000;\n\t" - "@P1 bra.uni DONE_RELAXED;\n\t" - "bra.uni LAB_WAIT_RELAXED;\n\t" - "DONE_RELAXED:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Exact source ports may request the PTX suspendTimeHint operand explicitly. -// The hint is expressed in nanoseconds and is kept separate from the canonical -// no-hint CTA helper so unrelated schedules retain their existing retry path. -__device__ __forceinline__ void mbarrier_wait_suspend( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_SUSPEND:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_SUSPEND;\n\t" - "bra.uni LAB_WAIT_SUSPEND;\n\t" - "DONE_SUSPEND:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE_CLUSTER;\n\t" - "bra.uni LAB_WAIT_CLUSTER;\n\t" - "DONE_CLUSTER:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - ".reg .u32 WAIT_ADDR;\n\t" - "mov.u32 WAIT_ADDR, %0;\n\t" - "LAB_WAIT_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [WAIT_ADDR], %1, %2;\n\t" - "@P1 bra.uni DONE_HINT;\n\t" - "bra.uni LAB_WAIT_HINT;\n\t" - "DONE_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -// Exact unqualified CTA wait used by source schedules whose PTX intentionally -// omits the acquire qualifier while retaining a typed suspendTimeHint operand. -__device__ __forceinline__ void mbarrier_wait_relaxed_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED_HINT:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra DONE_RELAXED_HINT;\n\t" - "bra LAB_WAIT_RELAXED_HINT;\n\t" - "DONE_RELAXED_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); -} - -__device__ __forceinline__ void mbarrier_wait_cluster_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" - "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" - "DONE_CLUSTER_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_suspend( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait_cluster(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); - } -} - - -__device__ __forceinline__ void tcgen05_mma_f16_cta2( - int taddr, uint64_t a_desc, uint64_t b_desc, - uint32_t i_desc, int enable_input_d) { - asm volatile( - "{\n\t" - ".reg .pred p;\n\t" - ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - "setp.ne.b32 p, %4, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" - "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "tcgen05.mma.cta_group::2.kind::f16 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" - "}\n" - :: "r"(taddr), "l"(a_desc), "l"(b_desc), - "r"(i_desc), "r"(enable_input_d) - : "memory"); -} - - -__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { - return (x & 0x3FFFFULL) >> 4ULL; -} - - -union MmaSmemDesc { - uint64_t u64; - uint32_t u32[2]; -}; - -__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { - MmaSmemDesc tmp; - tmp.u64 = smem_desc; - tmp.u32[0] += offset; - smem_desc = tmp.u64; -} - - -__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { - asm volatile( - "{\n\t" - ".reg .pred leader;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" - ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" - "}\n" - :: "r"(mbar_addr), "h"(cta_mask) : "memory"); -} - - -__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { - asm volatile( - "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" - :: "r"(mbar_addr), "r"(bytes) : "memory"); -} - - -__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7," - " %8, %9, %10, %11, %12, %13, %14, %15," - " %16, %17, %18, %19, %20, %21, %22, %23," - " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), - "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), - "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), - "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), - "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), - "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), - "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x16.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7," - " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), - "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), - "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ void tmem_ld_x4(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x4.b32" - " {%0, %1, %2, %3}, [%4];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { - asm volatile( - "tcgen05.st.sync.aligned.32x32b.x16.b32" - " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," - " %9, %10, %11, %12, %13, %14, %15, %16};" - :: "r"(tmem_addr), - "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), - "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), - "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), - "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); -} - - -__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { - asm volatile( - "tcgen05.st.sync.aligned.32x32b.x4.b32" - " [%0], {%1, %2, %3, %4};" - :: "r"(tmem_addr), - "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); -} - - -__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { - uint32_t addr; - asm("{\n\t" - ".reg .u64 u64addr;\n\t" - "cvta.to.shared.u64 u64addr, %1;\n\t" - "cvt.u32.u64 %0, u64addr;\n\t" - "}\n" : "=r"(addr) : "l"(ptr)); - return addr; -} - - -__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { - uint32_t remote; - asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" - : "=r"(remote) : "r"(local_addr), "r"(rank)); - return remote; -} - - -__device__ __forceinline__ float approx_exp2(float x) { - float y; - asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float approx_rcp(float x) { - float y; - asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float max_noftz(float a, float b) { - float c; - asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); - return c; -} - - -__device__ __forceinline__ float warp_reduce_max(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); - return val; -} - - -__device__ __forceinline__ float warp_reduce_sum(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val += __shfl_xor_sync(0xFFFFFFFF, val, offset); - return val; -} - - -__device__ __forceinline__ float row_max_reduce(float2 acc) { - return max_noftz(acc.x, acc.y); -} - - -__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { - #pragma unroll - for (int j = 0; j < 16; j++) { - if (j % 2 == 0) - acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); - else - acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); - } -} - - -__device__ __forceinline__ float2 ex2_emulation_f32x2_value(float2 value) { - const float c0 = 1.0f, c1 = 0.695146143436431884765625f; - const float c2 = 0.227564394474029541015625f, c3 = 0.077119089663028717041015625f; - const float magic = 12582912.0f; - float x0 = max_noftz(value.x, -127.0f), x1 = max_noftz(value.y, -127.0f); - float2 xc2 = make_float2(x0, x1), magic2 = make_float2(magic, magic); - float2 xr2; - asm("add.rm.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xr2) - : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&magic2)); - float2 c3_2 = make_float2(c3, c3), c2_2 = make_float2(c2, c2); - float2 c1_2 = make_float2(c1, c1), c0_2 = make_float2(c0, c0); - float2 xrb2, xfrac2; - asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xrb2) - : "l"(*(unsigned long long*)&xr2), "l"(*(unsigned long long*)&magic2)); - asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xfrac2) - : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&xrb2)); - float2 poly2; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&c3_2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c2_2)); - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c1_2)); - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c0_2)); - int x0r_i, x1r_i, p0_i, p1_i; - asm("mov.b64 {%0, %1}, %2;" : "=r"(x0r_i), "=r"(x1r_i) : "l"(*(unsigned long long*)&xr2)); - asm("mov.b64 {%0, %1}, %2;" : "=r"(p0_i), "=r"(p1_i) : "l"(*(unsigned long long*)&poly2)); - float r0, r1; - asm("mov.b32 %0, %1;" : "=f"(r0) : "r"((x0r_i << 23) + p0_i)); - asm("mov.b32 %0, %1;" : "=f"(r1) : "r"((x1r_i << 23) + p1_i)); - return make_float2(r0, r1); -} - -__device__ __forceinline__ void ex2_emulation_f32x2(float* x0_ptr, float* x1_ptr) { - float2 result = ex2_emulation_f32x2_value(make_float2(*x0_ptr, *x1_ptr)); - *x0_ptr = result.x; *x1_ptr = result.y; -} - -__device__ __forceinline__ void softmax_frag_exp2_cast( - float* sv, uint32_t* pv, int use_emu) -{ - #pragma unroll - for (int j = 0; j < 16; j++) { - if (use_emu && j >= 12) - ex2_emulation_f32x2(&sv[j*2], &sv[j*2+1]); - else { - sv[j*2] = approx_exp2(sv[j*2]); - sv[j*2+1] = approx_exp2(sv[j*2+1]); - } - } - #pragma unroll - for (int j = 0; j < 16; j++) { - __nv_bfloat162 bf = __float22bfloat162_rn({sv[j*2], sv[j*2+1]}); - pv[j] = reinterpret_cast(bf); - } -} - - - -__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { - const float2* sv2 = reinterpret_cast(sv); - #pragma unroll - for (int j = 0; j < 16; j++) { - asm("add.f32x2 %0, %1, %2;" - : "+l"(reinterpret_cast(*acc)) - : "l"(reinterpret_cast(*acc)), - "l"(reinterpret_cast(sv2[j]))); - } -} - - -__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { - asm("mul.rn.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { - asm("mul.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { - asm("add.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { - asm("add.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { - asm("sub.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { - asm("sub.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("add.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { - float2 r; - asm("sub.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("sub.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ void fma_scale_x32( - float* sv, const float2* scale2, const float2* neg_max2) -{ - float2* sv_2 = reinterpret_cast(sv); - #pragma unroll - for (int j = 0; j < 16; j++) - fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); -} - -__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("mul.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - - -// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) - -__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("add.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("add.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("add.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("add.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("add.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("add.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("add.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - - -__device__ __forceinline__ void fence_async_shared() { - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); -} - - -__device__ __forceinline__ uint64_t make_smem_desc(int addr) { - const int SBO = 1024; - return desc_encode(addr) - | (desc_encode(SBO) << 32ULL) - | (1ULL << 46ULL) - | (2ULL << 61ULL); -} - - -__device__ __forceinline__ void tma_2d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3}], [%4];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_4d_gmem2smem_cta2( - int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.4d.shared::cluster.global" - ".mbarrier::complete_tx::bytes.cta_group::2" - " [%0], [%1, {%2, %3, %4, %5}], [%6];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr) { - // Canonical .shared::cta form for non-multicast gather4, matching - // trtllm-gen / cuda_ptx and the PTX ISA qualifier order - // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, - // .shared::cluster is reserved for the multicast variant (ctaMask). - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem_cta2( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr) { - // Canonical .shared::cta form; see tma_gather4_gmem2smem above. - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" - ".mbarrier::complete_tx::bytes.cta_group::2" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem_mc( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr, unsigned short cta_mask) { - // Multicast variant: the PTX grammar ties the .shared::cluster - // destination to .multicast::cluster + ctaMask (cf. cuda_ptx / - // SM100_TMA_LOAD_MULTICAST_2D_GATHER4 in CUTLASS). - asm volatile( - "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" - ".mbarrier::complete_tx::bytes.multicast::cluster" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr), "h"(cta_mask) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem_mc_cta2( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr, unsigned short cta_mask) { - // Multicast + cta_group::2 variant; see tma_gather4_gmem2smem_mc. - asm volatile( - "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" - ".mbarrier::complete_tx::bytes.multicast::cluster.cta_group::2" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr), "h"(cta_mask) : "memory"); -} - - -__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { - asm volatile( - "{\n\t" - ".reg .b16 lo, hi;\n\t" - "mov.b32 {lo, hi}, %1;\n\t" - "tcgen05.commit.cta_group::2.mbarrier::arrive::one" - ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" - "}\n" - :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); -} - -extern "C" { - -__global__ __launch_bounds__(512, 1) __cluster_dims__(2,1,1) void -kernel_cake_dsv4_ba9ed6a6233761682cc3(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items) -{ - const int tid = threadIdx.x; - const int warp = make_warp_uniform(tid / 32); - const int lane = tid % 32; - - extern __shared__ __align__(1024) char smem_raw[]; - int smem; - asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); - smem = make_warp_uniform(smem); - - const int mbar_base = smem + 231488; - #define q_full_addr (mbar_base + 0) - #define q_empty_addr (mbar_base + 8) - #define k_full_addr (mbar_base + 16) - #define k_empty_addr (mbar_base + 48) - #define v_full_addr (mbar_base + 80) - #define v_empty_addr (mbar_base + 96) - #define s_full_addr (mbar_base + 112) - #define s_empty_addr (mbar_base + 128) - #define p_full_addr (mbar_base + 144) - #define p_empty_addr (mbar_base + 160) - #define o_empty_addr (mbar_base + 176) - #define stats_addr (mbar_base + 184) - #define stats_empty_addr (mbar_base + 200) - #define o_full_addr (mbar_base + 216) - #define o_first_slice_seeded_addr (mbar_base + 224) - #define tmem_dealloc_peer_addr (mbar_base + 232) - #define index_valid_full_addr (mbar_base + 240) - #define index_valid_empty_addr (mbar_base + 256) - - const int bid = blockIdx.x; - const int num_bids = gridDim.x; - const unsigned int clusters_x = gridDim.x / 2; - const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; - const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; - - int cta_rank; - asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); - - // Kernel setup ops - __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 0); - const int smem_q_addr = smem + 0; - __nv_bfloat16* smem_k = reinterpret_cast<__nv_bfloat16*>(smem_raw + 65536); - const int smem_k_addr = smem + 65536; - __nv_bfloat16* smem_v_lo = reinterpret_cast<__nv_bfloat16*>(smem_raw + 131072); - const int smem_v_lo_addr = smem + 131072; - __nv_bfloat16* smem_v_hi = reinterpret_cast<__nv_bfloat16*>(smem_raw + 147456); - const int smem_v_hi_addr = smem + 147456; - __nv_bfloat16* smem_p = reinterpret_cast<__nv_bfloat16*>(smem_raw + 196608); - const int smem_p_addr = smem + 196608; - int* smem_indices = reinterpret_cast(smem_raw + 229376); - const int smem_indices_addr = smem + 229376; - float* smem_softmax_exchange = reinterpret_cast(smem_raw + 230400); - const int smem_softmax_exchange_addr = smem + 230400; - float* smem_epilogue_exchange = reinterpret_cast(smem_raw + 230912); - const int smem_epilogue_exchange_addr = smem + 230912; - unsigned int* smem_index_valid = reinterpret_cast(smem_raw + 231424); - const int smem_index_valid_addr = smem + 231424; - asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); - asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); - - // Mbarrier init (18 pipeline groups, 0 ordered-sequence groups, 34 barriers) - // Mbarriers at smem_raw[231488..231760) - - if (warp == 0) { - uint32_t leader = elect_sync(); - if (leader) { - // q_full: 1 barriers, init_count=1 - mbarrier_init(smem + 231488, 1); - // q_empty: 1 barriers, init_count=1 - mbarrier_init(smem + 231496, 1); - // --- pipeline 'k_pipe' --- - // k_full: 4 barriers, init_count=1 - mbarrier_init(smem + 231504, 1); - mbarrier_init(smem + 231512, 1); - mbarrier_init(smem + 231520, 1); - mbarrier_init(smem + 231528, 1); - // k_empty: 4 barriers, init_count=1 - mbarrier_init(smem + 231536, 1); - mbarrier_init(smem + 231544, 1); - mbarrier_init(smem + 231552, 1); - mbarrier_init(smem + 231560, 1); - // --- pipeline 'v_pipe' --- - // v_full: 2 barriers, init_count=1 - mbarrier_init(smem + 231568, 1); - mbarrier_init(smem + 231576, 1); - // v_empty: 2 barriers, init_count=1 - mbarrier_init(smem + 231584, 1); - mbarrier_init(smem + 231592, 1); - // s_full: 2 barriers, init_count=1 - mbarrier_init(smem + 231600, 1); - mbarrier_init(smem + 231608, 1); - // s_empty: 2 barriers, init_count=256 - mbarrier_init(smem + 231616, 256); - mbarrier_init(smem + 231624, 256); - // p_full: 2 barriers, init_count=256 - mbarrier_init(smem + 231632, 256); - mbarrier_init(smem + 231640, 256); - // p_empty: 2 barriers, init_count=1 - mbarrier_init(smem + 231648, 1); - mbarrier_init(smem + 231656, 1); - // o_empty: 1 barriers, init_count=256 - mbarrier_init(smem + 231664, 256); - // stats: 2 barriers, init_count=128 - mbarrier_init(smem + 231672, 128); - mbarrier_init(smem + 231680, 128); - // stats_empty: 2 barriers, init_count=128 - mbarrier_init(smem + 231688, 128); - mbarrier_init(smem + 231696, 128); - // o_full: 1 barriers, init_count=1 - mbarrier_init(smem + 231704, 1); - // o_first_slice_seeded: 1 barriers, init_count=256 - mbarrier_init(smem + 231712, 256); - // tmem_dealloc_peer: 1 barriers, init_count=32 - mbarrier_init(smem + 231720, 32); - // index_valid_full: 2 barriers, init_count=32 - mbarrier_init(smem + 231728, 32); - mbarrier_init(smem + 231736, 32); - // index_valid_empty: 2 barriers, init_count=128 - mbarrier_init(smem + 231744, 128); - mbarrier_init(smem + 231752, 128); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - - asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); - asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); - - // TMEM alloc (512 columns, 512 used) - volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 231760); - if (warp == 0) { - int _tmem_hold = smem + 231760; - asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); - __syncwarp(); - asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); - } - - __syncthreads(); - asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); - asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); - - const int taddr = tmem_addr_storage[0]; - - // Kernel post-init ops - const int tmem_tmem_scratch = taddr; - - // ---- Ordered hardware-WG register redistribution ---- - // Dec phase frees registers before any WG attempts inc. - if (warp >= 8 && warp <= 15) { - asm volatile("setmaxnreg.dec.sync.aligned.u32 64;"); - } - - // ---- Role: softmax_wg ---- - if (warp <= 3) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); - { // softmax_wg_main - const int softmax_dummy = 0; - float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; - const int head_row_base = warp % 2 * 32; - const int score_row_base = warp % 4 * 32; - const int n_half = warp % 4 / 2; - const int my_row = head_row_base + lane; - const int exchange_idx = n_half * 64 + my_row; - int tile_cursor = 0; - int valid_iter = 0; - float seed_zero[4]; - #pragma unroll - for (int seed_i = 0; seed_i < 4; seed_i++) { - seed_zero[seed_i] = 0.0f; - } - #pragma unroll - for (int seed_stage = 0; seed_stage < 2; seed_stage++) { - #pragma unroll - for (int seed_col = 128 + seed_stage * 128; seed_col < 128 + (seed_stage + 1) * 128; seed_col += 4) { - int seed_addr = taddr + (unsigned int)seed_col + (unsigned int)(score_row_base << 16); - tmem_st_x4_f32(seed_addr, seed_zero); - } - } - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(o_first_slice_seeded_addr); - int peer_rank = cta_rank ^ 1; - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(o_first_slice_seeded_addr), "r"(peer_rank) : "memory"); - #pragma unroll 1 - for (unsigned int work_idx = cluster_id; work_idx < total_work_items; work_idx += num_clusters) { - int split_idx = work_idx % (unsigned int)(((1) ? 5 : 1)); - int query_idx = work_idx / (unsigned int)(((1) ? 5 : 1)); - int _max_0 = ((sparse_topk_lens[query_idx] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx] + sparse_topk_lens_offset) : (0)); - int _min_2 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); - int active_topk = _min_2; - int owner_capacity = ((sparse_topk + 128 - 1) / 128 + ((1) ? 5 : 1) - 1) / ((1) ? 5 : 1); - int first_owner_tile = work_idx % (unsigned int)(((1) ? 5 : 1)) * (unsigned int)owner_capacity; - int _min_3 = ((owner_capacity) < ((sparse_topk + 128 - 1) / 128 - first_owner_tile) ? (owner_capacity) : ((sparse_topk + 128 - 1) / 128 - first_owner_tile)); - int num_kv_tiles = _min_3; - int first_kv_tile = first_owner_tile; - int head_idx = cta_rank * 64 + my_row; - float row_max_val = -CAKE_INF; - float row_sum_val = 0.0f; - if (has_sinks != 0 && head_idx < num_heads && split_idx == 0) { - row_max_val = sinks[head_idx] * 1.4426950408889634f / softmax_scale_log2; - if (n_half == 0) { - row_sum_val = 1.0f; - } - } - int valid_slot = valid_iter & 1; - int valid_full_phase = valid_iter >> 1 & 1; - mbarrier_wait(index_valid_full_addr + (valid_slot) * 8, valid_full_phase); - int valid_word_base = valid_slot * 8 + n_half * 2; - #pragma unroll 1 - for (int tile = 0; tile < num_kv_tiles; tile++) { - int pipeline_tile = tile_cursor + tile; - int score_phase = pipeline_tile & 1; - int s_full_phase = pipeline_tile >> 1 & 1; - mbarrier_wait(s_full_addr + (score_phase) * 8, s_full_phase); - int score_offset = ((score_phase != 0) ? 64 : 0); - int score_base = taddr + (unsigned int)score_offset + (unsigned int)(score_row_base << 16); - float _tmem_load_0[64]; - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]) - : "r"(score_base)); - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]) - : "r"(score_base + 32)); - int valid_cols = active_topk - (first_kv_tile + tile) * 128 - n_half * 64; - if (valid_cols < 0) { - valid_cols = 0; - } - if (valid_cols > 64) { - valid_cols = 64; - } - unsigned int keep_lo = smem_index_valid[valid_word_base + tile * 4]; - unsigned int keep_hi = smem_index_valid[valid_word_base + tile * 4 + 1]; - uint32_t _thresh_mask_0; - { - int _lim_0 = valid_cols; - if (_lim_0 <= 0) { _thresh_mask_0 = 0u; } - else if (_lim_0 >= 32) { _thresh_mask_0 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_thresh_mask_0) : "r"(_lim_0)); - } - } - uint32_t _thresh_mask_1; - { - int _lim_1 = valid_cols - 32; - if (_lim_1 <= 0) { _thresh_mask_1 = 0u; } - else if (_lim_1 >= 32) { _thresh_mask_1 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_thresh_mask_1) : "r"(_lim_1)); - } - } - if (!(_thresh_mask_0 & keep_lo & (1u << 0))) _tmem_load_0[0] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 1))) _tmem_load_0[1] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 2))) _tmem_load_0[2] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 3))) _tmem_load_0[3] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 4))) _tmem_load_0[4] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 5))) _tmem_load_0[5] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 6))) _tmem_load_0[6] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 7))) _tmem_load_0[7] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 8))) _tmem_load_0[8] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 9))) _tmem_load_0[9] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 10))) _tmem_load_0[10] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 11))) _tmem_load_0[11] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 12))) _tmem_load_0[12] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 13))) _tmem_load_0[13] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 14))) _tmem_load_0[14] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 15))) _tmem_load_0[15] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 16))) _tmem_load_0[16] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 17))) _tmem_load_0[17] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 18))) _tmem_load_0[18] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 19))) _tmem_load_0[19] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 20))) _tmem_load_0[20] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 21))) _tmem_load_0[21] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 22))) _tmem_load_0[22] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 23))) _tmem_load_0[23] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 24))) _tmem_load_0[24] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 25))) _tmem_load_0[25] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 26))) _tmem_load_0[26] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 27))) _tmem_load_0[27] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 28))) _tmem_load_0[28] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 29))) _tmem_load_0[29] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 30))) _tmem_load_0[30] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 31))) _tmem_load_0[31] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 0))) _tmem_load_0[32] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 1))) _tmem_load_0[33] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 2))) _tmem_load_0[34] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 3))) _tmem_load_0[35] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 4))) _tmem_load_0[36] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 5))) _tmem_load_0[37] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 6))) _tmem_load_0[38] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 7))) _tmem_load_0[39] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 8))) _tmem_load_0[40] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 9))) _tmem_load_0[41] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 10))) _tmem_load_0[42] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 11))) _tmem_load_0[43] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 12))) _tmem_load_0[44] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 13))) _tmem_load_0[45] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 14))) _tmem_load_0[46] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 15))) _tmem_load_0[47] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 16))) _tmem_load_0[48] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 17))) _tmem_load_0[49] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 18))) _tmem_load_0[50] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 19))) _tmem_load_0[51] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 20))) _tmem_load_0[52] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 21))) _tmem_load_0[53] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 22))) _tmem_load_0[54] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 23))) _tmem_load_0[55] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 24))) _tmem_load_0[56] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 25))) _tmem_load_0[57] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 26))) _tmem_load_0[58] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 27))) _tmem_load_0[59] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 28))) _tmem_load_0[60] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 29))) _tmem_load_0[61] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 30))) _tmem_load_0[62] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 31))) _tmem_load_0[63] = -CAKE_INF; - float2 _reg_reduce_max2_2 = {-CAKE_INF, -CAKE_INF}; - row_max_x32_accum(&_tmem_load_0[0], _reg_reduce_max2_2); - row_max_x32_accum(&_tmem_load_0[32], _reg_reduce_max2_2); - float _tmem_load_0_max = row_max_reduce(_reg_reduce_max2_2); - float _max_1 = max_noftz(_tmem_load_0_max, row_max_val); - float new_max = _max_1; - smem_softmax_exchange[exchange_idx] = new_max; - asm volatile("barrier.sync 3, 128;" ::: "memory"); - float _max_2 = max_noftz(new_max, smem_softmax_exchange[exchange_idx ^ 64]); - new_max = _max_2; - asm volatile("barrier.sync 3, 128;" ::: "memory"); - float no_correction = (((new_max - row_max_val) * softmax_scale_log2 <= 0.0f) ? 1.0f : 0.0f); - float delta = softmax_scale_log2 * (row_max_val - new_max); - float _exp2_0 = approx_exp2(delta); - float exp_delta = _exp2_0; - float acc_scale = ((row_max_val > -CAKE_INF) ? exp_delta : 1.0f); - row_max_val = new_max; - float safe_max = ((new_max == -CAKE_INF) ? 0.0f : new_max); - float max_scaled = safe_max * softmax_scale_log2; - const float2 _fma_b2_3 = {softmax_scale_log2, softmax_scale_log2}; - const float2 _fma_c2_4 = {-max_scaled, -max_scaled}; - #pragma unroll - for (int _lf = 0; _lf < 32; _lf++) - fma_f32x2_inplace(&reinterpret_cast(_tmem_load_0)[_lf], _fma_b2_3, _fma_c2_4); - #pragma unroll - for (int _le = 0; _le < 64; _le++) { - _tmem_load_0[_le] = approx_exp2(_tmem_load_0[_le]); - } - float2 _reg_reduce_sum2_5 = make_float2(0.0f, 0.0f); - softmax_block_sum(&_tmem_load_0[0], &_reg_reduce_sum2_5); - softmax_block_sum(&_tmem_load_0[32], &_reg_reduce_sum2_5); - float _tmem_load_0_sum = _reg_reduce_sum2_5.x + _reg_reduce_sum2_5.y; - row_sum_val = row_sum_val * acc_scale + _tmem_load_0_sum; - int stats_empty_phase = pipeline_tile >> 1 & 1 ^ 1; - mbarrier_wait(stats_empty_addr + (score_phase) * 8, stats_empty_phase); - float meta[4]; - meta[0] = row_sum_val; - meta[1] = row_max_val; - meta[2] = acc_scale; - meta[3] = no_correction; - int meta_addr = taddr + 384 + (unsigned int)(score_phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(head_row_base << 16); - tmem_st_x4_f32(meta_addr, meta); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - mbarrier_arrive(stats_addr + (score_phase) * 8); - int p_stage = pipeline_tile % 2; - int p_empty_phase = pipeline_tile / 2 & 1 ^ 1; - mbarrier_wait(p_empty_addr + (p_stage) * 8, p_empty_phase); - uint32_t _tmem_load_0_bf16[32]; - #pragma unroll - for (int _lp = 0; _lp < 32; _lp++) { - __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_0[_lp*2 + 0], _tmem_load_0[_lp*2+1 + 0])); - _tmem_load_0_bf16[_lp] = *(uint32_t*)&_bf2; - } - int smem_p_stage = p_stage * 2 + n_half; - int p_base_smem = smem_p_addr + (unsigned int)(smem_p_stage * 8192); - #pragma unroll - for (int vec = 0; vec < 8; vec++) { - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_base_smem + (my_row * 128 + vec * 16 ^ (my_row * 128 + vec * 16 >> 7 & 7) << 4))), "r"(_tmem_load_0_bf16[vec * 4]), "r"(_tmem_load_0_bf16[vec * 4 + 1]), "r"(_tmem_load_0_bf16[vec * 4 + 2]), "r"(_tmem_load_0_bf16[vec * 4 + 3]) : "memory"); - } - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(p_full_addr + (unsigned int)(p_stage * 8)), "r"(0) : "memory"); - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(s_empty_addr + (unsigned int)(score_phase * 8)), "r"(0) : "memory"); - } - mbarrier_arrive(index_valid_empty_addr + (valid_slot) * 8); - valid_iter = valid_iter + 1; - tile_cursor = tile_cursor + num_kv_tiles; - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - } - } - // ---- Role: correction_wg ---- - if (warp >= 4 && warp <= 7) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); - { // correction_wg_main - const int correction_dummy = 0; - float softmax_scale_log2_1 = bmm1_scale[0] * 1.4426950408889634f; - float output_scale = bmm2_scale[0]; - const int head_row_base_1 = warp % 2 * 32; - const int output_row_base = warp % 4 * 32; - const int n_half_1 = warp % 4 / 2; - const int my_row_1 = head_row_base_1 + lane; - const int corr_row = head_row_base_1 << 16; - int tile_cursor_1 = 0; - unsigned int _phase_o_full_0 = 0; - #pragma unroll 1 - for (unsigned int work_idx_1 = cluster_id; work_idx_1 < total_work_items; work_idx_1 += num_clusters) { - int split_idx_1 = work_idx_1 % (unsigned int)(((1) ? 5 : 1)); - int query_idx_1 = work_idx_1 / (unsigned int)(((1) ? 5 : 1)); - int owner_capacity_1 = ((sparse_topk + 128 - 1) / 128 + ((1) ? 5 : 1) - 1) / ((1) ? 5 : 1); - int first_owner_tile_1 = work_idx_1 % (unsigned int)(((1) ? 5 : 1)) * (unsigned int)owner_capacity_1; - int _min_4 = ((owner_capacity_1) < ((sparse_topk + 128 - 1) / 128 - first_owner_tile_1) ? (owner_capacity_1) : ((sparse_topk + 128 - 1) / 128 - first_owner_tile_1)); - int num_kv_tiles_1 = _min_4; - float final_sum = 0.0f; - float final_max = -CAKE_INF; - #pragma unroll 1 - for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { - int pipeline_tile_1 = tile_cursor_1 + tile_1; - int score_phase_1 = pipeline_tile_1 & 1; - int stats_phase = pipeline_tile_1 >> 1 & 1; - mbarrier_wait(stats_addr + (score_phase_1) * 8, stats_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - int meta_addr_1 = taddr + 384 + (unsigned int)(score_phase_1 * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row; - float _tmem_load_1[4]; - tmem_ld_x4(&_tmem_load_1[0], meta_addr_1); - final_sum = _tmem_load_1[0]; - final_max = _tmem_load_1[1]; - float acc_scale_1 = _tmem_load_1[2]; - float no_correction_1 = _tmem_load_1[3]; - if (tile_1 > 0) { - mbarrier_wait(o_full_addr, _phase_o_full_0); - _phase_o_full_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - int _vote_0 = __all_sync(0xFFFFFFFF, no_correction_1 == 1.0f); - int skip_correction = _vote_0; - if (skip_correction == 0) { - #pragma unroll - for (int local_slice = 0; local_slice < 2; local_slice++) { - #pragma unroll - for (int acc_stage = 0; acc_stage < 2; acc_stage++) { - int o_base = taddr + 128 + (unsigned int)(acc_stage * 128) + (unsigned int)(local_slice * 64) + (unsigned int)(output_row_base << 16); - #pragma unroll - for (int c = 0; c < 64; c += 16) { - float _tmem_load_2[16]; - tmem_ld_x16(&_tmem_load_2[0], o_base + c); - const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; - #pragma unroll - for (int _ls = 0; _ls < 8; _ls++) - mul_f32x2_inplace(&reinterpret_cast(_tmem_load_2)[_ls], _scale2_0); - tmem_st_x16_f32(o_base + c, _tmem_load_2); - } - } - } - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - } - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(o_empty_addr), "r"(0) : "memory"); - } - mbarrier_arrive(stats_empty_addr + (score_phase_1) * 8); - } - mbarrier_wait(o_full_addr, _phase_o_full_0); - _phase_o_full_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - const int exchange_idx_1 = n_half_1 * 64 + my_row_1; - smem_epilogue_exchange[exchange_idx_1] = final_sum; - asm volatile("barrier.sync 4, 128;" ::: "memory"); - final_sum = final_sum + smem_epilogue_exchange[exchange_idx_1 ^ 64]; - asm volatile("barrier.sync 4, 128;" ::: "memory"); - float _rcp_0 = approx_rcp(final_sum); - float inv_sum = ((final_sum > 0.0f) ? _rcp_0 : 0.0f); - int head_idx_1 = cta_rank * 64 + my_row_1; - int output_offset = ((1) ? ((query_idx_1 * num_heads + head_idx_1) * ((1) ? 5 : 1) + split_idx_1) * 512 : (query_idx_1 * num_heads + head_idx_1) * 512); - if (n_half_1 == 0 && head_idx_1 < num_heads) { - int stat_offset = (query_idx_1 * num_heads + head_idx_1) * ((1) ? 5 : 1) + split_idx_1; - float _log2_0; - asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(final_sum)); - partial_lse[stat_offset] = ((final_sum > 0.0f) ? final_max * softmax_scale_log2_1 + _log2_0 : -CAKE_INF); - } - #pragma unroll - for (int acc_stage_1 = 0; acc_stage_1 < 2; acc_stage_1++) { - #pragma unroll - for (int local_slice_1 = 0; local_slice_1 < 2; local_slice_1++) { - int logical_slice = acc_stage_1 * 2 * 2 + n_half_1 * 2 + local_slice_1; - int o_base_1 = taddr + 128 + (unsigned int)(acc_stage_1 * 128) + (unsigned int)(local_slice_1 * 64) + (unsigned int)(output_row_base << 16); - #pragma unroll - for (int c_1 = 0; c_1 < 64; c_1 += 32) { - float _tmem_load_3[32]; - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]) - : "r"(o_base_1 + c_1)); - int gmem_base = output_offset + logical_slice * 64 + c_1; - #pragma unroll - for (int j = 0; j < 32; j += 16) { - { - const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; - #if __CUDA_ARCH__ >= 1000 - #pragma unroll - for (int _ps = 0; _ps < 8; _ps++) - mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_3[j])[_ps], _prescale2_1); - #else - #pragma unroll - for (int _ps = 0; _ps < 16; _ps++) - _tmem_load_3[j + _ps] *= inv_sum * output_scale; - #endif - { - __nv_bfloat162 _pk0 = __floats2bfloat162_rn(_tmem_load_3[j + 0], _tmem_load_3[j + 1]); - unsigned _pk_u0 = *reinterpret_cast(&_pk0); - __nv_bfloat162 _pk1 = __floats2bfloat162_rn(_tmem_load_3[j + 2], _tmem_load_3[j + 3]); - unsigned _pk_u1 = *reinterpret_cast(&_pk1); - __nv_bfloat162 _pk2 = __floats2bfloat162_rn(_tmem_load_3[j + 4], _tmem_load_3[j + 5]); - unsigned _pk_u2 = *reinterpret_cast(&_pk2); - __nv_bfloat162 _pk3 = __floats2bfloat162_rn(_tmem_load_3[j + 6], _tmem_load_3[j + 7]); - unsigned _pk_u3 = *reinterpret_cast(&_pk3); - __nv_bfloat162 _pk4 = __floats2bfloat162_rn(_tmem_load_3[j + 8], _tmem_load_3[j + 9]); - unsigned _pk_u4 = *reinterpret_cast(&_pk4); - __nv_bfloat162 _pk5 = __floats2bfloat162_rn(_tmem_load_3[j + 10], _tmem_load_3[j + 11]); - unsigned _pk_u5 = *reinterpret_cast(&_pk5); - __nv_bfloat162 _pk6 = __floats2bfloat162_rn(_tmem_load_3[j + 12], _tmem_load_3[j + 13]); - unsigned _pk_u6 = *reinterpret_cast(&_pk6); - __nv_bfloat162 _pk7 = __floats2bfloat162_rn(_tmem_load_3[j + 14], _tmem_load_3[j + 15]); - unsigned _pk_u7 = *reinterpret_cast(&_pk7); - asm volatile( - "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" - :: "l"((void*)(&((__nv_bfloat16*)(O + (gmem_base + j)))[0])), "r"(_pk_u0), "r"(_pk_u1), "r"(_pk_u2), "r"(_pk_u3), "r"(_pk_u4), "r"(_pk_u5), "r"(_pk_u6), "r"(_pk_u7) : "memory"); - } - } - } - } - } - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(o_empty_addr), "r"(0) : "memory"); - tile_cursor_1 = tile_cursor_1 + num_kv_tiles_1; - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - } - } - // ---- Role: mma_warp ---- - if (warp == 8) { - { // mma_warp_main - const int mma_dummy = 0; - unsigned int k_stage = 0; - unsigned int v_stage = 0; - int tile_cursor_2 = 0; - unsigned int _phase_o_first_slice_seeded_0 = 0; - mbarrier_wait(o_first_slice_seeded_addr, _phase_o_first_slice_seeded_0); - _phase_o_first_slice_seeded_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - unsigned int _phase_q_full_0 = 0; - unsigned int _phase_k_full = 0; - unsigned int _phase_o_empty_0 = 1; - unsigned int _phase_v_full = 0; - #pragma unroll 1 - for (unsigned int work_idx_2 = cluster_id; work_idx_2 < total_work_items; work_idx_2 += num_clusters) { - int owner_capacity_2 = ((sparse_topk + 128 - 1) / 128 + ((1) ? 5 : 1) - 1) / ((1) ? 5 : 1); - int first_owner_tile_2 = work_idx_2 % (unsigned int)(((1) ? 5 : 1)) * (unsigned int)owner_capacity_2; - int _min_1 = ((owner_capacity_2) < ((sparse_topk + 128 - 1) / 128 - first_owner_tile_2) ? (owner_capacity_2) : ((sparse_topk + 128 - 1) / 128 - first_owner_tile_2)); - int num_kv_tiles_2 = _min_1; - if (cta_rank == 0) { - mbarrier_wait(q_full_addr, _phase_q_full_0); - _phase_q_full_0 ^= 1; - int first_pv = 1; - #pragma unroll 1 - for (int tile_2 = 0; tile_2 < num_kv_tiles_2; tile_2++) { - int pipeline_tile_2 = tile_cursor_2 + tile_2; - int score_phase_2 = pipeline_tile_2 & 1; - int s_empty_phase = pipeline_tile_2 >> 1 & 1 ^ 1; - mbarrier_wait(s_empty_addr + (score_phase_2) * 8, s_empty_phase); - int score_col = ((score_phase_2 != 0) ? 64 : 0); - #pragma unroll 1 - for (int qk_stage = 0; qk_stage < 4; qk_stage++) { - mbarrier_wait(k_full_addr + (k_stage) * 8, _phase_k_full); - int _mma_a_lo_0 = (((smem_q_addr) >> 4) & 0x3FFF) + (qk_stage * 2) * 512; - int _mma_b_lo_0 = (((smem_k_addr) >> 4) & 0x3FFF) + (k_stage) * 1024; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 136316048;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (score_col))), "r"(((qk_stage == 0) ? 0 : 1))); - int _mma_a_lo_1 = (((smem_q_addr) >> 4) & 0x3FFF) + (qk_stage * 2 + 1) * 512; - int _mma_b_lo_1 = (((smem_k_addr + 8192) >> 4) & 0x3FFF) + (k_stage) * 1024; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 136316048;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); - elect_commit_cg2_multicast(k_empty_addr + (k_stage) * 8, (uint16_t)(3)); - k_stage += 1; - if (k_stage == 4) { k_stage = 0; _phase_k_full ^= 1; } - } - elect_commit_cg2_multicast(s_full_addr + (score_phase_2) * 8, (uint16_t)(3)); - if (tile_2 > 0) { - int prev_pipeline_tile = pipeline_tile_2 - 1; - int p_stage_1 = prev_pipeline_tile % 2; - int p_full_phase = prev_pipeline_tile / 2 & 1; - mbarrier_wait(o_empty_addr, _phase_o_empty_0); - _phase_o_empty_0 ^= 1; - mbarrier_wait(p_full_addr + (p_stage_1) * 8, p_full_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - #pragma unroll - for (int ps = 0; ps < 2; ps++) { - int smem_p_stage_1 = p_stage_1 * 2 + ps; - mbarrier_wait(v_full_addr + (v_stage) * 8, _phase_v_full); - #pragma unroll - for (int acc_stage_2 = 0; acc_stage_2 < 2; acc_stage_2++) { - int output_col = 128 + acc_stage_2 * 128; - if (acc_stage_2 == 0) { - int _mma_a_lo_2 = (((smem_p_addr) >> 4) & 0x3FFF) + (smem_p_stage_1) * 512; - int _mma_b_lo_2 = ((((smem_v_lo_addr) >> 4) & 0x3FFF) | 0x2000000) + (v_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 138478736;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (output_col))), "r"(((((ps == 0) ? first_pv : 0)) ? 0 : 1))); - } else { - int _mma_a_lo_3 = (((smem_p_addr) >> 4) & 0x3FFF) + (smem_p_stage_1) * 512; - int _mma_b_lo_3 = ((((smem_v_hi_addr) >> 4) & 0x3FFF) | 0x2000000) + (v_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 138478736;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"((tmem_tmem_scratch + (output_col))), "r"(((((ps == 0) ? first_pv : 0)) ? 0 : 1))); - } - } - elect_commit_cg2_multicast(v_empty_addr + (v_stage) * 8, (uint16_t)(3)); - v_stage += 1; - if (v_stage == 2) { v_stage = 0; _phase_v_full ^= 1; } - } - first_pv = 0; - elect_commit_cg2_multicast(p_empty_addr + (p_stage_1) * 8, (uint16_t)(3)); - elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); - } - } - int last_pipeline_tile = tile_cursor_2 + num_kv_tiles_2 - 1; - int last_p_stage = last_pipeline_tile % 2; - int last_p_phase = last_pipeline_tile / 2 & 1; - mbarrier_wait(o_empty_addr, _phase_o_empty_0); - _phase_o_empty_0 ^= 1; - mbarrier_wait(p_full_addr + (last_p_stage) * 8, last_p_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - #pragma unroll - for (int ps_1 = 0; ps_1 < 2; ps_1++) { - int smem_p_stage_2 = last_p_stage * 2 + ps_1; - mbarrier_wait(v_full_addr + (v_stage) * 8, _phase_v_full); - #pragma unroll - for (int acc_stage_3 = 0; acc_stage_3 < 2; acc_stage_3++) { - int output_col_1 = 128 + acc_stage_3 * 128; - if (acc_stage_3 == 0) { - int _mma_a_lo_4 = (((smem_p_addr) >> 4) & 0x3FFF) + (smem_p_stage_2) * 512; - int _mma_b_lo_4 = ((((smem_v_lo_addr) >> 4) & 0x3FFF) | 0x2000000) + (v_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 138478736;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"((tmem_tmem_scratch + (output_col_1))), "r"(((((ps_1 == 0) ? first_pv : 0)) ? 0 : 1))); - } else { - int _mma_a_lo_5 = (((smem_p_addr) >> 4) & 0x3FFF) + (smem_p_stage_2) * 512; - int _mma_b_lo_5 = ((((smem_v_hi_addr) >> 4) & 0x3FFF) | 0x2000000) + (v_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 138478736;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_5), "r"(_mma_b_lo_5), "r"((tmem_tmem_scratch + (output_col_1))), "r"(((((ps_1 == 0) ? first_pv : 0)) ? 0 : 1))); - } - } - elect_commit_cg2_multicast(v_empty_addr + (v_stage) * 8, (uint16_t)(3)); - v_stage += 1; - if (v_stage == 2) { v_stage = 0; _phase_v_full ^= 1; } - } - elect_commit_cg2_multicast(p_empty_addr + (last_p_stage) * 8, (uint16_t)(3)); - elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); - elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); - } - tile_cursor_2 = tile_cursor_2 + num_kv_tiles_2; - } - if (cta_rank == 0) { - #pragma unroll - for (int tail_offset = 0; tail_offset < 2; tail_offset++) { - int tail_tile = tile_cursor_2 + tail_offset; - int tail_stage = tail_tile & 1; - int tail_phase = tail_tile >> 1 & 1 ^ 1; - mbarrier_wait(s_empty_addr + (tail_stage) * 8, tail_phase); - } - mbarrier_wait(o_empty_addr, _phase_o_empty_0); - _phase_o_empty_0 ^= 1; - } - asm volatile("tcgen05.fence::before_thread_sync;"); - int peer_rank_1 = cta_rank ^ 1; - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(tmem_dealloc_peer_addr), "r"(peer_rank_1) : "memory"); - mbarrier_wait(tmem_dealloc_peer_addr, 0); - int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); - asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); - } - } - // ---- Role: load_wg ---- - if (warp >= 9 && warp <= 12) { - { // load_wg_main - const int load_dummy = 0; - const int load_warp_rank = warp - 9; - unsigned int k_stage_1 = 0; - unsigned int v_stage_1 = 0; - unsigned int _phase_q_empty_0 = 1; - unsigned int _phase_k_empty = 1; - unsigned int _phase_v_empty = 1; - #pragma unroll 1 - for (unsigned int work_idx_3 = cluster_id; work_idx_3 < total_work_items; work_idx_3 += num_clusters) { - int split_idx_2 = work_idx_3 % (unsigned int)(((1) ? 5 : 1)); - int query_idx_2 = work_idx_3 / (unsigned int)(((1) ? 5 : 1)); - int owner_capacity_3 = ((sparse_topk + 128 - 1) / 128 + ((1) ? 5 : 1) - 1) / ((1) ? 5 : 1); - int first_owner_tile_3 = work_idx_3 % (unsigned int)(((1) ? 5 : 1)) * (unsigned int)owner_capacity_3; - int _min_0 = ((owner_capacity_3) < ((sparse_topk + 128 - 1) / 128 - first_owner_tile_3) ? (owner_capacity_3) : ((sparse_topk + 128 - 1) / 128 - first_owner_tile_3)); - int num_kv_tiles_3 = _min_0; - int first_kv_tile_1 = first_owner_tile_3; - mbarrier_wait(q_empty_addr, _phase_q_empty_0); - _phase_q_empty_0 ^= 1; - if (load_warp_rank == 0) { - if (cta_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(131072)) : "memory"); - } - } - if (elect_sync()) { - #pragma unroll - for (int q_stage = 0; q_stage < 8; q_stage++) { - tma_4d_gmem2smem_cta2(smem_q_addr + (unsigned int)(q_stage * 8192), (&tmap_q), 0, cta_rank * 64, q_stage, query_idx_2, ((q_full_addr) & 0xFEFFFFFF)); - } - } - } - asm volatile("barrier.sync 2, 160;" ::: "memory"); - #pragma unroll - for (int qk_stage_1 = 0; qk_stage_1 < 4; qk_stage_1++) { - mbarrier_wait(k_empty_addr + (k_stage_1) * 8, _phase_k_empty); - if (load_warp_rank == 0) { - if (cta_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); - } - } - } - int group = load_warp_rank * 4 + lane; - if (lane < 4) { - int index_offset = cta_rank * 64 + group * 4; - int raw_rows[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows[0])), "=r"(*reinterpret_cast(&raw_rows[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows[(0) + 3])) - : "r"(smem_indices_addr + (unsigned int)(index_offset * 4))); - if (first_kv_tile_1 == 0) { - tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group * 512), (&tmap_swa_k), qk_stage_1 * 128, ((raw_rows[0] >= 0) ? raw_rows[0] : 0), ((raw_rows[1] >= 0) ? raw_rows[1] : 0), ((raw_rows[2] >= 0) ? raw_rows[2] : 0), ((raw_rows[3] >= 0) ? raw_rows[3] : 0), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group * 512) + 8192, (&tmap_swa_k), qk_stage_1 * 128 + 64, ((raw_rows[0] >= 0) ? raw_rows[0] : 0), ((raw_rows[1] >= 0) ? raw_rows[1] : 0), ((raw_rows[2] >= 0) ? raw_rows[2] : 0), ((raw_rows[3] >= 0) ? raw_rows[3] : 0), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - } else { - tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group * 512), (&tmap_compressed_k), qk_stage_1 * 128, ((raw_rows[0] >= 0) ? raw_rows[0] : 0), ((raw_rows[1] >= 0) ? raw_rows[1] : 0), ((raw_rows[2] >= 0) ? raw_rows[2] : 0), ((raw_rows[3] >= 0) ? raw_rows[3] : 0), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group * 512) + 8192, (&tmap_compressed_k), qk_stage_1 * 128 + 64, ((raw_rows[0] >= 0) ? raw_rows[0] : 0), ((raw_rows[1] >= 0) ? raw_rows[1] : 0), ((raw_rows[2] >= 0) ? raw_rows[2] : 0), ((raw_rows[3] >= 0) ? raw_rows[3] : 0), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - } - } - k_stage_1 += 1; - if (k_stage_1 == 4) { k_stage_1 = 0; _phase_k_empty ^= 1; } - } - #pragma unroll 1 - for (int tile_3 = 1; tile_3 < num_kv_tiles_3; tile_3++) { - #pragma unroll - for (int qk_stage_2 = 0; qk_stage_2 < 4; qk_stage_2++) { - mbarrier_wait(k_empty_addr + (k_stage_1) * 8, _phase_k_empty); - if (load_warp_rank == 0) { - if (cta_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); - } - } - } - int group_1 = load_warp_rank * 4 + lane; - if (lane < 4) { - int index_offset_1 = tile_3 * 128 + cta_rank * 64 + group_1 * 4; - int raw_rows_1[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_1[0])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 3])) - : "r"(smem_indices_addr + (unsigned int)(index_offset_1 * 4))); - tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group_1 * 512), (&tmap_compressed_k), qk_stage_2 * 128, ((raw_rows_1[0] >= 0) ? raw_rows_1[0] : 0), ((raw_rows_1[1] >= 0) ? raw_rows_1[1] : 0), ((raw_rows_1[2] >= 0) ? raw_rows_1[2] : 0), ((raw_rows_1[3] >= 0) ? raw_rows_1[3] : 0), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group_1 * 512) + 8192, (&tmap_compressed_k), qk_stage_2 * 128 + 64, ((raw_rows_1[0] >= 0) ? raw_rows_1[0] : 0), ((raw_rows_1[1] >= 0) ? raw_rows_1[1] : 0), ((raw_rows_1[2] >= 0) ? raw_rows_1[2] : 0), ((raw_rows_1[3] >= 0) ? raw_rows_1[3] : 0), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - } - k_stage_1 += 1; - if (k_stage_1 == 4) { k_stage_1 = 0; _phase_k_empty ^= 1; } - } - int prev_tile = tile_3 - 1; - #pragma unroll - for (int ps_2 = 0; ps_2 < 2; ps_2++) { - mbarrier_wait(v_empty_addr + (v_stage_1) * 8, _phase_v_empty); - if (load_warp_rank == 0) { - if (cta_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(65536)) : "memory"); - } - } - } - int group_2 = load_warp_rank * 4 + lane; - if (lane < 4) { - int index_offset_2 = prev_tile * 128 + ps_2 * 64 + group_2 * 4; - int raw_rows_2[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_2[0])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 3])) - : "r"(smem_indices_addr + (unsigned int)(index_offset_2 * 4))); - #pragma unroll - for (int acc_stage_4 = 0; acc_stage_4 < 2; acc_stage_4++) { - int dst_v = ((acc_stage_4 == 0) ? smem_v_lo_addr + v_stage_1 * 32768 : smem_v_hi_addr + v_stage_1 * 32768); - if (first_kv_tile_1 + prev_tile == 0) { - tma_gather4_gmem2smem_mc_cta2(dst_v + group_2 * 512, (&tmap_swa_v), acc_stage_4 * 128 * 2 + cta_rank * 128, ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : 0), ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : 0), ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : 0), ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : 0), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - tma_gather4_gmem2smem_mc_cta2(dst_v + group_2 * 512 + 8192, (&tmap_swa_v), acc_stage_4 * 128 * 2 + cta_rank * 128 + 64, ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : 0), ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : 0), ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : 0), ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : 0), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - } else { - tma_gather4_gmem2smem_mc_cta2(dst_v + group_2 * 512, (&tmap_compressed_v), acc_stage_4 * 128 * 2 + cta_rank * 128, ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : 0), ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : 0), ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : 0), ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : 0), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - tma_gather4_gmem2smem_mc_cta2(dst_v + group_2 * 512 + 8192, (&tmap_compressed_v), acc_stage_4 * 128 * 2 + cta_rank * 128 + 64, ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : 0), ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : 0), ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : 0), ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : 0), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - } - } - } - v_stage_1 += 1; - if (v_stage_1 == 2) { v_stage_1 = 0; _phase_v_empty ^= 1; } - } - } - int last_tile = num_kv_tiles_3 - 1; - #pragma unroll - for (int ps_3 = 0; ps_3 < 2; ps_3++) { - mbarrier_wait(v_empty_addr + (v_stage_1) * 8, _phase_v_empty); - if (load_warp_rank == 0) { - if (cta_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(65536)) : "memory"); - } - } - } - int group_3 = load_warp_rank * 4 + lane; - if (lane < 4) { - int index_offset_3 = last_tile * 128 + ps_3 * 64 + group_3 * 4; - int raw_rows_3[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_3[0])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 3])) - : "r"(smem_indices_addr + (unsigned int)(index_offset_3 * 4))); - #pragma unroll - for (int acc_stage_5 = 0; acc_stage_5 < 2; acc_stage_5++) { - int dst_v_1 = ((acc_stage_5 == 0) ? smem_v_lo_addr + v_stage_1 * 32768 : smem_v_hi_addr + v_stage_1 * 32768); - if (first_kv_tile_1 + last_tile == 0) { - tma_gather4_gmem2smem_mc_cta2(dst_v_1 + group_3 * 512, (&tmap_swa_v), acc_stage_5 * 128 * 2 + cta_rank * 128, ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : 0), ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : 0), ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : 0), ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : 0), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - tma_gather4_gmem2smem_mc_cta2(dst_v_1 + group_3 * 512 + 8192, (&tmap_swa_v), acc_stage_5 * 128 * 2 + cta_rank * 128 + 64, ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : 0), ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : 0), ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : 0), ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : 0), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - } else { - tma_gather4_gmem2smem_mc_cta2(dst_v_1 + group_3 * 512, (&tmap_compressed_v), acc_stage_5 * 128 * 2 + cta_rank * 128, ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : 0), ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : 0), ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : 0), ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : 0), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - tma_gather4_gmem2smem_mc_cta2(dst_v_1 + group_3 * 512 + 8192, (&tmap_compressed_v), acc_stage_5 * 128 * 2 + cta_rank * 128 + 64, ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : 0), ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : 0), ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : 0), ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : 0), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - } - } - } - v_stage_1 += 1; - if (v_stage_1 == 2) { v_stage_1 = 0; _phase_v_empty ^= 1; } - } - asm volatile("barrier.sync 2, 160;" ::: "memory"); - } - mbarrier_wait(q_empty_addr, _phase_q_empty_0); - _phase_q_empty_0 ^= 1; - } - } - // ---- Role: index_warp ---- - if (warp == 13) { - { // index_warp_main - const int index_dummy = 0; - int valid_iter_1 = 0; - #pragma unroll 1 - for (unsigned int work_idx_4 = cluster_id; work_idx_4 < total_work_items; work_idx_4 += num_clusters) { - int split_idx_3 = work_idx_4 % (unsigned int)(((1) ? 5 : 1)); - int query_idx_3 = work_idx_4 / (unsigned int)(((1) ? 5 : 1)); - int owner_col_base = split_idx_3 * ((9 + ((1) ? 5 : 1) - 1) / ((1) ? 5 : 1)) * 128; - int* swa_row = swa_indices + (query_idx_3 * swa_index_stride); - int* compressed_row = compressed_indices + (query_idx_3 * compressed_index_stride); - int valid_slot_1 = valid_iter_1 & 1; - int valid_empty_phase = valid_iter_1 >> 1 & 1 ^ 1; - mbarrier_wait(index_valid_empty_addr + (valid_slot_1) * 8, valid_empty_phase); - #pragma unroll - for (int index_pass = 0; index_pass < (9 + ((1) ? 5 : 1) - 1) / ((1) ? 5 : 1); index_pass++) { - int index_offset_4 = index_pass * 128 + lane * 4; - int combined_col = owner_col_base + index_offset_4; - int* index_src = ((combined_col < 128) ? swa_row + combined_col : compressed_row + (combined_col - 128)); - int values[4]; - #pragma unroll - for (int i = 0; i < 4; i++) { - values[i] = -1; - } - if (combined_col + 4 <= sparse_topk) { - int _vec_load_0[4]; - { - const int4* _ivptr_0 = reinterpret_cast(index_src + 0); - int4 _ivld_0; - _ivld_0 = *_ivptr_0; - _vec_load_0[0 + 0] = _ivld_0.x; - _vec_load_0[0 + 1] = _ivld_0.y; - _vec_load_0[0 + 2] = _ivld_0.z; - _vec_load_0[0 + 3] = _ivld_0.w; - } - #pragma unroll - for (int i_1 = 0; i_1 < 4; i_1++) { - values[i_1] = _vec_load_0[i_1]; - } - } - if (combined_col < sparse_topk && combined_col + 4 > sparse_topk) { - #pragma unroll - for (int i_2 = 0; i_2 < 4; i_2++) { - if (combined_col + i_2 < sparse_topk) { - values[i_2] = index_src[i_2]; - } - } - } - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(smem_indices_addr + (unsigned int)(index_offset_4 * 4)), "r"(values[0]), "r"(values[1]), "r"(values[2]), "r"(values[3]) : "memory"); - unsigned int valid_word = (((values[0] >= 0) ? 1 : 0) | ((values[1] >= 0) ? 1 : 0) << 1 | ((values[2] >= 0) ? 1 : 0) << 2 | ((values[3] >= 0) ? 1 : 0) << 3) << (lane & 7) * 4; - unsigned int _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, valid_word, 1); - valid_word = valid_word | _shfl_xor_0; - unsigned int _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, valid_word, 2); - valid_word = valid_word | _shfl_xor_1; - unsigned int _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, valid_word, 4); - valid_word = valid_word | _shfl_xor_2; - if ((lane & 7) == 0) { - smem_index_valid[valid_slot_1 * 8 + index_pass * 4 + (lane >> 3)] = valid_word; - } - } - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(index_valid_full_addr + (valid_slot_1) * 8); - asm volatile("barrier.sync 2, 160;" ::: "memory"); - asm volatile("barrier.sync 2, 160;" ::: "memory"); - valid_iter_1 = valid_iter_1 + 1; - } - } - } - // ---- Role: empty_warps ---- - if (warp >= 14 && warp <= 15) { - { // empty_warps_main - const int empty_dummy = 0; - #pragma unroll 1 - for (unsigned int work_idx_5 = cluster_id; work_idx_5 < total_work_items; work_idx_5 += num_clusters) { - asm volatile("tcgen05.fence::after_thread_sync;"); - } - } - } - - // Cleanup -} - -} // extern "C" diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_e1660a7f44d156286fea_binding.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_e1660a7f44d156286fea_binding.cu deleted file mode 100644 index d3bccd7047a..00000000000 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_e1660a7f44d156286fea_binding.cu +++ /dev/null @@ -1,555 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -// Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_e1660a7f44d156286fea'. -#include -#include -#include -#include - -#include "tvm_ffi_utils.h" - -#include -#include -#include -#include -#include -#include -#include - -extern "C" __global__ void kernel_cake_dsv4_e1660a7f44d156286fea(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); - - -namespace cake_host_shim_244c86188c16b5f1 { - -using tvm::ffi::Optional; -using tvm::ffi::TensorView; - -class ScopedCudaDevice { - public: - explicit ScopedCudaDevice(int device_id) { - cudaError_t error = cudaGetDevice(&previous_device_); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed before host-shim launch: cudaError=" - << static_cast(error); - if (previous_device_ != device_id) { - error = cudaSetDevice(device_id); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaSetDevice failed before host-shim launch for cuda:" - << device_id << ": cudaError=" << static_cast(error); - restore_ = true; - } - } - - ScopedCudaDevice(const ScopedCudaDevice&) = delete; - ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; - - ~ScopedCudaDevice() noexcept { - if (restore_) { - (void)cudaSetDevice(previous_device_); - } - } - - private: - int previous_device_ = -1; - bool restore_ = false; -}; - -inline int64_t CakeDeviceMultiprocessorCount(int device_id) { - constexpr int kMaxCachedCudaDevices = 64; - TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) - << "physical-SM-count cache does not cover cuda:" << device_id; - static std::atomic count_by_device[kMaxCachedCudaDevices]{}; - int cached = count_by_device[device_id].load(std::memory_order_acquire); - if (cached > 0) return cached; - - int count = 0; - cudaError_t error = cudaDeviceGetAttribute( - &count, cudaDevAttrMultiProcessorCount, device_id); - TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) - << "querying multiProcessorCount failed for cuda:" << device_id - << ": cudaError=" << static_cast(error) << ", count=" << count; - // Concurrent first launches may repeat the immutable device query, but all - // publication is atomic and every later hot-path lookup is one acquire load. - count_by_device[device_id].store(count, std::memory_order_release); - return count; -} - -inline void CheckCudaTensor(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) - << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; -} - -inline void CheckSameCudaDevice( - const TensorView& t, - const TensorView& reference, - const char* name, - const char* reference_name) { - TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) - << name << " must be on the same CUDA device as " << reference_name - << ": got cuda:" << t.device().device_id - << " versus cuda:" << reference.device().device_id; -} - -inline void CheckCurrentCudaDevice( - const TensorView& reference, - const char* reference_name) { - int current_device = -1; - cudaError_t error = cudaGetDevice(¤t_device); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed while validating " << reference_name - << ": cudaError=" << static_cast(error); - TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) - << "current CUDA device must match " << reference_name - << ": current=cuda:" << current_device - << ", tensor=cuda:" << reference.device().device_id; -} - -inline void CheckContiguous(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; -} - -inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { - DLDataType d = t.dtype(); - TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) - << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits - << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits - << ", lanes=" << (int)d.lanes << ")"; -} - -// A logical axis.outer(trailing) folds every source dim above the trailing -// dimensions. Shape products are independent of physical strides, so verify -// the leading dimensions form one dense row-major chain instead of inventing -// a "folded stride". The descriptor reads its exact adjacent physical step -// separately through stride[-(trailing + 1)]. -inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { - TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) - << name << " cannot fold leading dimensions above " << trailing - << " trailing dims from ndim=" << t.ndim(); - int outer_last = t.ndim() - trailing - 1; - if (outer_last <= 0) { - return; - } - int64_t step = t.stride(outer_last); - TVM_FFI_CHECK(step > 0, ValueError) - << name << " physical strides must be positive"; - int64_t expected = step; - for (int axis = outer_last - 1; axis >= 0; --axis) { - expected *= t.size(axis + 1); - if (t.size(axis) > 1) { - TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) - << name << " leading dims are not physically foldable above " << trailing - << " trailing dims: stride(" << axis << ")=" << t.stride(axis) - << ", expected " << expected; - } - } -} - -// 4D TMA descriptor for buffer 'tmap_q' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 3, ValueError) - << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - int64_t d2 = t.size(t.ndim() - 2); - int64_t d3 = t.size(t.ndim() - 3); - TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) - << "TMA source 'tmap_q' trailing dims must be positive"; - TVM_FFI_CHECK(d1 % 64 == 0, ValueError) - << "TMA source 'tmap_q' extent " << d1 - << " must divide exactly by " << 64; - uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(d3)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) - << "TMA box (64, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; - int64_t carrier_stride_0 = d1; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; - int64_t carrier_stride_1 = 64; - TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; - TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; - TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; - int64_t carrier_stride_2 = (d2 * d1); - TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; - TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; - TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; - uint64_t global_strides[3] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - (uint64_t)((carrier_stride_1 * 16) / 8), - (uint64_t)((carrier_stride_2 * 16) / 8), - }; - uint32_t box_dim[4] = {64u, 64u, 1u, 1u}; - uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_swa_k' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_swa_k(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_swa_k' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_swa_k' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_swa_k' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_swa_k"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_swa_k' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_swa_k' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_k'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_swa_k') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_compressed_k' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_compressed_k(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_compressed_k' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_compressed_k' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_compressed_k' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_compressed_k"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_compressed_k' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_compressed_k' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_k'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_k') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_swa_v' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_swa_v(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_swa_v' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_swa_v' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_swa_v' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_swa_v"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_swa_v' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_swa_v' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_v'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_swa_v') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_compressed_v' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_compressed_v(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_compressed_v' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_compressed_v' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_compressed_v' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_compressed_v"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_compressed_v' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_compressed_v' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_v'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_v') failed: CUresult=" << (int)r; - return tm; -} - -void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_compressed_k, TensorView arg_tmap_swa_v, TensorView arg_tmap_compressed_v, TensorView arg_O, TensorView arg_partial_lse, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_num_query_tokens, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, int64_t grid_x, int64_t grid_y, int64_t grid_z) { - DLDevice dev = arg_tmap_q.device(); - ScopedCudaDevice device_guard(dev.device_id); - CheckCudaTensor(arg_tmap_q, "tmap_q"); - CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); - CheckContiguous(arg_tmap_q, "tmap_q"); - CheckCudaTensor(arg_tmap_swa_k, "tmap_swa_k"); - CheckDtype(arg_tmap_swa_k, "tmap_swa_k", 4, 16, 1); - CheckCudaTensor(arg_tmap_compressed_k, "tmap_compressed_k"); - CheckDtype(arg_tmap_compressed_k, "tmap_compressed_k", 4, 16, 1); - CheckCudaTensor(arg_tmap_swa_v, "tmap_swa_v"); - CheckDtype(arg_tmap_swa_v, "tmap_swa_v", 4, 16, 1); - CheckCudaTensor(arg_tmap_compressed_v, "tmap_compressed_v"); - CheckDtype(arg_tmap_compressed_v, "tmap_compressed_v", 4, 16, 1); - CheckCudaTensor(arg_O, "O"); - CheckDtype(arg_O, "O", 4, 16, 1); - CheckContiguous(arg_O, "O"); - CheckCudaTensor(arg_partial_lse, "partial_lse"); - CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); - CheckContiguous(arg_partial_lse, "partial_lse"); - CheckCudaTensor(arg_swa_indices, "swa_indices"); - CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); - CheckContiguous(arg_swa_indices, "swa_indices"); - CheckCudaTensor(arg_compressed_indices, "compressed_indices"); - CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); - CheckContiguous(arg_compressed_indices, "compressed_indices"); - CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); - CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckCudaTensor(arg_sinks, "sinks"); - CheckDtype(arg_sinks, "sinks", 2, 32, 1); - CheckContiguous(arg_sinks, "sinks"); - CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); - CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); - CheckContiguous(arg_bmm1_scale, "bmm1_scale"); - CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); - CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); - CheckContiguous(arg_bmm2_scale, "bmm2_scale"); - TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) - << "scalar 'num_heads' value " << arg_num_heads - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_num_query_tokens >= -2147483648LL && arg_num_query_tokens <= 2147483647LL, ValueError) - << "scalar 'num_query_tokens' value " << arg_num_query_tokens - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) - << "scalar 'swa_index_stride' value " << arg_swa_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) - << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) - << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) - << "scalar 'sparse_topk' value " << arg_sparse_topk - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) - << "scalar 'has_sinks' value " << arg_has_sinks - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) - << "scalar 'total_work_items' value " << arg_total_work_items - << " is outside i32 range [-2147483648, 2147483647]"; - CheckSameCudaDevice(arg_tmap_swa_k, arg_tmap_q, "tmap_swa_k", "tmap_q"); - CheckSameCudaDevice(arg_tmap_compressed_k, arg_tmap_q, "tmap_compressed_k", "tmap_q"); - CheckSameCudaDevice(arg_tmap_swa_v, arg_tmap_q, "tmap_swa_v", "tmap_q"); - CheckSameCudaDevice(arg_tmap_compressed_v, arg_tmap_q, "tmap_compressed_v", "tmap_q"); - CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); - CheckSameCudaDevice(arg_partial_lse, arg_tmap_q, "partial_lse", "tmap_q"); - CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); - CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); - CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); - CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); - CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); - CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); - TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) - << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y - << ", " << grid_z << ")"; - TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) - << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z - << ") must be divisible by cluster dims (2, 1, 1)"; - - cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); - CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); - CUtensorMap p_tmap_swa_k = EncodeTma_tmap_swa_k(arg_tmap_swa_k); - CUtensorMap p_tmap_compressed_k = EncodeTma_tmap_compressed_k(arg_tmap_compressed_k); - CUtensorMap p_tmap_swa_v = EncodeTma_tmap_swa_v(arg_tmap_swa_v); - CUtensorMap p_tmap_compressed_v = EncodeTma_tmap_compressed_v(arg_tmap_compressed_v); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); - float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); - int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); - int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); - int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); - float* p_sinks = static_cast(arg_sinks.data_ptr()); - float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); - float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); - int32_t v_num_heads = (int32_t)arg_num_heads; - int32_t v_num_query_tokens = (int32_t)arg_num_query_tokens; - int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; - int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; - int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; - int32_t v_sparse_topk = (int32_t)arg_sparse_topk; - int32_t v_has_sinks = (int32_t)arg_has_sinks; - int32_t v_total_work_items = (int32_t)arg_total_work_items; - void* kargs[] = {&p_tmap_q, &p_tmap_swa_k, &p_tmap_compressed_k, &p_tmap_swa_v, &p_tmap_compressed_v, &p_O, &p_partial_lse, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_num_query_tokens, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_has_sinks, &v_total_work_items}; - - static std::mutex smem_mu; - static auto* smem_status_by_device = - new std::unordered_map(); - cudaError_t smem_status = cudaSuccess; - { - std::lock_guard lock(smem_mu); - auto it = smem_status_by_device->find(dev.device_id); - if (it == smem_status_by_device->end()) { - int current_device = -1; - smem_status = cudaGetDevice(¤t_device); - if (smem_status == cudaSuccess && current_device != dev.device_id) { - smem_status = cudaErrorInvalidDevice; - } - if (smem_status == cudaSuccess) { - smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_e1660a7f44d156286fea), - cudaFuncAttributeMaxDynamicSharedMemorySize, - 220032); - } - smem_status_by_device->emplace(dev.device_id, smem_status); - } else { - smem_status = it->second; - } - } - TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_e1660a7f44d156286fea failed: " - << cudaGetErrorString(smem_status); - dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); - dim3 block(512u, 1u, 1u); - - // Extended direct-source launch for a linked CUDA kernel symbol. - cudaLaunchAttribute attrs[2]{}; - int n = 0; - attrs[n].id = cudaLaunchAttributeClusterDimension; - attrs[n].val.clusterDim.x = 2u; - attrs[n].val.clusterDim.y = 1u; - attrs[n].val.clusterDim.z = 1u; - ++n; - attrs[n].id = cudaLaunchAttributeClusterSchedulingPolicyPreference; - attrs[n].val.clusterSchedulingPolicyPreference = cudaClusterSchedulingPolicySpread; - ++n; - cudaLaunchConfig_t config{}; - config.gridDim = grid; - config.blockDim = block; - config.dynamicSmemBytes = 220032u; - config.stream = stream; - config.attrs = attrs; - config.numAttrs = n; - cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_e1660a7f44d156286fea), kargs); - TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_e1660a7f44d156286fea failed: " - << cudaGetErrorString(launch_status); -} - -} // namespace cake_host_shim_244c86188c16b5f1 - -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_244c86188c16b5f1::Run); diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_e1660a7f44d156286fea_kernel.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_e1660a7f44d156286fea_kernel.cu deleted file mode 100644 index 5c3bf235129..00000000000 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_e1660a7f44d156286fea_kernel.cu +++ /dev/null @@ -1,2353 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -typedef signed char int8_t; -typedef unsigned char uint8_t; -typedef unsigned short uint16_t; -typedef unsigned int uint32_t; -#if defined(__CUDACC_RTC__) -typedef unsigned long long uint64_t; -#else -typedef unsigned long uint64_t; -#endif -static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); -typedef signed int int32_t; -typedef short int int16_t; -struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; -struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; -static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); -static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); - -#if defined(__CUDACC_RTC__) -typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; -#else -#include -#endif - -static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); -static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); -#include -#include - -__device__ __forceinline__ int make_warp_uniform(int x) { - int result; - asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" - : "=r"(result) : "r"(x)); - return result; -} - -#define CAKE_INF CUDART_INF_F -#define TMEM_NCOLS 512 -#define TMEM_TMEM_SCRATCH_OFFSET 0 -#define NUM_K_PIPE_STAGES 3 -#define NUM_V_PIPE_STAGES 2 -#define SMEM_SMEM_Q_OFF 1024 -#define SMEM_SMEM_Q_STAGE_BYTES 8192 -#define SMEM_SMEM_Q_STRIDE 8192 -#define SMEM_SMEM_K_OFF 66560 -#define SMEM_SMEM_K_STAGE_BYTES 16384 -#define SMEM_SMEM_K_STRIDE 16384 -#define SMEM_SMEM_V_LO_OFF 115712 -#define SMEM_SMEM_V_LO_STAGE_BYTES 16384 -#define SMEM_SMEM_V_LO_STRIDE 32768 -#define SMEM_SMEM_V_HI_OFF 132096 -#define SMEM_SMEM_V_HI_STAGE_BYTES 16384 -#define SMEM_SMEM_V_HI_STRIDE 32768 -#define SMEM_SMEM_P_OFF 181248 -#define SMEM_SMEM_P_STAGE_BYTES 8192 -#define SMEM_SMEM_P_STRIDE 8192 -#define SMEM_SMEM_INDICES_OFF 214016 -#define SMEM_SMEM_INDICES_STAGE_BYTES 2048 -#define SMEM_SMEM_INDICES_STRIDE 2048 -#define SMEM_SMEM_SOFTMAX_EXCHANGE_OFF 218624 -#define SMEM_SMEM_SOFTMAX_EXCHANGE_STAGE_BYTES 512 -#define SMEM_SMEM_SOFTMAX_EXCHANGE_STRIDE 512 -#define SMEM_SMEM_EPILOGUE_EXCHANGE_OFF 219136 -#define SMEM_SMEM_EPILOGUE_EXCHANGE_STAGE_BYTES 512 -#define SMEM_SMEM_EPILOGUE_EXCHANGE_STRIDE 512 -#define SMEM_SMEM_INDEX_VALID_OFF 219648 -#define SMEM_SMEM_INDEX_VALID_STAGE_BYTES 288 -#define SMEM_SMEM_INDEX_VALID_STRIDE 288 -#define SMEM_TOTAL 220032 -#define THREADS 512 - -#include - -__device__ __forceinline__ uint32_t elect_sync() { - uint32_t pred = 0; - asm volatile( - "{\n\t" - ".reg .pred %%px;\n\t" - "elect.sync _|%%px, %1;\n\t" - "@%%px mov.s32 %0, 1;\n\t" - "}\n" - : "+r"(pred) - : "r"(0xFFFFFFFF)); - return pred; -} - - -__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { - asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" - :: "r"(mbar_addr), "r"(count) : "memory"); -} - -__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { - asm volatile("mbarrier.init.b64 [%0], %1;" - :: "l"(mbar_addr), "r"(count) : "memory"); -} - - -__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - - -// CTA-local pipelines have short, resident producer/consumer edges. Omitting -// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the -// explicit loop still makes this helper blocking until acquire succeeds. -__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE;\n\t" - "bra.uni LAB_WAIT;\n\t" - "DONE:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Source-faithful relaxed CTA wait used only by a typed protocol that does -// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. -__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, 10000000;\n\t" - "@P1 bra.uni DONE_RELAXED;\n\t" - "bra.uni LAB_WAIT_RELAXED;\n\t" - "DONE_RELAXED:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Exact source ports may request the PTX suspendTimeHint operand explicitly. -// The hint is expressed in nanoseconds and is kept separate from the canonical -// no-hint CTA helper so unrelated schedules retain their existing retry path. -__device__ __forceinline__ void mbarrier_wait_suspend( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_SUSPEND:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_SUSPEND;\n\t" - "bra.uni LAB_WAIT_SUSPEND;\n\t" - "DONE_SUSPEND:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE_CLUSTER;\n\t" - "bra.uni LAB_WAIT_CLUSTER;\n\t" - "DONE_CLUSTER:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - ".reg .u32 WAIT_ADDR;\n\t" - "mov.u32 WAIT_ADDR, %0;\n\t" - "LAB_WAIT_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [WAIT_ADDR], %1, %2;\n\t" - "@P1 bra.uni DONE_HINT;\n\t" - "bra.uni LAB_WAIT_HINT;\n\t" - "DONE_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -// Exact unqualified CTA wait used by source schedules whose PTX intentionally -// omits the acquire qualifier while retaining a typed suspendTimeHint operand. -__device__ __forceinline__ void mbarrier_wait_relaxed_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED_HINT:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra DONE_RELAXED_HINT;\n\t" - "bra LAB_WAIT_RELAXED_HINT;\n\t" - "DONE_RELAXED_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); -} - -__device__ __forceinline__ void mbarrier_wait_cluster_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" - "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" - "DONE_CLUSTER_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_suspend( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait_cluster(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); - } -} - - -__device__ __forceinline__ void tcgen05_mma_f16_cta2( - int taddr, uint64_t a_desc, uint64_t b_desc, - uint32_t i_desc, int enable_input_d) { - asm volatile( - "{\n\t" - ".reg .pred p;\n\t" - ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - "setp.ne.b32 p, %4, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" - "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "tcgen05.mma.cta_group::2.kind::f16 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" - "}\n" - :: "r"(taddr), "l"(a_desc), "l"(b_desc), - "r"(i_desc), "r"(enable_input_d) - : "memory"); -} - - -__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { - return (x & 0x3FFFFULL) >> 4ULL; -} - - -union MmaSmemDesc { - uint64_t u64; - uint32_t u32[2]; -}; - -__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { - MmaSmemDesc tmp; - tmp.u64 = smem_desc; - tmp.u32[0] += offset; - smem_desc = tmp.u64; -} - - -__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { - asm volatile( - "{\n\t" - ".reg .pred leader;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" - ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" - "}\n" - :: "r"(mbar_addr), "h"(cta_mask) : "memory"); -} - - -__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { - asm volatile( - "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" - :: "r"(mbar_addr), "r"(bytes) : "memory"); -} - - -__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7," - " %8, %9, %10, %11, %12, %13, %14, %15," - " %16, %17, %18, %19, %20, %21, %22, %23," - " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), - "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), - "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), - "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), - "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), - "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), - "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x16.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7," - " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), - "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), - "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ void tmem_ld_x4(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x4.b32" - " {%0, %1, %2, %3}, [%4];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { - asm volatile( - "tcgen05.st.sync.aligned.32x32b.x16.b32" - " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," - " %9, %10, %11, %12, %13, %14, %15, %16};" - :: "r"(tmem_addr), - "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), - "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), - "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), - "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); -} - - -__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { - asm volatile( - "tcgen05.st.sync.aligned.32x32b.x4.b32" - " [%0], {%1, %2, %3, %4};" - :: "r"(tmem_addr), - "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); -} - - -__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { - uint32_t addr; - asm("{\n\t" - ".reg .u64 u64addr;\n\t" - "cvta.to.shared.u64 u64addr, %1;\n\t" - "cvt.u32.u64 %0, u64addr;\n\t" - "}\n" : "=r"(addr) : "l"(ptr)); - return addr; -} - - -__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { - uint32_t remote; - asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" - : "=r"(remote) : "r"(local_addr), "r"(rank)); - return remote; -} - - -__device__ __forceinline__ float approx_exp2(float x) { - float y; - asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float approx_rcp(float x) { - float y; - asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float max_noftz(float a, float b) { - float c; - asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); - return c; -} - - -__device__ __forceinline__ float warp_reduce_max(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); - return val; -} - - -__device__ __forceinline__ float warp_reduce_sum(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val += __shfl_xor_sync(0xFFFFFFFF, val, offset); - return val; -} - - -__device__ __forceinline__ float row_max_reduce(float2 acc) { - return max_noftz(acc.x, acc.y); -} - - -__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { - #pragma unroll - for (int j = 0; j < 16; j++) { - if (j % 2 == 0) - acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); - else - acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); - } -} - - -__device__ __forceinline__ float2 ex2_emulation_f32x2_value(float2 value) { - const float c0 = 1.0f, c1 = 0.695146143436431884765625f; - const float c2 = 0.227564394474029541015625f, c3 = 0.077119089663028717041015625f; - const float magic = 12582912.0f; - float x0 = max_noftz(value.x, -127.0f), x1 = max_noftz(value.y, -127.0f); - float2 xc2 = make_float2(x0, x1), magic2 = make_float2(magic, magic); - float2 xr2; - asm("add.rm.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xr2) - : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&magic2)); - float2 c3_2 = make_float2(c3, c3), c2_2 = make_float2(c2, c2); - float2 c1_2 = make_float2(c1, c1), c0_2 = make_float2(c0, c0); - float2 xrb2, xfrac2; - asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xrb2) - : "l"(*(unsigned long long*)&xr2), "l"(*(unsigned long long*)&magic2)); - asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xfrac2) - : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&xrb2)); - float2 poly2; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&c3_2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c2_2)); - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c1_2)); - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c0_2)); - int x0r_i, x1r_i, p0_i, p1_i; - asm("mov.b64 {%0, %1}, %2;" : "=r"(x0r_i), "=r"(x1r_i) : "l"(*(unsigned long long*)&xr2)); - asm("mov.b64 {%0, %1}, %2;" : "=r"(p0_i), "=r"(p1_i) : "l"(*(unsigned long long*)&poly2)); - float r0, r1; - asm("mov.b32 %0, %1;" : "=f"(r0) : "r"((x0r_i << 23) + p0_i)); - asm("mov.b32 %0, %1;" : "=f"(r1) : "r"((x1r_i << 23) + p1_i)); - return make_float2(r0, r1); -} - -__device__ __forceinline__ void ex2_emulation_f32x2(float* x0_ptr, float* x1_ptr) { - float2 result = ex2_emulation_f32x2_value(make_float2(*x0_ptr, *x1_ptr)); - *x0_ptr = result.x; *x1_ptr = result.y; -} - -__device__ __forceinline__ void softmax_frag_exp2_cast( - float* sv, uint32_t* pv, int use_emu) -{ - #pragma unroll - for (int j = 0; j < 16; j++) { - if (use_emu && j >= 12) - ex2_emulation_f32x2(&sv[j*2], &sv[j*2+1]); - else { - sv[j*2] = approx_exp2(sv[j*2]); - sv[j*2+1] = approx_exp2(sv[j*2+1]); - } - } - #pragma unroll - for (int j = 0; j < 16; j++) { - __nv_bfloat162 bf = __float22bfloat162_rn({sv[j*2], sv[j*2+1]}); - pv[j] = reinterpret_cast(bf); - } -} - - - -__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { - const float2* sv2 = reinterpret_cast(sv); - #pragma unroll - for (int j = 0; j < 16; j++) { - asm("add.f32x2 %0, %1, %2;" - : "+l"(reinterpret_cast(*acc)) - : "l"(reinterpret_cast(*acc)), - "l"(reinterpret_cast(sv2[j]))); - } -} - - -__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { - asm("mul.rn.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { - asm("mul.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { - asm("add.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { - asm("add.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { - asm("sub.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { - asm("sub.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("add.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { - float2 r; - asm("sub.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("sub.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ void fma_scale_x32( - float* sv, const float2* scale2, const float2* neg_max2) -{ - float2* sv_2 = reinterpret_cast(sv); - #pragma unroll - for (int j = 0; j < 16; j++) - fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); -} - -__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("mul.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - - -// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) - -__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("add.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("add.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("add.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("add.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("add.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("add.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("add.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - - -__device__ __forceinline__ void fence_async_shared() { - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); -} - - -__device__ __forceinline__ uint64_t make_smem_desc(int addr) { - const int SBO = 1024; - return desc_encode(addr) - | (desc_encode(SBO) << 32ULL) - | (1ULL << 46ULL) - | (2ULL << 61ULL); -} - - -__device__ __forceinline__ void tma_2d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3}], [%4];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_4d_gmem2smem_cta2( - int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.4d.shared::cluster.global" - ".mbarrier::complete_tx::bytes.cta_group::2" - " [%0], [%1, {%2, %3, %4, %5}], [%6];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr) { - // Canonical .shared::cta form for non-multicast gather4, matching - // trtllm-gen / cuda_ptx and the PTX ISA qualifier order - // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, - // .shared::cluster is reserved for the multicast variant (ctaMask). - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem_cta2( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr) { - // Canonical .shared::cta form; see tma_gather4_gmem2smem above. - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" - ".mbarrier::complete_tx::bytes.cta_group::2" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem_mc( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr, unsigned short cta_mask) { - // Multicast variant: the PTX grammar ties the .shared::cluster - // destination to .multicast::cluster + ctaMask (cf. cuda_ptx / - // SM100_TMA_LOAD_MULTICAST_2D_GATHER4 in CUTLASS). - asm volatile( - "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" - ".mbarrier::complete_tx::bytes.multicast::cluster" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr), "h"(cta_mask) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem_mc_cta2( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr, unsigned short cta_mask) { - // Multicast + cta_group::2 variant; see tma_gather4_gmem2smem_mc. - asm volatile( - "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" - ".mbarrier::complete_tx::bytes.multicast::cluster.cta_group::2" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr), "h"(cta_mask) : "memory"); -} - - -__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { - asm volatile( - "{\n\t" - ".reg .b16 lo, hi;\n\t" - "mov.b32 {lo, hi}, %1;\n\t" - "tcgen05.commit.cta_group::2.mbarrier::arrive::one" - ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" - "}\n" - :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); -} - -extern "C" { - -__global__ __launch_bounds__(512, 1) __cluster_dims__(2,1,1) void -kernel_cake_dsv4_e1660a7f44d156286fea(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items) -{ - const int tid = threadIdx.x; - const int warp = make_warp_uniform(tid / 32); - const int lane = tid % 32; - - extern __shared__ __align__(1024) char smem_raw[]; - int smem; - asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); - smem = make_warp_uniform(smem); - - const int mbar_base = smem; - #define q_full_addr (mbar_base + 0) - #define q_empty_addr (mbar_base + 8) - #define k_full_addr (mbar_base + 16) - #define k_empty_addr (mbar_base + 40) - #define v_full_addr (mbar_base + 64) - #define v_empty_addr (mbar_base + 80) - #define s_full_addr (mbar_base + 96) - #define s_empty_addr (mbar_base + 112) - #define p_full_addr (mbar_base + 128) - #define p_empty_addr (mbar_base + 144) - #define o_empty_addr (mbar_base + 160) - #define stats_addr (mbar_base + 168) - #define stats_empty_addr (mbar_base + 184) - #define o_full_addr (mbar_base + 200) - #define o_first_slice_seeded_addr (mbar_base + 208) - #define tmem_dealloc_peer_addr (mbar_base + 216) - #define index_valid_full_addr (mbar_base + 224) - #define index_valid_empty_addr (mbar_base + 240) - - const int bid = blockIdx.x; - const int num_bids = gridDim.x; - const unsigned int clusters_x = gridDim.x / 2; - const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; - const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; - - int cta_rank; - asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); - - // Kernel setup ops - __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); - const int smem_q_addr = smem + 1024; - __nv_bfloat16* smem_k = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); - const int smem_k_addr = smem + 66560; - __nv_bfloat16* smem_v_lo = reinterpret_cast<__nv_bfloat16*>(smem_raw + 115712); - const int smem_v_lo_addr = smem + 115712; - __nv_bfloat16* smem_v_hi = reinterpret_cast<__nv_bfloat16*>(smem_raw + 132096); - const int smem_v_hi_addr = smem + 132096; - __nv_bfloat16* smem_p = reinterpret_cast<__nv_bfloat16*>(smem_raw + 181248); - const int smem_p_addr = smem + 181248; - int* smem_indices = reinterpret_cast(smem_raw + 214016); - const int smem_indices_addr = smem + 214016; - float* smem_softmax_exchange = reinterpret_cast(smem_raw + 218624); - const int smem_softmax_exchange_addr = smem + 218624; - float* smem_epilogue_exchange = reinterpret_cast(smem_raw + 219136); - const int smem_epilogue_exchange_addr = smem + 219136; - unsigned int* smem_index_valid = reinterpret_cast(smem_raw + 219648); - const int smem_index_valid_addr = smem + 219648; - asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); - asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); - - // Mbarrier init (18 pipeline groups, 0 ordered-sequence groups, 32 barriers) - // Mbarriers at smem_raw[0..256) - - if (warp == 0) { - uint32_t leader = elect_sync(); - if (leader) { - // q_full: 1 barriers, init_count=1 - mbarrier_init(smem + 0, 1); - // q_empty: 1 barriers, init_count=1 - mbarrier_init(smem + 8, 1); - // --- pipeline 'k_pipe' --- - // k_full: 3 barriers, init_count=1 - mbarrier_init(smem + 16, 1); - mbarrier_init(smem + 24, 1); - mbarrier_init(smem + 32, 1); - // k_empty: 3 barriers, init_count=1 - mbarrier_init(smem + 40, 1); - mbarrier_init(smem + 48, 1); - mbarrier_init(smem + 56, 1); - // --- pipeline 'v_pipe' --- - // v_full: 2 barriers, init_count=1 - mbarrier_init(smem + 64, 1); - mbarrier_init(smem + 72, 1); - // v_empty: 2 barriers, init_count=1 - mbarrier_init(smem + 80, 1); - mbarrier_init(smem + 88, 1); - // s_full: 2 barriers, init_count=1 - mbarrier_init(smem + 96, 1); - mbarrier_init(smem + 104, 1); - // s_empty: 2 barriers, init_count=256 - mbarrier_init(smem + 112, 256); - mbarrier_init(smem + 120, 256); - // p_full: 2 barriers, init_count=256 - mbarrier_init(smem + 128, 256); - mbarrier_init(smem + 136, 256); - // p_empty: 2 barriers, init_count=1 - mbarrier_init(smem + 144, 1); - mbarrier_init(smem + 152, 1); - // o_empty: 1 barriers, init_count=256 - mbarrier_init(smem + 160, 256); - // stats: 2 barriers, init_count=128 - mbarrier_init(smem + 168, 128); - mbarrier_init(smem + 176, 128); - // stats_empty: 2 barriers, init_count=128 - mbarrier_init(smem + 184, 128); - mbarrier_init(smem + 192, 128); - // o_full: 1 barriers, init_count=1 - mbarrier_init(smem + 200, 1); - // o_first_slice_seeded: 1 barriers, init_count=256 - mbarrier_init(smem + 208, 256); - // tmem_dealloc_peer: 1 barriers, init_count=32 - mbarrier_init(smem + 216, 32); - // index_valid_full: 2 barriers, init_count=32 - mbarrier_init(smem + 224, 32); - mbarrier_init(smem + 232, 32); - // index_valid_empty: 2 barriers, init_count=128 - mbarrier_init(smem + 240, 128); - mbarrier_init(smem + 248, 128); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - - asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); - asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); - - // TMEM alloc (512 columns, 512 used) - volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 256); - if (warp == 0) { - int _tmem_hold = smem + 256; - asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); - __syncwarp(); - asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); - } - - __syncthreads(); - asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); - asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); - - const int taddr = tmem_addr_storage[0]; - - // Kernel post-init ops - const int tmem_tmem_scratch = taddr; - - // ---- Ordered hardware-WG register redistribution ---- - // Dec phase frees registers before any WG attempts inc. - if (warp >= 8 && warp <= 15) { - asm volatile("setmaxnreg.dec.sync.aligned.u32 56;"); - } - - // ---- Role: softmax_wg ---- - if (warp <= 3) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); - { // softmax_wg_main - const int softmax_dummy = 0; - float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; - const int head_row_base = warp % 2 * 32; - const int score_row_base = warp % 4 * 32; - const int n_half = warp % 4 / 2; - const int my_row = head_row_base + lane; - const int exchange_idx = n_half * 64 + my_row; - int tile_cursor = 0; - int valid_iter = 0; - float seed_zero[4]; - #pragma unroll - for (int seed_i = 0; seed_i < 4; seed_i++) { - seed_zero[seed_i] = 0.0f; - } - #pragma unroll - for (int seed_stage = 0; seed_stage < 2; seed_stage++) { - #pragma unroll - for (int seed_col = 128 + seed_stage * 128; seed_col < 128 + (seed_stage + 1) * 128; seed_col += 4) { - int seed_addr = taddr + (unsigned int)seed_col + (unsigned int)(score_row_base << 16); - tmem_st_x4_f32(seed_addr, seed_zero); - } - } - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(o_first_slice_seeded_addr); - int peer_rank = cta_rank ^ 1; - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(o_first_slice_seeded_addr), "r"(peer_rank) : "memory"); - #pragma unroll 1 - for (unsigned int work_idx = cluster_id; work_idx < total_work_items; work_idx += num_clusters) { - int split_idx = work_idx % (unsigned int)(((1) ? 4 : 1)); - int query_idx = work_idx / (unsigned int)(((1) ? 4 : 1)); - int _max_0 = ((sparse_topk_lens[query_idx] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx] + sparse_topk_lens_offset) : (0)); - int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); - int active_topk = _min_0; - int num_kv_tiles = ((sparse_topk + 128 - 1) / 128 + ((1) ? 4 : 1) - 1) / ((1) ? 4 : 1); - int first_kv_tile = split_idx * num_kv_tiles; - int head_idx = cta_rank * 64 + my_row; - float row_max_val = -CAKE_INF; - float row_sum_val = 0.0f; - if (has_sinks != 0 && head_idx < num_heads && split_idx == 0) { - row_max_val = sinks[head_idx] * 1.4426950408889634f / softmax_scale_log2; - if (n_half == 0) { - row_sum_val = 1.0f; - } - } - int valid_slot = valid_iter & 1; - int valid_full_phase = valid_iter >> 1 & 1; - mbarrier_wait(index_valid_full_addr + (valid_slot) * 8, valid_full_phase); - int valid_word_base = valid_slot * 36 + n_half * 2; - #pragma unroll 1 - for (int tile = 0; tile < num_kv_tiles; tile++) { - int pipeline_tile = tile_cursor + tile; - int score_phase = pipeline_tile & 1; - int s_full_phase = pipeline_tile >> 1 & 1; - mbarrier_wait(s_full_addr + (score_phase) * 8, s_full_phase); - int score_offset = ((score_phase != 0) ? 64 : 0); - int score_base = taddr + (unsigned int)score_offset + (unsigned int)(score_row_base << 16); - float _tmem_load_0[64]; - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]) - : "r"(score_base)); - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]) - : "r"(score_base + 32)); - int valid_cols = active_topk - (first_kv_tile + tile) * 128 - n_half * 64; - if (valid_cols < 0) { - valid_cols = 0; - } - if (valid_cols > 64) { - valid_cols = 64; - } - unsigned int keep_lo = smem_index_valid[valid_word_base + tile * 4]; - unsigned int keep_hi = smem_index_valid[valid_word_base + tile * 4 + 1]; - uint32_t _thresh_mask_0; - { - int _lim_0 = valid_cols; - if (_lim_0 <= 0) { _thresh_mask_0 = 0u; } - else if (_lim_0 >= 32) { _thresh_mask_0 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_thresh_mask_0) : "r"(_lim_0)); - } - } - uint32_t _thresh_mask_1; - { - int _lim_1 = valid_cols - 32; - if (_lim_1 <= 0) { _thresh_mask_1 = 0u; } - else if (_lim_1 >= 32) { _thresh_mask_1 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_thresh_mask_1) : "r"(_lim_1)); - } - } - if (!(_thresh_mask_0 & keep_lo & (1u << 0))) _tmem_load_0[0] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 1))) _tmem_load_0[1] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 2))) _tmem_load_0[2] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 3))) _tmem_load_0[3] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 4))) _tmem_load_0[4] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 5))) _tmem_load_0[5] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 6))) _tmem_load_0[6] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 7))) _tmem_load_0[7] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 8))) _tmem_load_0[8] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 9))) _tmem_load_0[9] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 10))) _tmem_load_0[10] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 11))) _tmem_load_0[11] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 12))) _tmem_load_0[12] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 13))) _tmem_load_0[13] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 14))) _tmem_load_0[14] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 15))) _tmem_load_0[15] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 16))) _tmem_load_0[16] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 17))) _tmem_load_0[17] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 18))) _tmem_load_0[18] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 19))) _tmem_load_0[19] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 20))) _tmem_load_0[20] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 21))) _tmem_load_0[21] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 22))) _tmem_load_0[22] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 23))) _tmem_load_0[23] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 24))) _tmem_load_0[24] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 25))) _tmem_load_0[25] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 26))) _tmem_load_0[26] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 27))) _tmem_load_0[27] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 28))) _tmem_load_0[28] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 29))) _tmem_load_0[29] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 30))) _tmem_load_0[30] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 31))) _tmem_load_0[31] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 0))) _tmem_load_0[32] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 1))) _tmem_load_0[33] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 2))) _tmem_load_0[34] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 3))) _tmem_load_0[35] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 4))) _tmem_load_0[36] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 5))) _tmem_load_0[37] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 6))) _tmem_load_0[38] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 7))) _tmem_load_0[39] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 8))) _tmem_load_0[40] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 9))) _tmem_load_0[41] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 10))) _tmem_load_0[42] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 11))) _tmem_load_0[43] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 12))) _tmem_load_0[44] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 13))) _tmem_load_0[45] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 14))) _tmem_load_0[46] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 15))) _tmem_load_0[47] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 16))) _tmem_load_0[48] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 17))) _tmem_load_0[49] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 18))) _tmem_load_0[50] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 19))) _tmem_load_0[51] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 20))) _tmem_load_0[52] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 21))) _tmem_load_0[53] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 22))) _tmem_load_0[54] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 23))) _tmem_load_0[55] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 24))) _tmem_load_0[56] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 25))) _tmem_load_0[57] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 26))) _tmem_load_0[58] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 27))) _tmem_load_0[59] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 28))) _tmem_load_0[60] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 29))) _tmem_load_0[61] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 30))) _tmem_load_0[62] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 31))) _tmem_load_0[63] = -CAKE_INF; - float2 _reg_reduce_max2_2 = {-CAKE_INF, -CAKE_INF}; - row_max_x32_accum(&_tmem_load_0[0], _reg_reduce_max2_2); - row_max_x32_accum(&_tmem_load_0[32], _reg_reduce_max2_2); - float _tmem_load_0_max = row_max_reduce(_reg_reduce_max2_2); - float _max_1 = max_noftz(_tmem_load_0_max, row_max_val); - float new_max = _max_1; - smem_softmax_exchange[exchange_idx] = new_max; - asm volatile("barrier.sync 3, 128;" ::: "memory"); - float _max_2 = max_noftz(new_max, smem_softmax_exchange[exchange_idx ^ 64]); - new_max = _max_2; - asm volatile("barrier.sync 3, 128;" ::: "memory"); - float no_correction = (((new_max - row_max_val) * softmax_scale_log2 <= 0.0f) ? 1.0f : 0.0f); - float delta = softmax_scale_log2 * (row_max_val - new_max); - float _exp2_0 = approx_exp2(delta); - float exp_delta = _exp2_0; - float acc_scale = ((row_max_val > -CAKE_INF) ? exp_delta : 1.0f); - row_max_val = new_max; - float safe_max = ((new_max == -CAKE_INF) ? 0.0f : new_max); - float max_scaled = safe_max * softmax_scale_log2; - const float2 _fma_b2_3 = {softmax_scale_log2, softmax_scale_log2}; - const float2 _fma_c2_4 = {-max_scaled, -max_scaled}; - #pragma unroll - for (int _lf = 0; _lf < 32; _lf++) - fma_f32x2_inplace(&reinterpret_cast(_tmem_load_0)[_lf], _fma_b2_3, _fma_c2_4); - #pragma unroll - for (int _le = 0; _le < 64; _le++) { - _tmem_load_0[_le] = approx_exp2(_tmem_load_0[_le]); - } - float2 _reg_reduce_sum2_5 = make_float2(0.0f, 0.0f); - softmax_block_sum(&_tmem_load_0[0], &_reg_reduce_sum2_5); - softmax_block_sum(&_tmem_load_0[32], &_reg_reduce_sum2_5); - float _tmem_load_0_sum = _reg_reduce_sum2_5.x + _reg_reduce_sum2_5.y; - row_sum_val = row_sum_val * acc_scale + _tmem_load_0_sum; - int stats_empty_phase = pipeline_tile >> 1 & 1 ^ 1; - mbarrier_wait(stats_empty_addr + (score_phase) * 8, stats_empty_phase); - float meta[4]; - meta[0] = row_sum_val; - meta[1] = row_max_val; - meta[2] = acc_scale; - meta[3] = no_correction; - int meta_addr = taddr + 384 + (unsigned int)(score_phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(head_row_base << 16); - tmem_st_x4_f32(meta_addr, meta); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - mbarrier_arrive(stats_addr + (score_phase) * 8); - int p_stage = pipeline_tile % 2; - int p_empty_phase = pipeline_tile / 2 & 1 ^ 1; - mbarrier_wait(p_empty_addr + (p_stage) * 8, p_empty_phase); - uint32_t _tmem_load_0_bf16[32]; - #pragma unroll - for (int _lp = 0; _lp < 32; _lp++) { - __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_0[_lp*2 + 0], _tmem_load_0[_lp*2+1 + 0])); - _tmem_load_0_bf16[_lp] = *(uint32_t*)&_bf2; - } - int smem_p_stage = p_stage * 2 + n_half; - int p_base_smem = smem_p_addr + (unsigned int)(smem_p_stage * 8192); - #pragma unroll - for (int vec = 0; vec < 8; vec++) { - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_base_smem + (my_row * 128 + vec * 16 ^ (my_row * 128 + vec * 16 >> 7 & 7) << 4))), "r"(_tmem_load_0_bf16[vec * 4]), "r"(_tmem_load_0_bf16[vec * 4 + 1]), "r"(_tmem_load_0_bf16[vec * 4 + 2]), "r"(_tmem_load_0_bf16[vec * 4 + 3]) : "memory"); - } - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(p_full_addr + (unsigned int)(p_stage * 8)), "r"(0) : "memory"); - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(s_empty_addr + (unsigned int)(score_phase * 8)), "r"(0) : "memory"); - } - mbarrier_arrive(index_valid_empty_addr + (valid_slot) * 8); - valid_iter = valid_iter + 1; - tile_cursor = tile_cursor + num_kv_tiles; - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - } - } - // ---- Role: correction_wg ---- - if (warp >= 4 && warp <= 7) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); - { // correction_wg_main - const int correction_dummy = 0; - float softmax_scale_log2_1 = bmm1_scale[0] * 1.4426950408889634f; - float output_scale = bmm2_scale[0]; - const int head_row_base_1 = warp % 2 * 32; - const int output_row_base = warp % 4 * 32; - const int n_half_1 = warp % 4 / 2; - const int my_row_1 = head_row_base_1 + lane; - const int corr_row = head_row_base_1 << 16; - int tile_cursor_1 = 0; - unsigned int _phase_o_full_0 = 0; - #pragma unroll 1 - for (unsigned int work_idx_1 = cluster_id; work_idx_1 < total_work_items; work_idx_1 += num_clusters) { - int split_idx_1 = work_idx_1 % (unsigned int)(((1) ? 4 : 1)); - int query_idx_1 = work_idx_1 / (unsigned int)(((1) ? 4 : 1)); - int num_kv_tiles_1 = ((sparse_topk + 128 - 1) / 128 + ((1) ? 4 : 1) - 1) / ((1) ? 4 : 1); - float final_sum = 0.0f; - float final_max = -CAKE_INF; - #pragma unroll 1 - for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { - int pipeline_tile_1 = tile_cursor_1 + tile_1; - int score_phase_1 = pipeline_tile_1 & 1; - int stats_phase = pipeline_tile_1 >> 1 & 1; - mbarrier_wait(stats_addr + (score_phase_1) * 8, stats_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - int meta_addr_1 = taddr + 384 + (unsigned int)(score_phase_1 * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row; - float _tmem_load_1[4]; - tmem_ld_x4(&_tmem_load_1[0], meta_addr_1); - final_sum = _tmem_load_1[0]; - final_max = _tmem_load_1[1]; - float acc_scale_1 = _tmem_load_1[2]; - float no_correction_1 = _tmem_load_1[3]; - if (tile_1 > 0) { - mbarrier_wait(o_full_addr, _phase_o_full_0); - _phase_o_full_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - int _vote_0 = __all_sync(0xFFFFFFFF, no_correction_1 == 1.0f); - int skip_correction = _vote_0; - if (skip_correction == 0) { - #pragma unroll - for (int local_slice = 0; local_slice < 2; local_slice++) { - #pragma unroll - for (int acc_stage = 0; acc_stage < 2; acc_stage++) { - int o_base = taddr + 128 + (unsigned int)(acc_stage * 128) + (unsigned int)(local_slice * 64) + (unsigned int)(output_row_base << 16); - #pragma unroll - for (int c = 0; c < 64; c += 16) { - float _tmem_load_2[16]; - tmem_ld_x16(&_tmem_load_2[0], o_base + c); - const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; - #pragma unroll - for (int _ls = 0; _ls < 8; _ls++) - mul_f32x2_inplace(&reinterpret_cast(_tmem_load_2)[_ls], _scale2_0); - tmem_st_x16_f32(o_base + c, _tmem_load_2); - } - } - } - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - } - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(o_empty_addr), "r"(0) : "memory"); - } - mbarrier_arrive(stats_empty_addr + (score_phase_1) * 8); - } - mbarrier_wait(o_full_addr, _phase_o_full_0); - _phase_o_full_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - const int exchange_idx_1 = n_half_1 * 64 + my_row_1; - smem_epilogue_exchange[exchange_idx_1] = final_sum; - asm volatile("barrier.sync 4, 128;" ::: "memory"); - final_sum = final_sum + smem_epilogue_exchange[exchange_idx_1 ^ 64]; - asm volatile("barrier.sync 4, 128;" ::: "memory"); - float _rcp_0 = approx_rcp(final_sum); - float inv_sum = ((final_sum > 0.0f) ? _rcp_0 : 0.0f); - int head_idx_1 = cta_rank * 64 + my_row_1; - int output_offset = ((1) ? ((query_idx_1 * num_heads + head_idx_1) * ((1) ? 4 : 1) + split_idx_1) * 512 : (query_idx_1 * num_heads + head_idx_1) * 512); - if (n_half_1 == 0 && head_idx_1 < num_heads) { - int stat_offset = (query_idx_1 * num_heads + head_idx_1) * ((1) ? 4 : 1) + split_idx_1; - float _log2_0; - asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(final_sum)); - partial_lse[stat_offset] = ((final_sum > 0.0f) ? final_max * softmax_scale_log2_1 + _log2_0 : -CAKE_INF); - } - #pragma unroll - for (int acc_stage_1 = 0; acc_stage_1 < 2; acc_stage_1++) { - #pragma unroll - for (int local_slice_1 = 0; local_slice_1 < 2; local_slice_1++) { - int logical_slice = acc_stage_1 * 2 * 2 + n_half_1 * 2 + local_slice_1; - int o_base_1 = taddr + 128 + (unsigned int)(acc_stage_1 * 128) + (unsigned int)(local_slice_1 * 64) + (unsigned int)(output_row_base << 16); - #pragma unroll - for (int c_1 = 0; c_1 < 64; c_1 += 32) { - float _tmem_load_3[32]; - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]) - : "r"(o_base_1 + c_1)); - int gmem_base = output_offset + logical_slice * 64 + c_1; - #pragma unroll - for (int j = 0; j < 32; j += 16) { - { - const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; - #if __CUDA_ARCH__ >= 1000 - #pragma unroll - for (int _ps = 0; _ps < 8; _ps++) - mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_3[j])[_ps], _prescale2_1); - #else - #pragma unroll - for (int _ps = 0; _ps < 16; _ps++) - _tmem_load_3[j + _ps] *= inv_sum * output_scale; - #endif - { - __nv_bfloat162 _pk0 = __floats2bfloat162_rn(_tmem_load_3[j + 0], _tmem_load_3[j + 1]); - unsigned _pk_u0 = *reinterpret_cast(&_pk0); - __nv_bfloat162 _pk1 = __floats2bfloat162_rn(_tmem_load_3[j + 2], _tmem_load_3[j + 3]); - unsigned _pk_u1 = *reinterpret_cast(&_pk1); - __nv_bfloat162 _pk2 = __floats2bfloat162_rn(_tmem_load_3[j + 4], _tmem_load_3[j + 5]); - unsigned _pk_u2 = *reinterpret_cast(&_pk2); - __nv_bfloat162 _pk3 = __floats2bfloat162_rn(_tmem_load_3[j + 6], _tmem_load_3[j + 7]); - unsigned _pk_u3 = *reinterpret_cast(&_pk3); - __nv_bfloat162 _pk4 = __floats2bfloat162_rn(_tmem_load_3[j + 8], _tmem_load_3[j + 9]); - unsigned _pk_u4 = *reinterpret_cast(&_pk4); - __nv_bfloat162 _pk5 = __floats2bfloat162_rn(_tmem_load_3[j + 10], _tmem_load_3[j + 11]); - unsigned _pk_u5 = *reinterpret_cast(&_pk5); - __nv_bfloat162 _pk6 = __floats2bfloat162_rn(_tmem_load_3[j + 12], _tmem_load_3[j + 13]); - unsigned _pk_u6 = *reinterpret_cast(&_pk6); - __nv_bfloat162 _pk7 = __floats2bfloat162_rn(_tmem_load_3[j + 14], _tmem_load_3[j + 15]); - unsigned _pk_u7 = *reinterpret_cast(&_pk7); - asm volatile( - "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" - :: "l"((void*)(&((__nv_bfloat16*)(O + (gmem_base + j)))[0])), "r"(_pk_u0), "r"(_pk_u1), "r"(_pk_u2), "r"(_pk_u3), "r"(_pk_u4), "r"(_pk_u5), "r"(_pk_u6), "r"(_pk_u7) : "memory"); - } - } - } - } - } - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(o_empty_addr), "r"(0) : "memory"); - tile_cursor_1 = tile_cursor_1 + num_kv_tiles_1; - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - } - } - // ---- Role: mma_warp ---- - if (warp == 8) { - { // mma_warp_main - const int mma_dummy = 0; - unsigned int k_stage = 0; - unsigned int v_stage = 0; - int tile_cursor_2 = 0; - unsigned int _phase_o_first_slice_seeded_0 = 0; - mbarrier_wait(o_first_slice_seeded_addr, _phase_o_first_slice_seeded_0); - _phase_o_first_slice_seeded_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - unsigned int _phase_q_full_0 = 0; - unsigned int _phase_k_full = 0; - unsigned int _phase_o_empty_0 = 1; - unsigned int _phase_v_full = 0; - #pragma unroll 1 - for (unsigned int work_idx_2 = cluster_id; work_idx_2 < total_work_items; work_idx_2 += num_clusters) { - int num_kv_tiles_2 = ((sparse_topk + 128 - 1) / 128 + ((1) ? 4 : 1) - 1) / ((1) ? 4 : 1); - if (cta_rank == 0) { - mbarrier_wait(q_full_addr, _phase_q_full_0); - _phase_q_full_0 ^= 1; - int first_pv = 1; - #pragma unroll 1 - for (int tile_2 = 0; tile_2 < num_kv_tiles_2; tile_2++) { - int pipeline_tile_2 = tile_cursor_2 + tile_2; - int score_phase_2 = pipeline_tile_2 & 1; - int s_empty_phase = pipeline_tile_2 >> 1 & 1 ^ 1; - mbarrier_wait(s_empty_addr + (score_phase_2) * 8, s_empty_phase); - int score_col = ((score_phase_2 != 0) ? 64 : 0); - #pragma unroll - for (int qk_stage = 0; qk_stage < 4; qk_stage++) { - mbarrier_wait(k_full_addr + (k_stage) * 8, _phase_k_full); - int _mma_a_lo_0 = (((smem_q_addr) >> 4) & 0x3FFF) + (qk_stage * 2) * 512; - int _mma_b_lo_0 = (((smem_k_addr) >> 4) & 0x3FFF) + (k_stage) * 1024; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 136316048;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (score_col))), "r"(((qk_stage == 0) ? 0 : 1))); - int _mma_a_lo_1 = (((smem_q_addr) >> 4) & 0x3FFF) + (qk_stage * 2 + 1) * 512; - int _mma_b_lo_1 = (((smem_k_addr + 8192) >> 4) & 0x3FFF) + (k_stage) * 1024; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 136316048;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); - elect_commit_cg2_multicast(k_empty_addr + (k_stage) * 8, (uint16_t)(3)); - k_stage += 1; - if (k_stage == 3) { k_stage = 0; _phase_k_full ^= 1; } - } - elect_commit_cg2_multicast(s_full_addr + (score_phase_2) * 8, (uint16_t)(3)); - if (tile_2 > 0) { - int prev_pipeline_tile = pipeline_tile_2 - 1; - int p_stage_1 = prev_pipeline_tile % 2; - int p_full_phase = prev_pipeline_tile / 2 & 1; - mbarrier_wait(o_empty_addr, _phase_o_empty_0); - _phase_o_empty_0 ^= 1; - mbarrier_wait(p_full_addr + (p_stage_1) * 8, p_full_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - #pragma unroll - for (int ps = 0; ps < 2; ps++) { - int smem_p_stage_1 = p_stage_1 * 2 + ps; - mbarrier_wait(v_full_addr + (v_stage) * 8, _phase_v_full); - #pragma unroll - for (int acc_stage_2 = 0; acc_stage_2 < 2; acc_stage_2++) { - int output_col = 128 + acc_stage_2 * 128; - if (acc_stage_2 == 0) { - int _mma_a_lo_2 = (((smem_p_addr) >> 4) & 0x3FFF) + (smem_p_stage_1) * 512; - int _mma_b_lo_2 = ((((smem_v_lo_addr) >> 4) & 0x3FFF) | 0x2000000) + (v_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 138478736;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (output_col))), "r"(((((ps == 0) ? first_pv : 0)) ? 0 : 1))); - } else { - int _mma_a_lo_3 = (((smem_p_addr) >> 4) & 0x3FFF) + (smem_p_stage_1) * 512; - int _mma_b_lo_3 = ((((smem_v_hi_addr) >> 4) & 0x3FFF) | 0x2000000) + (v_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 138478736;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"((tmem_tmem_scratch + (output_col))), "r"(((((ps == 0) ? first_pv : 0)) ? 0 : 1))); - } - } - elect_commit_cg2_multicast(v_empty_addr + (v_stage) * 8, (uint16_t)(3)); - v_stage += 1; - if (v_stage == 2) { v_stage = 0; _phase_v_full ^= 1; } - } - first_pv = 0; - elect_commit_cg2_multicast(p_empty_addr + (p_stage_1) * 8, (uint16_t)(3)); - elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); - } - } - int last_pipeline_tile = tile_cursor_2 + num_kv_tiles_2 - 1; - int last_p_stage = last_pipeline_tile % 2; - int last_p_phase = last_pipeline_tile / 2 & 1; - mbarrier_wait(o_empty_addr, _phase_o_empty_0); - _phase_o_empty_0 ^= 1; - mbarrier_wait(p_full_addr + (last_p_stage) * 8, last_p_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - #pragma unroll - for (int ps_1 = 0; ps_1 < 2; ps_1++) { - int smem_p_stage_2 = last_p_stage * 2 + ps_1; - mbarrier_wait(v_full_addr + (v_stage) * 8, _phase_v_full); - #pragma unroll - for (int acc_stage_3 = 0; acc_stage_3 < 2; acc_stage_3++) { - int output_col_1 = 128 + acc_stage_3 * 128; - if (acc_stage_3 == 0) { - int _mma_a_lo_4 = (((smem_p_addr) >> 4) & 0x3FFF) + (smem_p_stage_2) * 512; - int _mma_b_lo_4 = ((((smem_v_lo_addr) >> 4) & 0x3FFF) | 0x2000000) + (v_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 138478736;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"((tmem_tmem_scratch + (output_col_1))), "r"(((((ps_1 == 0) ? first_pv : 0)) ? 0 : 1))); - } else { - int _mma_a_lo_5 = (((smem_p_addr) >> 4) & 0x3FFF) + (smem_p_stage_2) * 512; - int _mma_b_lo_5 = ((((smem_v_hi_addr) >> 4) & 0x3FFF) | 0x2000000) + (v_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 138478736;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_5), "r"(_mma_b_lo_5), "r"((tmem_tmem_scratch + (output_col_1))), "r"(((((ps_1 == 0) ? first_pv : 0)) ? 0 : 1))); - } - } - elect_commit_cg2_multicast(v_empty_addr + (v_stage) * 8, (uint16_t)(3)); - v_stage += 1; - if (v_stage == 2) { v_stage = 0; _phase_v_full ^= 1; } - } - elect_commit_cg2_multicast(p_empty_addr + (last_p_stage) * 8, (uint16_t)(3)); - elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); - elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); - } - tile_cursor_2 = tile_cursor_2 + num_kv_tiles_2; - } - if (cta_rank == 0) { - #pragma unroll - for (int tail_offset = 0; tail_offset < 2; tail_offset++) { - int tail_tile = tile_cursor_2 + tail_offset; - int tail_stage = tail_tile & 1; - int tail_phase = tail_tile >> 1 & 1 ^ 1; - mbarrier_wait(s_empty_addr + (tail_stage) * 8, tail_phase); - } - mbarrier_wait(o_empty_addr, _phase_o_empty_0); - _phase_o_empty_0 ^= 1; - } - asm volatile("tcgen05.fence::before_thread_sync;"); - int peer_rank_1 = cta_rank ^ 1; - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(tmem_dealloc_peer_addr), "r"(peer_rank_1) : "memory"); - mbarrier_wait(tmem_dealloc_peer_addr, 0); - int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); - asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); - } - } - // ---- Role: load_wg ---- - if (warp >= 9 && warp <= 12) { - { // load_wg_main - const int load_dummy = 0; - const int load_warp_rank = warp - 9; - unsigned int k_stage_1 = 0; - unsigned int v_stage_1 = 0; - unsigned int _phase_q_empty_0 = 1; - unsigned int _phase_k_empty = 1; - unsigned int _phase_v_empty = 1; - #pragma unroll 1 - for (unsigned int work_idx_3 = cluster_id; work_idx_3 < total_work_items; work_idx_3 += num_clusters) { - int split_idx_2 = work_idx_3 % (unsigned int)(((1) ? 4 : 1)); - int query_idx_2 = work_idx_3 / (unsigned int)(((1) ? 4 : 1)); - int num_kv_tiles_3 = ((sparse_topk + 128 - 1) / 128 + ((1) ? 4 : 1) - 1) / ((1) ? 4 : 1); - int first_kv_tile_1 = split_idx_2 * num_kv_tiles_3; - mbarrier_wait(q_empty_addr, _phase_q_empty_0); - _phase_q_empty_0 ^= 1; - if (load_warp_rank == 0) { - if (cta_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(131072)) : "memory"); - } - } - if (elect_sync()) { - #pragma unroll - for (int q_stage = 0; q_stage < 8; q_stage++) { - tma_4d_gmem2smem_cta2(smem_q_addr + (unsigned int)(q_stage * 8192), (&tmap_q), 0, cta_rank * 64, q_stage, query_idx_2, ((q_full_addr) & 0xFEFFFFFF)); - } - } - } - asm volatile("barrier.sync 2, 160;" ::: "memory"); - #pragma unroll - for (int qk_stage_1 = 0; qk_stage_1 < 4; qk_stage_1++) { - mbarrier_wait(k_empty_addr + (k_stage_1) * 8, _phase_k_empty); - if (load_warp_rank == 0) { - if (cta_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); - } - } - } - int group = load_warp_rank * 4 + lane; - if (lane < 4) { - int index_offset = cta_rank * 64 + group * 4; - int raw_rows[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows[0])), "=r"(*reinterpret_cast(&raw_rows[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows[(0) + 3])) - : "r"(smem_indices_addr + (unsigned int)(index_offset * 4))); - if (first_kv_tile_1 == 0) { - tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group * 512), (&tmap_swa_k), qk_stage_1 * 128, ((raw_rows[0] >= 0) ? raw_rows[0] : ((1) ? group * 4 : 0)), ((raw_rows[1] >= 0) ? raw_rows[1] : ((1) ? group * 4 + 1 : 0)), ((raw_rows[2] >= 0) ? raw_rows[2] : ((1) ? group * 4 + 2 : 0)), ((raw_rows[3] >= 0) ? raw_rows[3] : ((1) ? group * 4 + 3 : 0)), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group * 512) + 8192, (&tmap_swa_k), qk_stage_1 * 128 + 64, ((raw_rows[0] >= 0) ? raw_rows[0] : ((1) ? group * 4 : 0)), ((raw_rows[1] >= 0) ? raw_rows[1] : ((1) ? group * 4 + 1 : 0)), ((raw_rows[2] >= 0) ? raw_rows[2] : ((1) ? group * 4 + 2 : 0)), ((raw_rows[3] >= 0) ? raw_rows[3] : ((1) ? group * 4 + 3 : 0)), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - } else { - tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group * 512), (&tmap_compressed_k), qk_stage_1 * 128, ((raw_rows[0] >= 0) ? raw_rows[0] : ((1) ? group * 4 : 0)), ((raw_rows[1] >= 0) ? raw_rows[1] : ((1) ? group * 4 + 1 : 0)), ((raw_rows[2] >= 0) ? raw_rows[2] : ((1) ? group * 4 + 2 : 0)), ((raw_rows[3] >= 0) ? raw_rows[3] : ((1) ? group * 4 + 3 : 0)), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group * 512) + 8192, (&tmap_compressed_k), qk_stage_1 * 128 + 64, ((raw_rows[0] >= 0) ? raw_rows[0] : ((1) ? group * 4 : 0)), ((raw_rows[1] >= 0) ? raw_rows[1] : ((1) ? group * 4 + 1 : 0)), ((raw_rows[2] >= 0) ? raw_rows[2] : ((1) ? group * 4 + 2 : 0)), ((raw_rows[3] >= 0) ? raw_rows[3] : ((1) ? group * 4 + 3 : 0)), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - } - } - k_stage_1 += 1; - if (k_stage_1 == 3) { k_stage_1 = 0; _phase_k_empty ^= 1; } - } - #pragma unroll 1 - for (int tile_3 = 1; tile_3 < num_kv_tiles_3; tile_3++) { - #pragma unroll - for (int qk_stage_2 = 0; qk_stage_2 < 4; qk_stage_2++) { - mbarrier_wait(k_empty_addr + (k_stage_1) * 8, _phase_k_empty); - if (load_warp_rank == 0) { - if (cta_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); - } - } - } - int group_1 = load_warp_rank * 4 + lane; - if (lane < 4) { - int index_offset_1 = tile_3 * 128 + cta_rank * 64 + group_1 * 4; - int raw_rows_1[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_1[0])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 3])) - : "r"(smem_indices_addr + (unsigned int)(index_offset_1 * 4))); - tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group_1 * 512), (&tmap_compressed_k), qk_stage_2 * 128, ((raw_rows_1[0] >= 0) ? raw_rows_1[0] : ((1) ? group_1 * 4 : 0)), ((raw_rows_1[1] >= 0) ? raw_rows_1[1] : ((1) ? group_1 * 4 + 1 : 0)), ((raw_rows_1[2] >= 0) ? raw_rows_1[2] : ((1) ? group_1 * 4 + 2 : 0)), ((raw_rows_1[3] >= 0) ? raw_rows_1[3] : ((1) ? group_1 * 4 + 3 : 0)), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group_1 * 512) + 8192, (&tmap_compressed_k), qk_stage_2 * 128 + 64, ((raw_rows_1[0] >= 0) ? raw_rows_1[0] : ((1) ? group_1 * 4 : 0)), ((raw_rows_1[1] >= 0) ? raw_rows_1[1] : ((1) ? group_1 * 4 + 1 : 0)), ((raw_rows_1[2] >= 0) ? raw_rows_1[2] : ((1) ? group_1 * 4 + 2 : 0)), ((raw_rows_1[3] >= 0) ? raw_rows_1[3] : ((1) ? group_1 * 4 + 3 : 0)), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - } - k_stage_1 += 1; - if (k_stage_1 == 3) { k_stage_1 = 0; _phase_k_empty ^= 1; } - } - int prev_tile = tile_3 - 1; - #pragma unroll - for (int ps_2 = 0; ps_2 < 2; ps_2++) { - mbarrier_wait(v_empty_addr + (v_stage_1) * 8, _phase_v_empty); - if (load_warp_rank == 0) { - if (cta_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(65536)) : "memory"); - } - } - } - int group_2 = load_warp_rank * 4 + lane; - if (lane < 4) { - int index_offset_2 = prev_tile * 128 + ps_2 * 64 + group_2 * 4; - int raw_rows_2[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_2[0])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 3])) - : "r"(smem_indices_addr + (unsigned int)(index_offset_2 * 4))); - #pragma unroll - for (int acc_stage_4 = 0; acc_stage_4 < 2; acc_stage_4++) { - int dst_v = ((acc_stage_4 == 0) ? smem_v_lo_addr + v_stage_1 * 32768 : smem_v_hi_addr + v_stage_1 * 32768); - if (first_kv_tile_1 + prev_tile == 0) { - tma_gather4_gmem2smem_mc_cta2(dst_v + group_2 * 512, (&tmap_swa_v), acc_stage_4 * 128 * 2 + cta_rank * 128, ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : ((1) ? group_2 * 4 : 0)), ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : ((1) ? group_2 * 4 + 1 : 0)), ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : ((1) ? group_2 * 4 + 2 : 0)), ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : ((1) ? group_2 * 4 + 3 : 0)), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - tma_gather4_gmem2smem_mc_cta2(dst_v + group_2 * 512 + 8192, (&tmap_swa_v), acc_stage_4 * 128 * 2 + cta_rank * 128 + 64, ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : ((1) ? group_2 * 4 : 0)), ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : ((1) ? group_2 * 4 + 1 : 0)), ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : ((1) ? group_2 * 4 + 2 : 0)), ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : ((1) ? group_2 * 4 + 3 : 0)), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - } else { - tma_gather4_gmem2smem_mc_cta2(dst_v + group_2 * 512, (&tmap_compressed_v), acc_stage_4 * 128 * 2 + cta_rank * 128, ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : ((1) ? group_2 * 4 : 0)), ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : ((1) ? group_2 * 4 + 1 : 0)), ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : ((1) ? group_2 * 4 + 2 : 0)), ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : ((1) ? group_2 * 4 + 3 : 0)), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - tma_gather4_gmem2smem_mc_cta2(dst_v + group_2 * 512 + 8192, (&tmap_compressed_v), acc_stage_4 * 128 * 2 + cta_rank * 128 + 64, ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : ((1) ? group_2 * 4 : 0)), ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : ((1) ? group_2 * 4 + 1 : 0)), ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : ((1) ? group_2 * 4 + 2 : 0)), ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : ((1) ? group_2 * 4 + 3 : 0)), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - } - } - } - v_stage_1 += 1; - if (v_stage_1 == 2) { v_stage_1 = 0; _phase_v_empty ^= 1; } - } - } - int last_tile = num_kv_tiles_3 - 1; - #pragma unroll - for (int ps_3 = 0; ps_3 < 2; ps_3++) { - mbarrier_wait(v_empty_addr + (v_stage_1) * 8, _phase_v_empty); - if (load_warp_rank == 0) { - if (cta_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(65536)) : "memory"); - } - } - } - int group_3 = load_warp_rank * 4 + lane; - if (lane < 4) { - int index_offset_3 = last_tile * 128 + ps_3 * 64 + group_3 * 4; - int raw_rows_3[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_3[0])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 3])) - : "r"(smem_indices_addr + (unsigned int)(index_offset_3 * 4))); - #pragma unroll - for (int acc_stage_5 = 0; acc_stage_5 < 2; acc_stage_5++) { - int dst_v_1 = ((acc_stage_5 == 0) ? smem_v_lo_addr + v_stage_1 * 32768 : smem_v_hi_addr + v_stage_1 * 32768); - if (first_kv_tile_1 + last_tile == 0) { - tma_gather4_gmem2smem_mc_cta2(dst_v_1 + group_3 * 512, (&tmap_swa_v), acc_stage_5 * 128 * 2 + cta_rank * 128, ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : ((1) ? group_3 * 4 : 0)), ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : ((1) ? group_3 * 4 + 1 : 0)), ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : ((1) ? group_3 * 4 + 2 : 0)), ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : ((1) ? group_3 * 4 + 3 : 0)), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - tma_gather4_gmem2smem_mc_cta2(dst_v_1 + group_3 * 512 + 8192, (&tmap_swa_v), acc_stage_5 * 128 * 2 + cta_rank * 128 + 64, ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : ((1) ? group_3 * 4 : 0)), ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : ((1) ? group_3 * 4 + 1 : 0)), ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : ((1) ? group_3 * 4 + 2 : 0)), ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : ((1) ? group_3 * 4 + 3 : 0)), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - } else { - tma_gather4_gmem2smem_mc_cta2(dst_v_1 + group_3 * 512, (&tmap_compressed_v), acc_stage_5 * 128 * 2 + cta_rank * 128, ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : ((1) ? group_3 * 4 : 0)), ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : ((1) ? group_3 * 4 + 1 : 0)), ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : ((1) ? group_3 * 4 + 2 : 0)), ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : ((1) ? group_3 * 4 + 3 : 0)), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - tma_gather4_gmem2smem_mc_cta2(dst_v_1 + group_3 * 512 + 8192, (&tmap_compressed_v), acc_stage_5 * 128 * 2 + cta_rank * 128 + 64, ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : ((1) ? group_3 * 4 : 0)), ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : ((1) ? group_3 * 4 + 1 : 0)), ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : ((1) ? group_3 * 4 + 2 : 0)), ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : ((1) ? group_3 * 4 + 3 : 0)), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - } - } - } - v_stage_1 += 1; - if (v_stage_1 == 2) { v_stage_1 = 0; _phase_v_empty ^= 1; } - } - asm volatile("barrier.sync 2, 160;" ::: "memory"); - } - mbarrier_wait(q_empty_addr, _phase_q_empty_0); - _phase_q_empty_0 ^= 1; - } - } - // ---- Role: index_warp ---- - if (warp == 13) { - { // index_warp_main - const int index_dummy = 0; - int valid_iter_1 = 0; - #pragma unroll 1 - for (unsigned int work_idx_4 = cluster_id; work_idx_4 < total_work_items; work_idx_4 += num_clusters) { - int split_idx_3 = work_idx_4 % (unsigned int)(((1) ? 4 : 1)); - int query_idx_3 = work_idx_4 / (unsigned int)(((1) ? 4 : 1)); - int owner_col_base = split_idx_3 * (4 / ((1) ? 4 : 1)) * 128; - int* swa_row = swa_indices + (query_idx_3 * swa_index_stride); - int* compressed_row = compressed_indices + (query_idx_3 * compressed_index_stride); - int valid_slot_1 = valid_iter_1 & 1; - int valid_empty_phase = valid_iter_1 >> 1 & 1 ^ 1; - mbarrier_wait(index_valid_empty_addr + (valid_slot_1) * 8, valid_empty_phase); - #pragma unroll - for (int index_pass = 0; index_pass < 4 / ((1) ? 4 : 1); index_pass++) { - int index_offset_4 = index_pass * 128 + lane * 4; - int combined_col = owner_col_base + index_offset_4; - int* index_src = ((combined_col < 128) ? swa_row + combined_col : compressed_row + (combined_col - 128)); - int values[4]; - #pragma unroll - for (int i = 0; i < 4; i++) { - values[i] = -1; - } - if (combined_col + 4 <= sparse_topk) { - int _vec_load_0[4]; - { - const int4* _ivptr_0 = reinterpret_cast(index_src + 0); - int4 _ivld_0; - _ivld_0 = *_ivptr_0; - _vec_load_0[0 + 0] = _ivld_0.x; - _vec_load_0[0 + 1] = _ivld_0.y; - _vec_load_0[0 + 2] = _ivld_0.z; - _vec_load_0[0 + 3] = _ivld_0.w; - } - #pragma unroll - for (int i_1 = 0; i_1 < 4; i_1++) { - values[i_1] = _vec_load_0[i_1]; - } - } - if (combined_col < sparse_topk && combined_col + 4 > sparse_topk) { - #pragma unroll - for (int i_2 = 0; i_2 < 4; i_2++) { - if (combined_col + i_2 < sparse_topk) { - values[i_2] = index_src[i_2]; - } - } - } - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(smem_indices_addr + (unsigned int)(index_offset_4 * 4)), "r"(values[0]), "r"(values[1]), "r"(values[2]), "r"(values[3]) : "memory"); - unsigned int valid_word = (((values[0] >= 0) ? 1 : 0) | ((values[1] >= 0) ? 1 : 0) << 1 | ((values[2] >= 0) ? 1 : 0) << 2 | ((values[3] >= 0) ? 1 : 0) << 3) << (lane & 7) * 4; - unsigned int _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, valid_word, 1); - valid_word = valid_word | _shfl_xor_0; - unsigned int _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, valid_word, 2); - valid_word = valid_word | _shfl_xor_1; - unsigned int _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, valid_word, 4); - valid_word = valid_word | _shfl_xor_2; - if ((lane & 7) == 0) { - smem_index_valid[valid_slot_1 * 36 + index_pass * 4 + (lane >> 3)] = valid_word; - } - } - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(index_valid_full_addr + (valid_slot_1) * 8); - asm volatile("barrier.sync 2, 160;" ::: "memory"); - asm volatile("barrier.sync 2, 160;" ::: "memory"); - valid_iter_1 = valid_iter_1 + 1; - } - } - } - // ---- Role: empty_warps ---- - if (warp >= 14 && warp <= 15) { - { // empty_warps_main - const int empty_dummy = 0; - #pragma unroll 1 - for (unsigned int work_idx_5 = cluster_id; work_idx_5 < total_work_items; work_idx_5 += num_clusters) { - asm volatile("tcgen05.fence::after_thread_sync;"); - } - } - } - - // Cleanup -} - -} // extern "C" diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_family_068c73b21d7cbcd958b8_binding.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_family_068c73b21d7cbcd958b8_binding.cu deleted file mode 100644 index 4515647e55f..00000000000 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_family_068c73b21d7cbcd958b8_binding.cu +++ /dev/null @@ -1,652 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -// Generated by cake.runtime.ffi_dispatch — do not edit. -// tvm-ffi dispatch module for kernel family 'cake_dsv4_family_068c73b21d7cbcd958b8'. -// dispatch-policy-revision: unpinned -#include -#include -#include -#include - -#include "tvm_ffi_utils.h" - -#include -#include -#include -#include -#include -#include -#include -#include - -extern "C" __global__ void kernel_cake_dsv4_d906a03bd66a7488fc29(__nv_bfloat16* __restrict__ partial_O, float* __restrict__ partial_lse, __nv_bfloat16* __restrict__ O, int num_q_heads, int num_split); -extern "C" __global__ void kernel_cake_dsv4_ba9ed6a6233761682cc3(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); - -namespace cake_host_shim_3cfd00f6b2ce7022 { - -using tvm::ffi::Optional; -using tvm::ffi::TensorView; - -class ScopedCudaDevice { - public: - explicit ScopedCudaDevice(int device_id) { - cudaError_t error = cudaGetDevice(&previous_device_); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed before host-shim launch: cudaError=" - << static_cast(error); - if (previous_device_ != device_id) { - error = cudaSetDevice(device_id); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaSetDevice failed before host-shim launch for cuda:" - << device_id << ": cudaError=" << static_cast(error); - restore_ = true; - } - } - - ScopedCudaDevice(const ScopedCudaDevice&) = delete; - ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; - - ~ScopedCudaDevice() noexcept { - if (restore_) { - (void)cudaSetDevice(previous_device_); - } - } - - private: - int previous_device_ = -1; - bool restore_ = false; -}; - -inline int64_t CakeDeviceMultiprocessorCount(int device_id) { - constexpr int kMaxCachedCudaDevices = 64; - TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) - << "physical-SM-count cache does not cover cuda:" << device_id; - static std::atomic count_by_device[kMaxCachedCudaDevices]{}; - int cached = count_by_device[device_id].load(std::memory_order_acquire); - if (cached > 0) return cached; - - int count = 0; - cudaError_t error = cudaDeviceGetAttribute( - &count, cudaDevAttrMultiProcessorCount, device_id); - TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) - << "querying multiProcessorCount failed for cuda:" << device_id - << ": cudaError=" << static_cast(error) << ", count=" << count; - // Concurrent first launches may repeat the immutable device query, but all - // publication is atomic and every later hot-path lookup is one acquire load. - count_by_device[device_id].store(count, std::memory_order_release); - return count; -} - -inline void CheckCudaTensor(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) - << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; -} - -inline void CheckSameCudaDevice( - const TensorView& t, - const TensorView& reference, - const char* name, - const char* reference_name) { - TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) - << name << " must be on the same CUDA device as " << reference_name - << ": got cuda:" << t.device().device_id - << " versus cuda:" << reference.device().device_id; -} - -inline void CheckCurrentCudaDevice( - const TensorView& reference, - const char* reference_name) { - int current_device = -1; - cudaError_t error = cudaGetDevice(¤t_device); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed while validating " << reference_name - << ": cudaError=" << static_cast(error); - TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) - << "current CUDA device must match " << reference_name - << ": current=cuda:" << current_device - << ", tensor=cuda:" << reference.device().device_id; -} - -inline void CheckContiguous(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; -} - -inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { - DLDataType d = t.dtype(); - TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) - << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits - << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits - << ", lanes=" << (int)d.lanes << ")"; -} - -// A logical axis.outer(trailing) folds every source dim above the trailing -// dimensions. Shape products are independent of physical strides, so verify -// the leading dimensions form one dense row-major chain instead of inventing -// a "folded stride". The descriptor reads its exact adjacent physical step -// separately through stride[-(trailing + 1)]. -inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { - TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) - << name << " cannot fold leading dimensions above " << trailing - << " trailing dims from ndim=" << t.ndim(); - int outer_last = t.ndim() - trailing - 1; - if (outer_last <= 0) { - return; - } - int64_t step = t.stride(outer_last); - TVM_FFI_CHECK(step > 0, ValueError) - << name << " physical strides must be positive"; - int64_t expected = step; - for (int axis = outer_last - 1; axis >= 0; --axis) { - expected *= t.size(axis + 1); - if (t.size(axis) > 1) { - TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) - << name << " leading dims are not physically foldable above " << trailing - << " trailing dims: stride(" << axis << ")=" << t.stride(axis) - << ", expected " << expected; - } - } -} - -namespace variant_cake_dsv4_ba9ed6a6233761682cc3_sm_100a_f23ac2823b80 { - -// 4D TMA descriptor for buffer 'tmap_q' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 3, ValueError) - << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - int64_t d2 = t.size(t.ndim() - 2); - int64_t d3 = t.size(t.ndim() - 3); - TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) - << "TMA source 'tmap_q' trailing dims must be positive"; - TVM_FFI_CHECK(d1 % 64 == 0, ValueError) - << "TMA source 'tmap_q' extent " << d1 - << " must divide exactly by " << 64; - uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(d3)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) - << "TMA box (64, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; - int64_t carrier_stride_0 = d1; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; - int64_t carrier_stride_1 = 64; - TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; - TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; - TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; - int64_t carrier_stride_2 = (d2 * d1); - TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; - TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; - TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; - uint64_t global_strides[3] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - (uint64_t)((carrier_stride_1 * 16) / 8), - (uint64_t)((carrier_stride_2 * 16) / 8), - }; - uint32_t box_dim[4] = {64u, 64u, 1u, 1u}; - uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_swa_k' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_swa_k(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_swa_k' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_swa_k' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_swa_k' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_swa_k"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_swa_k' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_swa_k' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_k'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_swa_k') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_compressed_k' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_compressed_k(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_compressed_k' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_compressed_k' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_compressed_k' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_compressed_k"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_compressed_k' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_compressed_k' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_k'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_k') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_swa_v' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_swa_v(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_swa_v' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_swa_v' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_swa_v' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_swa_v"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_swa_v' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_swa_v' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_v'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_swa_v') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_compressed_v' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_compressed_v(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_compressed_v' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_compressed_v' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_compressed_v' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_compressed_v"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_compressed_v' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_compressed_v' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_v'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_v') failed: CUresult=" << (int)r; - return tm; -} - -void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_compressed_k, TensorView arg_tmap_swa_v, TensorView arg_tmap_compressed_v, TensorView arg_O, TensorView arg_partial_lse, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_num_query_tokens, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { - DLDevice dev = arg_tmap_q.device(); - ScopedCudaDevice device_guard(dev.device_id); - CheckCudaTensor(arg_tmap_q, "tmap_q"); - CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); - CheckContiguous(arg_tmap_q, "tmap_q"); - CheckCudaTensor(arg_tmap_swa_k, "tmap_swa_k"); - CheckDtype(arg_tmap_swa_k, "tmap_swa_k", 4, 16, 1); - CheckCudaTensor(arg_tmap_compressed_k, "tmap_compressed_k"); - CheckDtype(arg_tmap_compressed_k, "tmap_compressed_k", 4, 16, 1); - CheckCudaTensor(arg_tmap_swa_v, "tmap_swa_v"); - CheckDtype(arg_tmap_swa_v, "tmap_swa_v", 4, 16, 1); - CheckCudaTensor(arg_tmap_compressed_v, "tmap_compressed_v"); - CheckDtype(arg_tmap_compressed_v, "tmap_compressed_v", 4, 16, 1); - CheckCudaTensor(arg_O, "O"); - CheckDtype(arg_O, "O", 4, 16, 1); - CheckContiguous(arg_O, "O"); - CheckCudaTensor(arg_partial_lse, "partial_lse"); - CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); - CheckContiguous(arg_partial_lse, "partial_lse"); - CheckCudaTensor(arg_swa_indices, "swa_indices"); - CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); - CheckContiguous(arg_swa_indices, "swa_indices"); - CheckCudaTensor(arg_compressed_indices, "compressed_indices"); - CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); - CheckContiguous(arg_compressed_indices, "compressed_indices"); - CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); - CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckCudaTensor(arg_sinks, "sinks"); - CheckDtype(arg_sinks, "sinks", 2, 32, 1); - CheckContiguous(arg_sinks, "sinks"); - CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); - CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); - CheckContiguous(arg_bmm1_scale, "bmm1_scale"); - CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); - CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); - CheckContiguous(arg_bmm2_scale, "bmm2_scale"); - TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) - << "scalar 'num_heads' value " << arg_num_heads - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_num_query_tokens >= -2147483648LL && arg_num_query_tokens <= 2147483647LL, ValueError) - << "scalar 'num_query_tokens' value " << arg_num_query_tokens - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) - << "scalar 'swa_index_stride' value " << arg_swa_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) - << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) - << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) - << "scalar 'sparse_topk' value " << arg_sparse_topk - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) - << "scalar 'has_sinks' value " << arg_has_sinks - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) - << "scalar 'total_work_items' value " << arg_total_work_items - << " is outside i32 range [-2147483648, 2147483647]"; - CheckSameCudaDevice(arg_tmap_swa_k, arg_tmap_q, "tmap_swa_k", "tmap_q"); - CheckSameCudaDevice(arg_tmap_compressed_k, arg_tmap_q, "tmap_compressed_k", "tmap_q"); - CheckSameCudaDevice(arg_tmap_swa_v, arg_tmap_q, "tmap_swa_v", "tmap_q"); - CheckSameCudaDevice(arg_tmap_compressed_v, arg_tmap_q, "tmap_compressed_v", "tmap_q"); - CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); - CheckSameCudaDevice(arg_partial_lse, arg_tmap_q, "partial_lse", "tmap_q"); - CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); - CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); - CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); - CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); - CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); - CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); - TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) - << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y - << ", " << grid_z << ")"; - TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) - << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z - << ") must be divisible by cluster dims (2, 1, 1)"; - - - CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); - CUtensorMap p_tmap_swa_k = EncodeTma_tmap_swa_k(arg_tmap_swa_k); - CUtensorMap p_tmap_compressed_k = EncodeTma_tmap_compressed_k(arg_tmap_compressed_k); - CUtensorMap p_tmap_swa_v = EncodeTma_tmap_swa_v(arg_tmap_swa_v); - CUtensorMap p_tmap_compressed_v = EncodeTma_tmap_compressed_v(arg_tmap_compressed_v); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); - float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); - int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); - int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); - int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); - float* p_sinks = static_cast(arg_sinks.data_ptr()); - float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); - float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); - int32_t v_num_heads = (int32_t)arg_num_heads; - int32_t v_num_query_tokens = (int32_t)arg_num_query_tokens; - int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; - int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; - int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; - int32_t v_sparse_topk = (int32_t)arg_sparse_topk; - int32_t v_has_sinks = (int32_t)arg_has_sinks; - int32_t v_total_work_items = (int32_t)arg_total_work_items; - void* kargs[] = {&p_tmap_q, &p_tmap_swa_k, &p_tmap_compressed_k, &p_tmap_swa_v, &p_tmap_compressed_v, &p_O, &p_partial_lse, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_num_query_tokens, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_has_sinks, &v_total_work_items}; - - static std::mutex smem_mu; - static auto* smem_status_by_device = - new std::unordered_map(); - cudaError_t smem_status = cudaSuccess; - { - std::lock_guard lock(smem_mu); - auto it = smem_status_by_device->find(dev.device_id); - if (it == smem_status_by_device->end()) { - int current_device = -1; - smem_status = cudaGetDevice(¤t_device); - if (smem_status == cudaSuccess && current_device != dev.device_id) { - smem_status = cudaErrorInvalidDevice; - } - if (smem_status == cudaSuccess) { - smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_ba9ed6a6233761682cc3), - cudaFuncAttributeMaxDynamicSharedMemorySize, - 231808); - } - smem_status_by_device->emplace(dev.device_id, smem_status); - } else { - smem_status = it->second; - } - } - TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_ba9ed6a6233761682cc3 failed: " - << cudaGetErrorString(smem_status); - dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); - dim3 block(512u, 1u, 1u); - - // Extended direct-source launch for a linked CUDA kernel symbol. - cudaLaunchAttribute attrs[1]{}; - int n = 0; - attrs[n].id = cudaLaunchAttributeClusterDimension; - attrs[n].val.clusterDim.x = 2u; - attrs[n].val.clusterDim.y = 1u; - attrs[n].val.clusterDim.z = 1u; - ++n; - cudaLaunchConfig_t config{}; - config.gridDim = grid; - config.blockDim = block; - config.dynamicSmemBytes = 231808u; - config.stream = stream; - config.attrs = attrs; - config.numAttrs = n; - cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_ba9ed6a6233761682cc3), kargs); - TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_ba9ed6a6233761682cc3 failed: " - << cudaGetErrorString(launch_status); -} -} // namespace variant_cake_dsv4_ba9ed6a6233761682cc3_sm_100a_f23ac2823b80 - -namespace variant_cake_dsv4_d906a03bd66a7488fc29_sm_100a_cd6fb31da15f { - -void Run(TensorView arg_partial_O, TensorView arg_partial_lse, TensorView arg_O, int64_t arg_num_q_heads, int64_t arg_num_split, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { - DLDevice dev = arg_partial_O.device(); - ScopedCudaDevice device_guard(dev.device_id); - CheckCudaTensor(arg_partial_O, "partial_O"); - CheckDtype(arg_partial_O, "partial_O", 4, 16, 1); - CheckContiguous(arg_partial_O, "partial_O"); - CheckCudaTensor(arg_partial_lse, "partial_lse"); - CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); - CheckContiguous(arg_partial_lse, "partial_lse"); - CheckCudaTensor(arg_O, "O"); - CheckDtype(arg_O, "O", 4, 16, 1); - CheckContiguous(arg_O, "O"); - TVM_FFI_CHECK(arg_num_q_heads >= -2147483648LL && arg_num_q_heads <= 2147483647LL, ValueError) - << "scalar 'num_q_heads' value " << arg_num_q_heads - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_num_split >= -2147483648LL && arg_num_split <= 2147483647LL, ValueError) - << "scalar 'num_split' value " << arg_num_split - << " is outside i32 range [-2147483648, 2147483647]"; - CheckSameCudaDevice(arg_partial_lse, arg_partial_O, "partial_lse", "partial_O"); - CheckSameCudaDevice(arg_O, arg_partial_O, "O", "partial_O"); - TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) - << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y - << ", " << grid_z << ")"; - - - __nv_bfloat16* p_partial_O = static_cast<__nv_bfloat16*>(arg_partial_O.data_ptr()); - float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); - int32_t v_num_q_heads = (int32_t)arg_num_q_heads; - int32_t v_num_split = (int32_t)arg_num_split; - void* kargs[] = {&p_partial_O, &p_partial_lse, &p_O, &v_num_q_heads, &v_num_split}; - - - dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); - dim3 block(128u, 1u, 1u); - - cudaLaunchConfig_t config{}; - config.gridDim = grid; - config.blockDim = block; - config.dynamicSmemBytes = 0u; - config.stream = stream; - cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_d906a03bd66a7488fc29), kargs); - TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_d906a03bd66a7488fc29 failed: " - << cudaGetErrorString(launch_status); -} -} // namespace variant_cake_dsv4_d906a03bd66a7488fc29_sm_100a_cd6fb31da15f - -int SelectVariant(int32_t num_heads, int32_t num_query_tokens, int32_t swa_index_stride, int32_t compressed_index_stride, int32_t sparse_topk_lens_offset, int32_t sparse_topk, int32_t has_sinks, int32_t total_work_items, int32_t num_split) { - if ((num_split == 5)) return 0; - return -1; -} - -void Dispatch(tvm::ffi::TensorView arg_Q, tvm::ffi::TensorView arg_SWA_cache, tvm::ffi::TensorView arg_compressed_KV_cache, tvm::ffi::TensorView arg_partial_O, tvm::ffi::TensorView arg_partial_lse, tvm::ffi::TensorView arg_swa_indices, tvm::ffi::TensorView arg_compressed_indices, tvm::ffi::TensorView arg_sparse_topk_lens, tvm::ffi::TensorView arg_sinks, tvm::ffi::TensorView arg_bmm1_scale, tvm::ffi::TensorView arg_bmm2_scale, tvm::ffi::TensorView arg_O, int32_t num_heads, int32_t num_query_tokens, int32_t swa_index_stride, int32_t compressed_index_stride, int32_t sparse_topk_lens_offset, int32_t sparse_topk, int32_t has_sinks, int32_t total_work_items, int32_t num_split, int64_t cuda_stream_ptr) { - CheckCudaTensor(arg_Q, "Q"); - CheckCudaTensor(arg_SWA_cache, "SWA_cache"); - CheckCudaTensor(arg_compressed_KV_cache, "compressed_KV_cache"); - CheckCudaTensor(arg_partial_O, "partial_O"); - CheckCudaTensor(arg_partial_lse, "partial_lse"); - CheckCudaTensor(arg_swa_indices, "swa_indices"); - CheckCudaTensor(arg_compressed_indices, "compressed_indices"); - CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckCudaTensor(arg_sinks, "sinks"); - CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); - CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); - CheckCudaTensor(arg_O, "O"); - cudaStream_t stream = reinterpret_cast(static_cast(cuda_stream_ptr)); - int v = SelectVariant(num_heads, num_query_tokens, swa_index_stride, compressed_index_stride, sparse_topk_lens_offset, sparse_topk, has_sinks, total_work_items, num_split); - switch (v) { - case 0: { - { - int64_t gx = (total_work_items * 2); - int64_t gy = 1; - int64_t gz = 1; - cake_host_shim_3cfd00f6b2ce7022::variant_cake_dsv4_ba9ed6a6233761682cc3_sm_100a_f23ac2823b80::Run(arg_Q, arg_SWA_cache, arg_compressed_KV_cache, arg_SWA_cache, arg_compressed_KV_cache, arg_partial_O, arg_partial_lse, arg_swa_indices, arg_compressed_indices, arg_sparse_topk_lens, arg_sinks, arg_bmm1_scale, arg_bmm2_scale, num_heads, num_query_tokens, swa_index_stride, compressed_index_stride, sparse_topk_lens_offset, sparse_topk, has_sinks, total_work_items, gx, gy, gz, stream); - } - { - int64_t gx = num_query_tokens; - int64_t gy = num_heads; - int64_t gz = 1; - cake_host_shim_3cfd00f6b2ce7022::variant_cake_dsv4_d906a03bd66a7488fc29_sm_100a_cd6fb31da15f::Run(arg_partial_O, arg_partial_lse, arg_O, num_heads, num_split, gx, gy, gz, stream); - } - return; - } - default: - TVM_FFI_CHECK(false, ValueError) << "no dispatch route for cake_dsv4_family_068c73b21d7cbcd958b8 shape (num_heads=" << num_heads << ", num_query_tokens=" << num_query_tokens << ", swa_index_stride=" << swa_index_stride << ", compressed_index_stride=" << compressed_index_stride << ", sparse_topk_lens_offset=" << sparse_topk_lens_offset << ", sparse_topk=" << sparse_topk << ", has_sinks=" << has_sinks << ", total_work_items=" << total_work_items << ", num_split=" << num_split << ")"; - } -} - -int64_t RouteOf(int32_t num_heads, int32_t num_query_tokens, int32_t swa_index_stride, int32_t compressed_index_stride, int32_t sparse_topk_lens_offset, int32_t sparse_topk, int32_t has_sinks, int32_t total_work_items, int32_t num_split) { return SelectVariant(num_heads, num_query_tokens, swa_index_stride, compressed_index_stride, sparse_topk_lens_offset, sparse_topk, has_sinks, total_work_items, num_split); } - -} // namespace cake_host_shim_3cfd00f6b2ce7022 - -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_3cfd00f6b2ce7022::Dispatch); -TVM_FFI_DLL_EXPORT_TYPED_FUNC(route_of, cake_host_shim_3cfd00f6b2ce7022::RouteOf); diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_family_a84dbf9944ef3cfb088b_binding.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_family_a84dbf9944ef3cfb088b_binding.cu deleted file mode 100644 index c24f61f5b90..00000000000 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_family_a84dbf9944ef3cfb088b_binding.cu +++ /dev/null @@ -1,652 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -// Generated by cake.runtime.ffi_dispatch — do not edit. -// tvm-ffi dispatch module for kernel family 'cake_dsv4_family_a84dbf9944ef3cfb088b'. -// dispatch-policy-revision: unpinned -#include -#include -#include -#include - -#include "tvm_ffi_utils.h" - -#include -#include -#include -#include -#include -#include -#include -#include - -extern "C" __global__ void kernel_cake_dsv4_d50dd820272ad24f7027(__nv_bfloat16* __restrict__ partial_O, float* __restrict__ partial_lse, __nv_bfloat16* __restrict__ O, int num_q_heads, int num_split); -extern "C" __global__ void kernel_cake_dsv4_e1660a7f44d156286fea(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); - -namespace cake_host_shim_c06b4c8dc8a15cae { - -using tvm::ffi::Optional; -using tvm::ffi::TensorView; - -class ScopedCudaDevice { - public: - explicit ScopedCudaDevice(int device_id) { - cudaError_t error = cudaGetDevice(&previous_device_); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed before host-shim launch: cudaError=" - << static_cast(error); - if (previous_device_ != device_id) { - error = cudaSetDevice(device_id); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaSetDevice failed before host-shim launch for cuda:" - << device_id << ": cudaError=" << static_cast(error); - restore_ = true; - } - } - - ScopedCudaDevice(const ScopedCudaDevice&) = delete; - ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; - - ~ScopedCudaDevice() noexcept { - if (restore_) { - (void)cudaSetDevice(previous_device_); - } - } - - private: - int previous_device_ = -1; - bool restore_ = false; -}; - -inline int64_t CakeDeviceMultiprocessorCount(int device_id) { - constexpr int kMaxCachedCudaDevices = 64; - TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) - << "physical-SM-count cache does not cover cuda:" << device_id; - static std::atomic count_by_device[kMaxCachedCudaDevices]{}; - int cached = count_by_device[device_id].load(std::memory_order_acquire); - if (cached > 0) return cached; - - int count = 0; - cudaError_t error = cudaDeviceGetAttribute( - &count, cudaDevAttrMultiProcessorCount, device_id); - TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) - << "querying multiProcessorCount failed for cuda:" << device_id - << ": cudaError=" << static_cast(error) << ", count=" << count; - // Concurrent first launches may repeat the immutable device query, but all - // publication is atomic and every later hot-path lookup is one acquire load. - count_by_device[device_id].store(count, std::memory_order_release); - return count; -} - -inline void CheckCudaTensor(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) - << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; -} - -inline void CheckSameCudaDevice( - const TensorView& t, - const TensorView& reference, - const char* name, - const char* reference_name) { - TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) - << name << " must be on the same CUDA device as " << reference_name - << ": got cuda:" << t.device().device_id - << " versus cuda:" << reference.device().device_id; -} - -inline void CheckCurrentCudaDevice( - const TensorView& reference, - const char* reference_name) { - int current_device = -1; - cudaError_t error = cudaGetDevice(¤t_device); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed while validating " << reference_name - << ": cudaError=" << static_cast(error); - TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) - << "current CUDA device must match " << reference_name - << ": current=cuda:" << current_device - << ", tensor=cuda:" << reference.device().device_id; -} - -inline void CheckContiguous(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; -} - -inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { - DLDataType d = t.dtype(); - TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) - << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits - << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits - << ", lanes=" << (int)d.lanes << ")"; -} - -// A logical axis.outer(trailing) folds every source dim above the trailing -// dimensions. Shape products are independent of physical strides, so verify -// the leading dimensions form one dense row-major chain instead of inventing -// a "folded stride". The descriptor reads its exact adjacent physical step -// separately through stride[-(trailing + 1)]. -inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { - TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) - << name << " cannot fold leading dimensions above " << trailing - << " trailing dims from ndim=" << t.ndim(); - int outer_last = t.ndim() - trailing - 1; - if (outer_last <= 0) { - return; - } - int64_t step = t.stride(outer_last); - TVM_FFI_CHECK(step > 0, ValueError) - << name << " physical strides must be positive"; - int64_t expected = step; - for (int axis = outer_last - 1; axis >= 0; --axis) { - expected *= t.size(axis + 1); - if (t.size(axis) > 1) { - TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) - << name << " leading dims are not physically foldable above " << trailing - << " trailing dims: stride(" << axis << ")=" << t.stride(axis) - << ", expected " << expected; - } - } -} - -namespace variant_cake_dsv4_e1660a7f44d156286fea_sm_100a_994fe0d2939f { - -// 4D TMA descriptor for buffer 'tmap_q' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 3, ValueError) - << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - int64_t d2 = t.size(t.ndim() - 2); - int64_t d3 = t.size(t.ndim() - 3); - TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) - << "TMA source 'tmap_q' trailing dims must be positive"; - TVM_FFI_CHECK(d1 % 64 == 0, ValueError) - << "TMA source 'tmap_q' extent " << d1 - << " must divide exactly by " << 64; - uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(d3)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) - << "TMA box (64, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; - int64_t carrier_stride_0 = d1; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; - int64_t carrier_stride_1 = 64; - TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; - TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; - TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; - int64_t carrier_stride_2 = (d2 * d1); - TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; - TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; - TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; - uint64_t global_strides[3] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - (uint64_t)((carrier_stride_1 * 16) / 8), - (uint64_t)((carrier_stride_2 * 16) / 8), - }; - uint32_t box_dim[4] = {64u, 64u, 1u, 1u}; - uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_swa_k' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_swa_k(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_swa_k' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_swa_k' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_swa_k' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_swa_k"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_swa_k' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_swa_k' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_k'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_swa_k') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_compressed_k' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_compressed_k(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_compressed_k' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_compressed_k' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_compressed_k' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_compressed_k"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_compressed_k' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_compressed_k' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_k'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_k') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_swa_v' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_swa_v(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_swa_v' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_swa_v' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_swa_v' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_swa_v"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_swa_v' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_swa_v' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_v'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_swa_v') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_compressed_v' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_compressed_v(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_compressed_v' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_compressed_v' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_compressed_v' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_compressed_v"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_compressed_v' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_compressed_v' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_v'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_v') failed: CUresult=" << (int)r; - return tm; -} - -void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_compressed_k, TensorView arg_tmap_swa_v, TensorView arg_tmap_compressed_v, TensorView arg_O, TensorView arg_partial_lse, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_num_query_tokens, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { - DLDevice dev = arg_tmap_q.device(); - ScopedCudaDevice device_guard(dev.device_id); - CheckCudaTensor(arg_tmap_q, "tmap_q"); - CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); - CheckContiguous(arg_tmap_q, "tmap_q"); - CheckCudaTensor(arg_tmap_swa_k, "tmap_swa_k"); - CheckDtype(arg_tmap_swa_k, "tmap_swa_k", 4, 16, 1); - CheckCudaTensor(arg_tmap_compressed_k, "tmap_compressed_k"); - CheckDtype(arg_tmap_compressed_k, "tmap_compressed_k", 4, 16, 1); - CheckCudaTensor(arg_tmap_swa_v, "tmap_swa_v"); - CheckDtype(arg_tmap_swa_v, "tmap_swa_v", 4, 16, 1); - CheckCudaTensor(arg_tmap_compressed_v, "tmap_compressed_v"); - CheckDtype(arg_tmap_compressed_v, "tmap_compressed_v", 4, 16, 1); - CheckCudaTensor(arg_O, "O"); - CheckDtype(arg_O, "O", 4, 16, 1); - CheckContiguous(arg_O, "O"); - CheckCudaTensor(arg_partial_lse, "partial_lse"); - CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); - CheckContiguous(arg_partial_lse, "partial_lse"); - CheckCudaTensor(arg_swa_indices, "swa_indices"); - CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); - CheckContiguous(arg_swa_indices, "swa_indices"); - CheckCudaTensor(arg_compressed_indices, "compressed_indices"); - CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); - CheckContiguous(arg_compressed_indices, "compressed_indices"); - CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); - CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckCudaTensor(arg_sinks, "sinks"); - CheckDtype(arg_sinks, "sinks", 2, 32, 1); - CheckContiguous(arg_sinks, "sinks"); - CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); - CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); - CheckContiguous(arg_bmm1_scale, "bmm1_scale"); - CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); - CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); - CheckContiguous(arg_bmm2_scale, "bmm2_scale"); - TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) - << "scalar 'num_heads' value " << arg_num_heads - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_num_query_tokens >= -2147483648LL && arg_num_query_tokens <= 2147483647LL, ValueError) - << "scalar 'num_query_tokens' value " << arg_num_query_tokens - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) - << "scalar 'swa_index_stride' value " << arg_swa_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) - << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) - << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) - << "scalar 'sparse_topk' value " << arg_sparse_topk - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) - << "scalar 'has_sinks' value " << arg_has_sinks - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) - << "scalar 'total_work_items' value " << arg_total_work_items - << " is outside i32 range [-2147483648, 2147483647]"; - CheckSameCudaDevice(arg_tmap_swa_k, arg_tmap_q, "tmap_swa_k", "tmap_q"); - CheckSameCudaDevice(arg_tmap_compressed_k, arg_tmap_q, "tmap_compressed_k", "tmap_q"); - CheckSameCudaDevice(arg_tmap_swa_v, arg_tmap_q, "tmap_swa_v", "tmap_q"); - CheckSameCudaDevice(arg_tmap_compressed_v, arg_tmap_q, "tmap_compressed_v", "tmap_q"); - CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); - CheckSameCudaDevice(arg_partial_lse, arg_tmap_q, "partial_lse", "tmap_q"); - CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); - CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); - CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); - CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); - CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); - CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); - TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) - << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y - << ", " << grid_z << ")"; - TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) - << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z - << ") must be divisible by cluster dims (2, 1, 1)"; - - - CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); - CUtensorMap p_tmap_swa_k = EncodeTma_tmap_swa_k(arg_tmap_swa_k); - CUtensorMap p_tmap_compressed_k = EncodeTma_tmap_compressed_k(arg_tmap_compressed_k); - CUtensorMap p_tmap_swa_v = EncodeTma_tmap_swa_v(arg_tmap_swa_v); - CUtensorMap p_tmap_compressed_v = EncodeTma_tmap_compressed_v(arg_tmap_compressed_v); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); - float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); - int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); - int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); - int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); - float* p_sinks = static_cast(arg_sinks.data_ptr()); - float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); - float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); - int32_t v_num_heads = (int32_t)arg_num_heads; - int32_t v_num_query_tokens = (int32_t)arg_num_query_tokens; - int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; - int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; - int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; - int32_t v_sparse_topk = (int32_t)arg_sparse_topk; - int32_t v_has_sinks = (int32_t)arg_has_sinks; - int32_t v_total_work_items = (int32_t)arg_total_work_items; - void* kargs[] = {&p_tmap_q, &p_tmap_swa_k, &p_tmap_compressed_k, &p_tmap_swa_v, &p_tmap_compressed_v, &p_O, &p_partial_lse, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_num_query_tokens, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_has_sinks, &v_total_work_items}; - - static std::mutex smem_mu; - static auto* smem_status_by_device = - new std::unordered_map(); - cudaError_t smem_status = cudaSuccess; - { - std::lock_guard lock(smem_mu); - auto it = smem_status_by_device->find(dev.device_id); - if (it == smem_status_by_device->end()) { - int current_device = -1; - smem_status = cudaGetDevice(¤t_device); - if (smem_status == cudaSuccess && current_device != dev.device_id) { - smem_status = cudaErrorInvalidDevice; - } - if (smem_status == cudaSuccess) { - smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_e1660a7f44d156286fea), - cudaFuncAttributeMaxDynamicSharedMemorySize, - 220032); - } - smem_status_by_device->emplace(dev.device_id, smem_status); - } else { - smem_status = it->second; - } - } - TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_e1660a7f44d156286fea failed: " - << cudaGetErrorString(smem_status); - dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); - dim3 block(512u, 1u, 1u); - - // Extended direct-source launch for a linked CUDA kernel symbol. - cudaLaunchAttribute attrs[1]{}; - int n = 0; - attrs[n].id = cudaLaunchAttributeClusterDimension; - attrs[n].val.clusterDim.x = 2u; - attrs[n].val.clusterDim.y = 1u; - attrs[n].val.clusterDim.z = 1u; - ++n; - cudaLaunchConfig_t config{}; - config.gridDim = grid; - config.blockDim = block; - config.dynamicSmemBytes = 220032u; - config.stream = stream; - config.attrs = attrs; - config.numAttrs = n; - cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_e1660a7f44d156286fea), kargs); - TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_e1660a7f44d156286fea failed: " - << cudaGetErrorString(launch_status); -} -} // namespace variant_cake_dsv4_e1660a7f44d156286fea_sm_100a_994fe0d2939f - -namespace variant_cake_dsv4_d50dd820272ad24f7027_sm_100a_aa3105f3b9de { - -void Run(TensorView arg_partial_O, TensorView arg_partial_lse, TensorView arg_O, int64_t arg_num_q_heads, int64_t arg_num_split, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { - DLDevice dev = arg_partial_O.device(); - ScopedCudaDevice device_guard(dev.device_id); - CheckCudaTensor(arg_partial_O, "partial_O"); - CheckDtype(arg_partial_O, "partial_O", 4, 16, 1); - CheckContiguous(arg_partial_O, "partial_O"); - CheckCudaTensor(arg_partial_lse, "partial_lse"); - CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); - CheckContiguous(arg_partial_lse, "partial_lse"); - CheckCudaTensor(arg_O, "O"); - CheckDtype(arg_O, "O", 4, 16, 1); - CheckContiguous(arg_O, "O"); - TVM_FFI_CHECK(arg_num_q_heads >= -2147483648LL && arg_num_q_heads <= 2147483647LL, ValueError) - << "scalar 'num_q_heads' value " << arg_num_q_heads - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_num_split >= -2147483648LL && arg_num_split <= 2147483647LL, ValueError) - << "scalar 'num_split' value " << arg_num_split - << " is outside i32 range [-2147483648, 2147483647]"; - CheckSameCudaDevice(arg_partial_lse, arg_partial_O, "partial_lse", "partial_O"); - CheckSameCudaDevice(arg_O, arg_partial_O, "O", "partial_O"); - TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) - << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y - << ", " << grid_z << ")"; - - - __nv_bfloat16* p_partial_O = static_cast<__nv_bfloat16*>(arg_partial_O.data_ptr()); - float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); - int32_t v_num_q_heads = (int32_t)arg_num_q_heads; - int32_t v_num_split = (int32_t)arg_num_split; - void* kargs[] = {&p_partial_O, &p_partial_lse, &p_O, &v_num_q_heads, &v_num_split}; - - - dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); - dim3 block(128u, 1u, 1u); - - cudaLaunchConfig_t config{}; - config.gridDim = grid; - config.blockDim = block; - config.dynamicSmemBytes = 0u; - config.stream = stream; - cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_d50dd820272ad24f7027), kargs); - TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_d50dd820272ad24f7027 failed: " - << cudaGetErrorString(launch_status); -} -} // namespace variant_cake_dsv4_d50dd820272ad24f7027_sm_100a_aa3105f3b9de - -int SelectVariant(int32_t num_heads, int32_t num_query_tokens, int32_t swa_index_stride, int32_t compressed_index_stride, int32_t sparse_topk_lens_offset, int32_t sparse_topk, int32_t has_sinks, int32_t total_work_items, int32_t num_split) { - if ((num_split == 4)) return 0; - return -1; -} - -void Dispatch(tvm::ffi::TensorView arg_Q, tvm::ffi::TensorView arg_SWA_cache, tvm::ffi::TensorView arg_compressed_KV_cache, tvm::ffi::TensorView arg_partial_O, tvm::ffi::TensorView arg_partial_lse, tvm::ffi::TensorView arg_swa_indices, tvm::ffi::TensorView arg_compressed_indices, tvm::ffi::TensorView arg_sparse_topk_lens, tvm::ffi::TensorView arg_sinks, tvm::ffi::TensorView arg_bmm1_scale, tvm::ffi::TensorView arg_bmm2_scale, tvm::ffi::TensorView arg_O, int32_t num_heads, int32_t num_query_tokens, int32_t swa_index_stride, int32_t compressed_index_stride, int32_t sparse_topk_lens_offset, int32_t sparse_topk, int32_t has_sinks, int32_t total_work_items, int32_t num_split, int64_t cuda_stream_ptr) { - CheckCudaTensor(arg_Q, "Q"); - CheckCudaTensor(arg_SWA_cache, "SWA_cache"); - CheckCudaTensor(arg_compressed_KV_cache, "compressed_KV_cache"); - CheckCudaTensor(arg_partial_O, "partial_O"); - CheckCudaTensor(arg_partial_lse, "partial_lse"); - CheckCudaTensor(arg_swa_indices, "swa_indices"); - CheckCudaTensor(arg_compressed_indices, "compressed_indices"); - CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckCudaTensor(arg_sinks, "sinks"); - CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); - CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); - CheckCudaTensor(arg_O, "O"); - cudaStream_t stream = reinterpret_cast(static_cast(cuda_stream_ptr)); - int v = SelectVariant(num_heads, num_query_tokens, swa_index_stride, compressed_index_stride, sparse_topk_lens_offset, sparse_topk, has_sinks, total_work_items, num_split); - switch (v) { - case 0: { - { - int64_t gx = (total_work_items * 2); - int64_t gy = 1; - int64_t gz = 1; - cake_host_shim_c06b4c8dc8a15cae::variant_cake_dsv4_e1660a7f44d156286fea_sm_100a_994fe0d2939f::Run(arg_Q, arg_SWA_cache, arg_compressed_KV_cache, arg_SWA_cache, arg_compressed_KV_cache, arg_partial_O, arg_partial_lse, arg_swa_indices, arg_compressed_indices, arg_sparse_topk_lens, arg_sinks, arg_bmm1_scale, arg_bmm2_scale, num_heads, num_query_tokens, swa_index_stride, compressed_index_stride, sparse_topk_lens_offset, sparse_topk, has_sinks, total_work_items, gx, gy, gz, stream); - } - { - int64_t gx = num_query_tokens; - int64_t gy = num_heads; - int64_t gz = 1; - cake_host_shim_c06b4c8dc8a15cae::variant_cake_dsv4_d50dd820272ad24f7027_sm_100a_aa3105f3b9de::Run(arg_partial_O, arg_partial_lse, arg_O, num_heads, num_split, gx, gy, gz, stream); - } - return; - } - default: - TVM_FFI_CHECK(false, ValueError) << "no dispatch route for cake_dsv4_family_a84dbf9944ef3cfb088b shape (num_heads=" << num_heads << ", num_query_tokens=" << num_query_tokens << ", swa_index_stride=" << swa_index_stride << ", compressed_index_stride=" << compressed_index_stride << ", sparse_topk_lens_offset=" << sparse_topk_lens_offset << ", sparse_topk=" << sparse_topk << ", has_sinks=" << has_sinks << ", total_work_items=" << total_work_items << ", num_split=" << num_split << ")"; - } -} - -int64_t RouteOf(int32_t num_heads, int32_t num_query_tokens, int32_t swa_index_stride, int32_t compressed_index_stride, int32_t sparse_topk_lens_offset, int32_t sparse_topk, int32_t has_sinks, int32_t total_work_items, int32_t num_split) { return SelectVariant(num_heads, num_query_tokens, swa_index_stride, compressed_index_stride, sparse_topk_lens_offset, sparse_topk, has_sinks, total_work_items, num_split); } - -} // namespace cake_host_shim_c06b4c8dc8a15cae - -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_c06b4c8dc8a15cae::Dispatch); -TVM_FFI_DLL_EXPORT_TYPED_FUNC(route_of, cake_host_shim_c06b4c8dc8a15cae::RouteOf); diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_fb6409904b35ed0f7a1a_binding.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_fb6409904b35ed0f7a1a_binding.cu deleted file mode 100644 index 45f5eb50419..00000000000 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_fb6409904b35ed0f7a1a_binding.cu +++ /dev/null @@ -1,461 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -// Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_fb6409904b35ed0f7a1a'. -#include -#include -#include -#include - -#include "tvm_ffi_utils.h" - -#include -#include -#include -#include -#include -#include -#include - -extern "C" __global__ void kernel_cake_dsv4_fb6409904b35ed0f7a1a(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items); - - -namespace cake_host_shim_c54faed3b4eb93e3 { - -using tvm::ffi::Optional; -using tvm::ffi::TensorView; - -class ScopedCudaDevice { - public: - explicit ScopedCudaDevice(int device_id) { - cudaError_t error = cudaGetDevice(&previous_device_); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed before host-shim launch: cudaError=" - << static_cast(error); - if (previous_device_ != device_id) { - error = cudaSetDevice(device_id); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaSetDevice failed before host-shim launch for cuda:" - << device_id << ": cudaError=" << static_cast(error); - restore_ = true; - } - } - - ScopedCudaDevice(const ScopedCudaDevice&) = delete; - ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; - - ~ScopedCudaDevice() noexcept { - if (restore_) { - (void)cudaSetDevice(previous_device_); - } - } - - private: - int previous_device_ = -1; - bool restore_ = false; -}; - -inline int64_t CakeDeviceMultiprocessorCount(int device_id) { - constexpr int kMaxCachedCudaDevices = 64; - TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) - << "physical-SM-count cache does not cover cuda:" << device_id; - static std::atomic count_by_device[kMaxCachedCudaDevices]{}; - int cached = count_by_device[device_id].load(std::memory_order_acquire); - if (cached > 0) return cached; - - int count = 0; - cudaError_t error = cudaDeviceGetAttribute( - &count, cudaDevAttrMultiProcessorCount, device_id); - TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) - << "querying multiProcessorCount failed for cuda:" << device_id - << ": cudaError=" << static_cast(error) << ", count=" << count; - // Concurrent first launches may repeat the immutable device query, but all - // publication is atomic and every later hot-path lookup is one acquire load. - count_by_device[device_id].store(count, std::memory_order_release); - return count; -} - -inline void CheckCudaTensor(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) - << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; -} - -inline void CheckSameCudaDevice( - const TensorView& t, - const TensorView& reference, - const char* name, - const char* reference_name) { - TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) - << name << " must be on the same CUDA device as " << reference_name - << ": got cuda:" << t.device().device_id - << " versus cuda:" << reference.device().device_id; -} - -inline void CheckCurrentCudaDevice( - const TensorView& reference, - const char* reference_name) { - int current_device = -1; - cudaError_t error = cudaGetDevice(¤t_device); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed while validating " << reference_name - << ": cudaError=" << static_cast(error); - TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) - << "current CUDA device must match " << reference_name - << ": current=cuda:" << current_device - << ", tensor=cuda:" << reference.device().device_id; -} - -inline void CheckContiguous(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; -} - -inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { - DLDataType d = t.dtype(); - TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) - << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits - << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits - << ", lanes=" << (int)d.lanes << ")"; -} - -// A logical axis.outer(trailing) folds every source dim above the trailing -// dimensions. Shape products are independent of physical strides, so verify -// the leading dimensions form one dense row-major chain instead of inventing -// a "folded stride". The descriptor reads its exact adjacent physical step -// separately through stride[-(trailing + 1)]. -inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { - TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) - << name << " cannot fold leading dimensions above " << trailing - << " trailing dims from ndim=" << t.ndim(); - int outer_last = t.ndim() - trailing - 1; - if (outer_last <= 0) { - return; - } - int64_t step = t.stride(outer_last); - TVM_FFI_CHECK(step > 0, ValueError) - << name << " physical strides must be positive"; - int64_t expected = step; - for (int axis = outer_last - 1; axis >= 0; --axis) { - expected *= t.size(axis + 1); - if (t.size(axis) > 1) { - TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) - << name << " leading dims are not physically foldable above " << trailing - << " trailing dims: stride(" << axis << ")=" << t.stride(axis) - << ", expected " << expected; - } - } -} - -// 4D TMA descriptor for buffer 'tmap_q' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 3, ValueError) - << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - int64_t d2 = t.size(t.ndim() - 2); - int64_t d3 = t.size(t.ndim() - 3); - TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) - << "TMA source 'tmap_q' trailing dims must be positive"; - TVM_FFI_CHECK(d1 % 128 == 0, ValueError) - << "TMA source 'tmap_q' extent " << d1 - << " must divide exactly by " << 128; - uint64_t global_dim[4] = {(uint64_t)(128), (uint64_t)(d2), (uint64_t)((d1 / 128)), (uint64_t)(d3)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; - TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) - << "TMA box (128, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; - int64_t carrier_stride_0 = d1; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; - int64_t carrier_stride_1 = 128; - TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; - TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; - TVM_FFI_CHECK((carrier_stride_1 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; - int64_t carrier_stride_2 = (d2 * d1); - TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; - TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; - TVM_FFI_CHECK((carrier_stride_2 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; - uint64_t global_strides[3] = { - (uint64_t)((carrier_stride_0 * 8) / 8), - (uint64_t)((carrier_stride_1 * 8) / 8), - (uint64_t)((carrier_stride_2 * 8) / 8), - }; - uint32_t box_dim[4] = {128u, 64u, 1u, 1u}; - uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_swa_kv' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_swa_kv' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; - TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (128, 1) exceeds resolved global dims for 'tmap_swa_kv'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 8) / 8), - }; - uint32_t box_dim[2] = {128u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_compressed_kv' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_compressed_kv(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_compressed_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_compressed_kv' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_compressed_kv' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_compressed_kv"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_compressed_kv' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved a non-positive global dim"; - TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (128, 1) exceeds resolved global dims for 'tmap_compressed_kv'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 8) / 8), - }; - uint32_t box_dim[2] = {128u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_kv') failed: CUresult=" << (int)r; - return tm; -} - -void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_compressed_kv, TensorView arg_O, TensorView arg_partial_lse, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_num_query_tokens, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, int64_t grid_x, int64_t grid_y, int64_t grid_z) { - DLDevice dev = arg_tmap_q.device(); - ScopedCudaDevice device_guard(dev.device_id); - CheckCudaTensor(arg_tmap_q, "tmap_q"); - CheckDtype(arg_tmap_q, "tmap_q", 1, 8, 1); - CheckContiguous(arg_tmap_q, "tmap_q"); - CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); - CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 1, 8, 1); - CheckCudaTensor(arg_tmap_compressed_kv, "tmap_compressed_kv"); - CheckDtype(arg_tmap_compressed_kv, "tmap_compressed_kv", 1, 8, 1); - CheckCudaTensor(arg_O, "O"); - CheckDtype(arg_O, "O", 4, 16, 1); - CheckContiguous(arg_O, "O"); - CheckCudaTensor(arg_partial_lse, "partial_lse"); - CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); - CheckContiguous(arg_partial_lse, "partial_lse"); - CheckCudaTensor(arg_swa_indices, "swa_indices"); - CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); - CheckContiguous(arg_swa_indices, "swa_indices"); - CheckCudaTensor(arg_compressed_indices, "compressed_indices"); - CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); - CheckContiguous(arg_compressed_indices, "compressed_indices"); - CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); - CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckCudaTensor(arg_sinks, "sinks"); - CheckDtype(arg_sinks, "sinks", 2, 32, 1); - CheckContiguous(arg_sinks, "sinks"); - CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); - CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); - CheckContiguous(arg_bmm1_scale, "bmm1_scale"); - CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); - CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); - CheckContiguous(arg_bmm2_scale, "bmm2_scale"); - TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) - << "scalar 'num_heads' value " << arg_num_heads - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) - << "scalar 'swa_index_stride' value " << arg_swa_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) - << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) - << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_num_query_tokens >= -2147483648LL && arg_num_query_tokens <= 2147483647LL, ValueError) - << "scalar 'num_query_tokens' value " << arg_num_query_tokens - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) - << "scalar 'sparse_topk' value " << arg_sparse_topk - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) - << "scalar 'has_sinks' value " << arg_has_sinks - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) - << "scalar 'total_work_items' value " << arg_total_work_items - << " is outside i32 range [-2147483648, 2147483647]"; - CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); - CheckSameCudaDevice(arg_tmap_compressed_kv, arg_tmap_q, "tmap_compressed_kv", "tmap_q"); - CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); - CheckSameCudaDevice(arg_partial_lse, arg_tmap_q, "partial_lse", "tmap_q"); - CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); - CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); - CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); - CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); - CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); - CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); - TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) - << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y - << ", " << grid_z << ")"; - TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) - << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z - << ") must be divisible by cluster dims (2, 1, 1)"; - - cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); - CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); - CUtensorMap p_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); - CUtensorMap p_tmap_compressed_kv = EncodeTma_tmap_compressed_kv(arg_tmap_compressed_kv); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); - float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); - int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); - int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); - int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); - float* p_sinks = static_cast(arg_sinks.data_ptr()); - float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); - float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); - int32_t v_num_heads = (int32_t)arg_num_heads; - int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; - int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; - int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; - int32_t v_num_query_tokens = (int32_t)arg_num_query_tokens; - int32_t v_sparse_topk = (int32_t)arg_sparse_topk; - int32_t v_has_sinks = (int32_t)arg_has_sinks; - int32_t v_total_work_items = (int32_t)arg_total_work_items; - void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_tmap_compressed_kv, &p_O, &p_partial_lse, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_num_query_tokens, &v_sparse_topk, &v_has_sinks, &v_total_work_items}; - - static std::mutex smem_mu; - static auto* smem_status_by_device = - new std::unordered_map(); - cudaError_t smem_status = cudaSuccess; - { - std::lock_guard lock(smem_mu); - auto it = smem_status_by_device->find(dev.device_id); - if (it == smem_status_by_device->end()) { - int current_device = -1; - smem_status = cudaGetDevice(¤t_device); - if (smem_status == cudaSuccess && current_device != dev.device_id) { - smem_status = cudaErrorInvalidDevice; - } - if (smem_status == cudaSuccess) { - smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_fb6409904b35ed0f7a1a), - cudaFuncAttributeMaxDynamicSharedMemorySize, - 190976); - } - smem_status_by_device->emplace(dev.device_id, smem_status); - } else { - smem_status = it->second; - } - } - TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_fb6409904b35ed0f7a1a failed: " - << cudaGetErrorString(smem_status); - dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); - dim3 block(512u, 1u, 1u); - - // Extended direct-source launch for a linked CUDA kernel symbol. - cudaLaunchAttribute attrs[2]{}; - int n = 0; - attrs[n].id = cudaLaunchAttributeClusterDimension; - attrs[n].val.clusterDim.x = 2u; - attrs[n].val.clusterDim.y = 1u; - attrs[n].val.clusterDim.z = 1u; - ++n; - attrs[n].id = cudaLaunchAttributeClusterSchedulingPolicyPreference; - attrs[n].val.clusterSchedulingPolicyPreference = cudaClusterSchedulingPolicySpread; - ++n; - cudaLaunchConfig_t config{}; - config.gridDim = grid; - config.blockDim = block; - config.dynamicSmemBytes = 190976u; - config.stream = stream; - config.attrs = attrs; - config.numAttrs = n; - cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_fb6409904b35ed0f7a1a), kargs); - TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_fb6409904b35ed0f7a1a failed: " - << cudaGetErrorString(launch_status); -} - -} // namespace cake_host_shim_c54faed3b4eb93e3 - -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_c54faed3b4eb93e3::Run); diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_fb6409904b35ed0f7a1a_kernel.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_fb6409904b35ed0f7a1a_kernel.cu deleted file mode 100644 index bcc67b72a4c..00000000000 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_fb6409904b35ed0f7a1a_kernel.cu +++ /dev/null @@ -1,2812 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -typedef signed char int8_t; -typedef unsigned char uint8_t; -typedef unsigned short uint16_t; -typedef unsigned int uint32_t; -#if defined(__CUDACC_RTC__) -typedef unsigned long long uint64_t; -#else -typedef unsigned long uint64_t; -#endif -static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); -typedef signed int int32_t; -typedef short int int16_t; -struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; -struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; -static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); -static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); - -#if defined(__CUDACC_RTC__) -typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; -#else -#include -#endif - -static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); -static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); -#include -#include - -#define CAKE_INF CUDART_INF_F -#define TMEM_NCOLS 512 -#define TMEM_TMEM_SCRATCH_OFFSET 0 -#define NUM_K_PIPE_STAGES 2 -#define NUM_V_PIPE_STAGES 2 -#define NUM_KV_PIPE_STAGES 1 -#define NUM_INDEX_PIPE_STAGES 6 -#define SMEM_SMEM_Q_OFF 1024 -#define SMEM_SMEM_Q_STAGE_BYTES 8192 -#define SMEM_SMEM_Q_STRIDE 8192 -#define SMEM_SMEM_KV_OFF 33792 -#define SMEM_SMEM_KV_STAGE_BYTES 16384 -#define SMEM_SMEM_KV_STRIDE 16384 -#define SMEM_SMEM_V_OFF 33792 -#define SMEM_SMEM_V_STAGE_BYTES 16384 -#define SMEM_SMEM_V_STRIDE 16384 -#define SMEM_SMEM_Q_FULL_OFF 1024 -#define SMEM_SMEM_Q_FULL_STAGE_BYTES 32768 -#define SMEM_SMEM_Q_FULL_STRIDE 32768 -#define SMEM_SMEM_K_FULL_OFF 33792 -#define SMEM_SMEM_K_FULL_STAGE_BYTES 32768 -#define SMEM_SMEM_K_FULL_STRIDE 32768 -#define SMEM_SMEM_V_FULL_OFF 99328 -#define SMEM_SMEM_V_FULL_STAGE_BYTES 32768 -#define SMEM_SMEM_V_FULL_STRIDE 32768 -#define SMEM_SMEM_STATS_MAX_OFF 164864 -#define SMEM_SMEM_STATS_MAX_STAGE_BYTES 1024 -#define SMEM_SMEM_STATS_MAX_STRIDE 1024 -#define SMEM_SMEM_STATS_SUM_OFF 165888 -#define SMEM_SMEM_STATS_SUM_STAGE_BYTES 512 -#define SMEM_SMEM_STATS_SUM_STRIDE 512 -#define SMEM_SMEM_STATS_FINAL_MAX_OFF 166400 -#define SMEM_SMEM_STATS_FINAL_MAX_STAGE_BYTES 512 -#define SMEM_SMEM_STATS_FINAL_MAX_STRIDE 512 -#define SMEM_SMEM_SOFTMAX_WARP_PAIR_EXCHANGE_OFF 166912 -#define SMEM_SMEM_SOFTMAX_WARP_PAIR_EXCHANGE_STAGE_BYTES 1024 -#define SMEM_SMEM_SOFTMAX_WARP_PAIR_EXCHANGE_STRIDE 1024 -#define SMEM_SMEM_CORR_WARP_PAIR_EXCHANGE_OFF 167936 -#define SMEM_SMEM_CORR_WARP_PAIR_EXCHANGE_STAGE_BYTES 512 -#define SMEM_SMEM_CORR_WARP_PAIR_EXCHANGE_STRIDE 512 -#define SMEM_SMEM_SPARSE_INDICES_OFF 168448 -#define SMEM_SMEM_SPARSE_INDICES_STAGE_BYTES 1024 -#define SMEM_SMEM_SPARSE_INDICES_STRIDE 1024 -#define SMEM_SMEM_P_FP8_OFF 174592 -#define SMEM_SMEM_P_FP8_STAGE_BYTES 8192 -#define SMEM_SMEM_P_FP8_STRIDE 8192 -#define SMEM_TOTAL 190976 -#define THREADS 512 - -#include - -__device__ __forceinline__ uint32_t elect_sync() { - uint32_t pred = 0; - asm volatile( - "{\n\t" - ".reg .pred %%px;\n\t" - "elect.sync _|%%px, %1;\n\t" - "@%%px mov.s32 %0, 1;\n\t" - "}\n" - : "+r"(pred) - : "r"(0xFFFFFFFF)); - return pred; -} - - -__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { - asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" - :: "r"(mbar_addr), "r"(count) : "memory"); -} - -__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { - asm volatile("mbarrier.init.b64 [%0], %1;" - :: "l"(mbar_addr), "r"(count) : "memory"); -} - - -__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - - -// CTA-local pipelines have short, resident producer/consumer edges. Omitting -// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the -// explicit loop still makes this helper blocking until acquire succeeds. -__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE;\n\t" - "bra.uni LAB_WAIT;\n\t" - "DONE:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Source-faithful relaxed CTA wait used only by a typed protocol that does -// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. -__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, 10000000;\n\t" - "@P1 bra.uni DONE_RELAXED;\n\t" - "bra.uni LAB_WAIT_RELAXED;\n\t" - "DONE_RELAXED:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Exact source ports may request the PTX suspendTimeHint operand explicitly. -// The hint is expressed in nanoseconds and is kept separate from the canonical -// no-hint CTA helper so unrelated schedules retain their existing retry path. -__device__ __forceinline__ void mbarrier_wait_suspend( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_SUSPEND:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_SUSPEND;\n\t" - "bra.uni LAB_WAIT_SUSPEND;\n\t" - "DONE_SUSPEND:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE_CLUSTER;\n\t" - "bra.uni LAB_WAIT_CLUSTER;\n\t" - "DONE_CLUSTER:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - ".reg .u32 WAIT_ADDR;\n\t" - "mov.u32 WAIT_ADDR, %0;\n\t" - "LAB_WAIT_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [WAIT_ADDR], %1, %2;\n\t" - "@P1 bra.uni DONE_HINT;\n\t" - "bra.uni LAB_WAIT_HINT;\n\t" - "DONE_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -// Exact unqualified CTA wait used by source schedules whose PTX intentionally -// omits the acquire qualifier while retaining a typed suspendTimeHint operand. -__device__ __forceinline__ void mbarrier_wait_relaxed_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED_HINT:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra DONE_RELAXED_HINT;\n\t" - "bra LAB_WAIT_RELAXED_HINT;\n\t" - "DONE_RELAXED_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); -} - -__device__ __forceinline__ void mbarrier_wait_cluster_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" - "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" - "DONE_CLUSTER_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_suspend( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait_cluster(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); - } -} - - -__device__ __forceinline__ void tcgen05_mma_f8f6f4_cta2( - int taddr, uint64_t a_desc, uint64_t b_desc, - uint32_t i_desc, int enable_input_d) { - asm volatile( - "{\n\t" - ".reg .pred p;\n\t" - ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - "setp.ne.b32 p, %4, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" - "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "tcgen05.mma.cta_group::2.kind::f8f6f4 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" - "}\n" - :: "r"(taddr), "l"(a_desc), "l"(b_desc), - "r"(i_desc), "r"(enable_input_d) - : "memory"); -} - - -__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { - return (x & 0x3FFFFULL) >> 4ULL; -} - - -union MmaSmemDesc { - uint64_t u64; - uint32_t u32[2]; -}; - -__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { - MmaSmemDesc tmp; - tmp.u64 = smem_desc; - tmp.u32[0] += offset; - smem_desc = tmp.u64; -} - - -__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { - asm volatile( - "{\n\t" - ".reg .pred leader;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" - ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" - "}\n" - :: "r"(mbar_addr), "h"(cta_mask) : "memory"); -} - - -__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { - asm volatile( - "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" - :: "r"(mbar_addr), "r"(bytes) : "memory"); -} - - -__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7," - " %8, %9, %10, %11, %12, %13, %14, %15," - " %16, %17, %18, %19, %20, %21, %22, %23," - " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), - "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), - "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), - "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), - "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), - "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), - "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x16.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7," - " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), - "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), - "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ void tmem_ld_x4(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x4.b32" - " {%0, %1, %2, %3}, [%4];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { - asm volatile( - "tcgen05.st.sync.aligned.32x32b.x16.b32" - " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," - " %9, %10, %11, %12, %13, %14, %15, %16};" - :: "r"(tmem_addr), - "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), - "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), - "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), - "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); -} - - -__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { - asm volatile( - "tcgen05.st.sync.aligned.32x32b.x4.b32" - " [%0], {%1, %2, %3, %4};" - :: "r"(tmem_addr), - "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); -} - - -__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { - uint32_t addr; - asm("{\n\t" - ".reg .u64 u64addr;\n\t" - "cvta.to.shared.u64 u64addr, %1;\n\t" - "cvt.u32.u64 %0, u64addr;\n\t" - "}\n" : "=r"(addr) : "l"(ptr)); - return addr; -} - - -__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { - uint32_t remote; - asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" - : "=r"(remote) : "r"(local_addr), "r"(rank)); - return remote; -} - - -__device__ __forceinline__ float approx_exp2(float x) { - float y; - asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float approx_rcp(float x) { - float y; - asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float max_noftz(float a, float b) { - float c; - asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); - return c; -} - - -__device__ __forceinline__ float warp_reduce_max(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); - return val; -} - - -__device__ __forceinline__ float warp_reduce_sum(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val += __shfl_xor_sync(0xFFFFFFFF, val, offset); - return val; -} - - -__device__ __forceinline__ float row_max_reduce(float2 acc) { - return max_noftz(acc.x, acc.y); -} - - -__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { - #pragma unroll - for (int j = 0; j < 16; j++) { - if (j % 2 == 0) - acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); - else - acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); - } -} - - -__device__ __forceinline__ float2 ex2_emulation_f32x2_value(float2 value) { - const float c0 = 1.0f, c1 = 0.695146143436431884765625f; - const float c2 = 0.227564394474029541015625f, c3 = 0.077119089663028717041015625f; - const float magic = 12582912.0f; - float x0 = max_noftz(value.x, -127.0f), x1 = max_noftz(value.y, -127.0f); - float2 xc2 = make_float2(x0, x1), magic2 = make_float2(magic, magic); - float2 xr2; - asm("add.rm.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xr2) - : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&magic2)); - float2 c3_2 = make_float2(c3, c3), c2_2 = make_float2(c2, c2); - float2 c1_2 = make_float2(c1, c1), c0_2 = make_float2(c0, c0); - float2 xrb2, xfrac2; - asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xrb2) - : "l"(*(unsigned long long*)&xr2), "l"(*(unsigned long long*)&magic2)); - asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xfrac2) - : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&xrb2)); - float2 poly2; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&c3_2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c2_2)); - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c1_2)); - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c0_2)); - int x0r_i, x1r_i, p0_i, p1_i; - asm("mov.b64 {%0, %1}, %2;" : "=r"(x0r_i), "=r"(x1r_i) : "l"(*(unsigned long long*)&xr2)); - asm("mov.b64 {%0, %1}, %2;" : "=r"(p0_i), "=r"(p1_i) : "l"(*(unsigned long long*)&poly2)); - float r0, r1; - asm("mov.b32 %0, %1;" : "=f"(r0) : "r"((x0r_i << 23) + p0_i)); - asm("mov.b32 %0, %1;" : "=f"(r1) : "r"((x1r_i << 23) + p1_i)); - return make_float2(r0, r1); -} - -__device__ __forceinline__ void ex2_emulation_f32x2(float* x0_ptr, float* x1_ptr) { - float2 result = ex2_emulation_f32x2_value(make_float2(*x0_ptr, *x1_ptr)); - *x0_ptr = result.x; *x1_ptr = result.y; -} - -__device__ __forceinline__ void softmax_frag_exp2_cast( - float* sv, uint32_t* pv, int use_emu) -{ - #pragma unroll - for (int j = 0; j < 16; j++) { - if (use_emu && j >= 12) - ex2_emulation_f32x2(&sv[j*2], &sv[j*2+1]); - else { - sv[j*2] = approx_exp2(sv[j*2]); - sv[j*2+1] = approx_exp2(sv[j*2+1]); - } - } - #pragma unroll - for (int j = 0; j < 16; j++) { - __nv_bfloat162 bf = __float22bfloat162_rn({sv[j*2], sv[j*2+1]}); - pv[j] = reinterpret_cast(bf); - } -} - - - -__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { - const float2* sv2 = reinterpret_cast(sv); - #pragma unroll - for (int j = 0; j < 16; j++) { - asm("add.f32x2 %0, %1, %2;" - : "+l"(reinterpret_cast(*acc)) - : "l"(reinterpret_cast(*acc)), - "l"(reinterpret_cast(sv2[j]))); - } -} - - -__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { - asm("mul.rn.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { - asm("mul.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { - asm("add.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { - asm("add.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { - asm("sub.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { - asm("sub.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("add.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { - float2 r; - asm("sub.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("sub.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ void fma_scale_x32( - float* sv, const float2* scale2, const float2* neg_max2) -{ - float2* sv_2 = reinterpret_cast(sv); - #pragma unroll - for (int j = 0; j < 16; j++) - fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); -} - -__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("mul.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - - -// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) - -__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("add.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("add.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("add.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("add.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("add.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("add.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("add.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - - -__device__ __forceinline__ void fence_async_shared() { - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); -} - - -__device__ __forceinline__ uint64_t make_smem_desc(int addr) { - const int SBO = 1024; - return desc_encode(addr) - | (desc_encode(SBO) << 32ULL) - | (1ULL << 46ULL) - | (2ULL << 61ULL); -} - - -__device__ __forceinline__ void tma_2d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3}], [%4];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_4d_gmem2smem_cta2( - int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.4d.shared::cluster.global" - ".mbarrier::complete_tx::bytes.cta_group::2" - " [%0], [%1, {%2, %3, %4, %5}], [%6];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr) { - // Canonical .shared::cta form for non-multicast gather4, matching - // trtllm-gen / cuda_ptx and the PTX ISA qualifier order - // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, - // .shared::cluster is reserved for the multicast variant (ctaMask). - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem_cta2( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr) { - // Canonical .shared::cta form; see tma_gather4_gmem2smem above. - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" - ".mbarrier::complete_tx::bytes.cta_group::2" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem_mc( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr, unsigned short cta_mask) { - // Multicast variant: the PTX grammar ties the .shared::cluster - // destination to .multicast::cluster + ctaMask (cf. cuda_ptx / - // SM100_TMA_LOAD_MULTICAST_2D_GATHER4 in CUTLASS). - asm volatile( - "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" - ".mbarrier::complete_tx::bytes.multicast::cluster" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr), "h"(cta_mask) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem_mc_cta2( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr, unsigned short cta_mask) { - // Multicast + cta_group::2 variant; see tma_gather4_gmem2smem_mc. - asm volatile( - "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" - ".mbarrier::complete_tx::bytes.multicast::cluster.cta_group::2" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr), "h"(cta_mask) : "memory"); -} - - -__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { - asm volatile( - "{\n\t" - ".reg .b16 lo, hi;\n\t" - "mov.b32 {lo, hi}, %1;\n\t" - "tcgen05.commit.cta_group::2.mbarrier::arrive::one" - ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" - "}\n" - :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); -} - - -__device__ __forceinline__ void tmem_st_x8_u32(int addr, const uint32_t* src) { - asm volatile( - "tcgen05.st.sync.aligned.32x32b.x8.b32" - " [%0], {%1,%2,%3,%4,%5,%6,%7,%8};" - :: "r"(addr), - "r"(src[0]), "r"(src[1]), "r"(src[2]), "r"(src[3]), - "r"(src[4]), "r"(src[5]), "r"(src[6]), "r"(src[7])); -} - - -__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { - uint32_t result; - asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" - : "=r"(result) : "r"(val)); - return result; -} - -extern "C" { - -__global__ __launch_bounds__(512, 1) __cluster_dims__(2,1,1) void -kernel_cake_dsv4_fb6409904b35ed0f7a1a(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items) -{ - const int tid = threadIdx.x; - const int warp = make_warp_uniform(tid / 32); - const int lane = tid % 32; - - extern __shared__ __align__(1024) char smem_raw[]; - int smem; - asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); - smem = make_warp_uniform(smem); - - const int mbar_base = smem; - #define q_full_addr (mbar_base + 0) - #define q_empty_addr (mbar_base + 8) - #define k_full_addr (mbar_base + 16) - #define k_empty_addr (mbar_base + 32) - #define v_full_addr (mbar_base + 48) - #define v_empty_addr (mbar_base + 64) - #define index_full_addr (mbar_base + 80) - #define index_empty_addr (mbar_base + 128) - #define s_full_addr (mbar_base + 176) - #define p_full_addr (mbar_base + 192) - #define corr_done_addr (mbar_base + 208) - #define stats_addr (mbar_base + 224) - #define sum_ready_addr (mbar_base + 240) - #define o_done_addr (mbar_base + 248) - #define pv_done_addr (mbar_base + 256) - #define s_seeded_addr (mbar_base + 264) - #define q_pair_ready_addr (mbar_base + 272) - #define kv_pair_ready_addr (mbar_base + 280) - #define pv_pair_ready_addr (mbar_base + 288) - #define tmem_dealloc_addr (mbar_base + 304) - #define tmem_dealloc_peer_addr (mbar_base + 312) - - const int bid = blockIdx.x; - const int num_bids = gridDim.x; - const unsigned int clusters_x = gridDim.x / 2; - const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; - const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; - - int cta_rank; - asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); - - // Kernel setup ops - uint8_t* smem_q = reinterpret_cast(smem_raw + 1024); - const int smem_q_addr = smem + 1024; - uint8_t* smem_kv = reinterpret_cast(smem_raw + 33792); - const int smem_kv_addr = smem + 33792; - uint8_t* smem_v = reinterpret_cast(smem_raw + 33792); - const int smem_v_addr = smem + 33792; - uint8_t* smem_q_full = reinterpret_cast(smem_raw + 1024); - const int smem_q_full_addr = smem + 1024; - uint8_t* smem_k_full = reinterpret_cast(smem_raw + 33792); - const int smem_k_full_addr = smem + 33792; - uint8_t* smem_v_full = reinterpret_cast(smem_raw + 99328); - const int smem_v_full_addr = smem + 99328; - float* smem_stats_max = reinterpret_cast(smem_raw + 164864); - const int smem_stats_max_addr = smem + 164864; - float* smem_stats_sum = reinterpret_cast(smem_raw + 165888); - const int smem_stats_sum_addr = smem + 165888; - float* smem_stats_final_max = reinterpret_cast(smem_raw + 166400); - const int smem_stats_final_max_addr = smem + 166400; - float* smem_softmax_warp_pair_exchange = reinterpret_cast(smem_raw + 166912); - const int smem_softmax_warp_pair_exchange_addr = smem + 166912; - float* smem_corr_warp_pair_exchange = reinterpret_cast(smem_raw + 167936); - const int smem_corr_warp_pair_exchange_addr = smem + 167936; - int* smem_sparse_indices = reinterpret_cast(smem_raw + 168448); - const int smem_sparse_indices_addr = smem + 168448; - uint8_t* smem_p_fp8 = reinterpret_cast(smem_raw + 174592); - const int smem_p_fp8_addr = smem + 174592; - { - asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); - asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); - } - - // Mbarrier init (21 pipeline groups, 0 ordered-sequence groups, 40 barriers) - // Mbarriers at smem_raw[0..320) - - if (warp == 12) { - uint32_t leader = elect_sync(); - if (leader) { - // q_full: 1 barriers, init_count=2 - mbarrier_init(smem + 0, 2); - // q_empty: 1 barriers, init_count=1 - mbarrier_init(smem + 8, 1); - // --- pipeline 'k_pipe' --- - // k_full: 2 barriers, init_count=2 - mbarrier_init(smem + 16, 2); - mbarrier_init(smem + 24, 2); - // k_empty: 2 barriers, init_count=1 - mbarrier_init(smem + 32, 1); - mbarrier_init(smem + 40, 1); - // --- pipeline 'v_pipe' --- - // v_full: 2 barriers, init_count=2 - mbarrier_init(smem + 48, 2); - mbarrier_init(smem + 56, 2); - // v_empty: 2 barriers, init_count=1 - mbarrier_init(smem + 64, 1); - mbarrier_init(smem + 72, 1); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - if (warp == 9) { - uint32_t leader = elect_sync(); - if (leader) { - // --- pipeline 'index_pipe' --- - // index_full: 6 barriers, init_count=32 - mbarrier_init(smem + 80, 32); - mbarrier_init(smem + 88, 32); - mbarrier_init(smem + 96, 32); - mbarrier_init(smem + 104, 32); - mbarrier_init(smem + 112, 32); - mbarrier_init(smem + 120, 32); - // index_empty: 6 barriers, init_count=192 - mbarrier_init(smem + 128, 192); - mbarrier_init(smem + 136, 192); - mbarrier_init(smem + 144, 192); - mbarrier_init(smem + 152, 192); - mbarrier_init(smem + 160, 192); - mbarrier_init(smem + 168, 192); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - if (warp == 0) { - uint32_t leader = elect_sync(); - if (leader) { - // s_full: stages (0,), init_count=1 - mbarrier_init(smem + 176, 1); - // p_full: stages (0,), init_count=64 - mbarrier_init(smem + 192, 64); - // o_done: 1 barriers, init_count=1 - mbarrier_init(smem + 248, 1); - // pv_done: 1 barriers, init_count=1 - mbarrier_init(smem + 256, 1); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - if (warp == 1) { - uint32_t leader = elect_sync(); - if (leader) { - // s_full: stages (1,), init_count=1 - mbarrier_init(smem + 184, 1); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - if (warp == 2) { - uint32_t leader = elect_sync(); - if (leader) { - // p_full: stages (1,), init_count=64 - mbarrier_init(smem + 200, 64); - // s_seeded: 1 barriers, init_count=256 - mbarrier_init(smem + 264, 256); - // q_pair_ready: 1 barriers, init_count=64 - mbarrier_init(smem + 272, 64); - // --- pipeline 'kv_pipe' --- - // kv_pair_ready: 1 barriers, init_count=64 - mbarrier_init(smem + 280, 64); - // pv_pair_ready: stages (0,), init_count=64 - mbarrier_init(smem + 288, 64); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - if (warp == 4) { - uint32_t leader = elect_sync(); - if (leader) { - // corr_done: stages (0,), init_count=64 - mbarrier_init(smem + 208, 64); - // stats: stages (0,), init_count=64 - mbarrier_init(smem + 224, 64); - // sum_ready: 1 barriers, init_count=64 - mbarrier_init(smem + 240, 64); - // pv_pair_ready: stages (1,), init_count=64 - mbarrier_init(smem + 296, 64); - // tmem_dealloc: 1 barriers, init_count=288 - mbarrier_init(smem + 304, 288); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - if (warp == 6) { - uint32_t leader = elect_sync(); - if (leader) { - // corr_done: stages (1,), init_count=64 - mbarrier_init(smem + 216, 64); - // stats: stages (1,), init_count=64 - mbarrier_init(smem + 232, 64); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - if (warp == 5) { - uint32_t leader = elect_sync(); - if (leader) { - // tmem_dealloc_peer: 1 barriers, init_count=32 - mbarrier_init(smem + 312, 32); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - - __syncwarp(); - - asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); - asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); - - // TMEM alloc (512 columns, 512 used) - volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 320); - if (warp == 0) { - int _tmem_hold = smem + 320; - asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); - __syncwarp(); - asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); - } - - __syncthreads(); - asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); - asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); - - const int taddr = tmem_addr_storage[0]; - - // Kernel post-init ops - const int tmem_tmem_scratch = taddr; - - // ---- Ordered hardware-WG register redistribution ---- - // Inc phase consumes the registers released above. - if (warp >= 0 && warp <= 3) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 152;"); - } else if (warp >= 4 && warp <= 7) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 144;"); - } else if (warp >= 8 && warp <= 11) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 136;"); - } - - // ---- Role: index_warp ---- - if (warp == 9) { - { // index_warp_main - const int index_dummy = 0; - unsigned int _phase_index_empty = 1; - unsigned int _phase_q_full_0 = 0; - { - int all_num_kv_tiles = (sparse_topk + 128 - 1) / 128; - unsigned int index_prod_stage = 0; - #pragma unroll - for (unsigned int static_work = 0; static_work < 1; static_work++) { - int split_idx = 0; - int query_idx = bid / 2; - int tiles_per_split = all_num_kv_tiles + 1 - 1; - int first_tile = split_idx * tiles_per_split; - int num_index_passes = (tiles_per_split + 1) / 2; - #pragma unroll 1 - for (int index_pass = 0; index_pass < num_index_passes; index_pass++) { - mbarrier_wait(index_empty_addr + (index_prod_stage) * 8, _phase_index_empty); - int index_stage_base = smem_sparse_indices_addr + index_prod_stage * 1024; - int stage_row_base = index_pass * 256; - #pragma unroll - for (int index_half = 0; index_half < 2; index_half++) { - int index_offset = index_half * 128 + lane * 4; - int tile_col_base = first_tile * 128 + stage_row_base + index_half * 128; - int* row_tile_ptr = ((tile_col_base < 128) ? swa_indices + (query_idx * swa_index_stride) : compressed_indices + (query_idx * compressed_index_stride + (tile_col_base - 128))); - int lane_col = tile_col_base + lane * 4; - if (lane_col + 4 <= sparse_topk) { - asm volatile("cp.async.cg.shared::cta.global.L2::128B [%0], [%1], 16;" - :: "r"(index_stage_base + index_offset * 4), "l"(row_tile_ptr + (lane * 4))); - } else { - int tail_rows[4]; - #pragma unroll - for (int row_i = 0; row_i < 4; row_i++) { - tail_rows[row_i] = -1; - } - #pragma unroll - for (int row_i_1 = 0; row_i_1 < 4; row_i_1++) { - if (lane_col + row_i_1 < sparse_topk) { - tail_rows[row_i_1] = row_tile_ptr[lane * 4 + row_i_1]; - } - } - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(smem_sparse_indices_addr + ((unsigned int)index_stage_base - smem_sparse_indices_addr + (unsigned int)(index_offset * 4))), "r"(tail_rows[0]), "r"(tail_rows[1]), "r"(tail_rows[2]), "r"(tail_rows[3]) : "memory"); - } - } - asm volatile("cp.async.commit_group;"); - asm volatile("cp.async.wait_group 0;"); - mbarrier_arrive(index_full_addr + (index_prod_stage) * 8); - index_prod_stage += 1; - if (index_prod_stage == 6) { index_prod_stage = 0; _phase_index_empty ^= 1; } - } - } - } - } - } - // ---- Role: load_warp ---- - if (warp >= 12 && warp <= 15) { - asm volatile("setmaxnreg.dec.sync.aligned.u32 72;"); - { // load_warp_main - const int wg2_dummy = 0; - const int load_warp_rank = warp - ((1) ? 12 : 9); - int all_num_kv_tiles_1 = (sparse_topk + 128 - 1) / 128; - unsigned int load_k_stage = 0; - unsigned int load_v_stage = 0; - unsigned int load_kv_stage = 0; - unsigned int load_k_index_stage = 0; - unsigned int load_v_index_stage = 0; - int k_cta_offset = bid % 2 * 64; - unsigned int _phase_q_empty_0 = 1; - unsigned int _phase_index_full = 0; - unsigned int _phase_k_empty = 1; - unsigned int _phase_k_empty_1 = 1; - unsigned int _phase_v_empty = 1; - #pragma unroll - for (unsigned int work_idx = 0; work_idx < 1; work_idx++) { - int split_idx_1 = 0; - int query_idx_1 = bid / 2 >> 1; - int v_chunk = bid / 2 & 1; - { - split_idx_1 = 0; - query_idx_1 = bid / 2; - v_chunk = bid % 2; - } - int tiles_per_split_1 = all_num_kv_tiles_1 + 1 - 1; - int first_tile_1 = split_idx_1 * tiles_per_split_1; - int num_kv_tiles = tiles_per_split_1; - int sparse_extent = num_kv_tiles * 128; - mbarrier_wait(q_empty_addr, _phase_q_empty_0); - _phase_q_empty_0 ^= 1; - if (load_warp_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); - #pragma unroll - for (int q_stage = 0; q_stage < 4; q_stage++) { - asm volatile( - "cp.async.bulk.tensor.4d.shared::cluster.global.mbarrier::complete_tx::bytes.multicast::cluster.cta_group::2" - " [%0], [%1, {%2, %3, %4, %5}], [%6], %7;" - :: "r"(smem_q_full_addr + (unsigned int)(q_stage * 8192)), "l"((&tmap_q)), "r"(0), "r"(bid % 2 * 64), "r"(q_stage), "r"(query_idx_1), - "r"(((q_full_addr) & 0xFEFFFFFF)), "h"((uint16_t)(1 << bid % 2)) : "memory"); - } - } - } - { - mbarrier_wait(index_full_addr + (load_k_index_stage) * 8, _phase_index_full); - asm volatile("tcgen05.fence::after_thread_sync;"); - } - int k_index_stage_base = smem_sparse_indices_addr + load_k_index_stage * 1024; - { - mbarrier_wait(k_empty_addr + (load_k_stage) * 8, _phase_k_empty); - if (load_warp_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); - } - } - int k_dst = smem_k_full_addr + load_k_stage * 32768 + (unsigned int)(load_warp_rank * 8192); - int group = lane; - if (group < 16) { - int group_offset = k_cta_offset + group * 4; - int raw_rows[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows[0])), "=r"(*reinterpret_cast(&raw_rows[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows[(0) + 3])) - : "r"(k_index_stage_base + group_offset * 4)); - int row0 = ((raw_rows[0] >= 0) ? raw_rows[0] : 0); - int row1 = ((raw_rows[1] >= 0) ? raw_rows[1] : 0); - int row2 = ((raw_rows[2] >= 0) ? raw_rows[2] : 0); - int row3 = ((raw_rows[3] >= 0) ? raw_rows[3] : 0); - if (first_tile_1 == 0) { - tma_gather4_gmem2smem_mc_cta2(k_dst + group * 512, (&tmap_swa_kv), load_warp_rank * 128, row0, row1, row2, row3, ((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } else { - tma_gather4_gmem2smem_mc_cta2(k_dst + group * 512, (&tmap_compressed_kv), load_warp_rank * 128, row0, row1, row2, row3, ((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } - } - load_k_stage += 1; - if (load_k_stage == 2) { load_k_stage = 0; _phase_k_empty ^= 1; } - } - #pragma unroll 1 - for (int tile = 1; tile < num_kv_tiles; tile++) { - if ((1 & (int)((tile & 1) == 0)) != 0) { - mbarrier_wait(index_full_addr + (load_k_index_stage) * 8, _phase_index_full); - asm volatile("tcgen05.fence::after_thread_sync;"); - } - k_index_stage_base = smem_sparse_indices_addr + load_k_index_stage * 1024; - { - mbarrier_wait(k_empty_addr + (load_k_stage) * 8, _phase_k_empty); - if (load_warp_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); - } - } - int k_dst_1 = smem_k_full_addr + load_k_stage * 32768 + (unsigned int)(load_warp_rank * 8192); - int group_1 = lane; - if (group_1 < 16) { - int group_offset_1 = (tile & 1) * 128 + k_cta_offset + group_1 * 4; - int raw_rows_1[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_1[0])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 3])) - : "r"(k_index_stage_base + group_offset_1 * 4)); - int row0_1 = ((raw_rows_1[0] >= 0) ? raw_rows_1[0] : 0); - int row1_1 = ((raw_rows_1[1] >= 0) ? raw_rows_1[1] : 0); - int row2_1 = ((raw_rows_1[2] >= 0) ? raw_rows_1[2] : 0); - int row3_1 = ((raw_rows_1[3] >= 0) ? raw_rows_1[3] : 0); - tma_gather4_gmem2smem_mc_cta2(k_dst_1 + group_1 * 512, (&tmap_compressed_kv), load_warp_rank * 128, row0_1, row1_1, row2_1, row3_1, ((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } - load_k_stage += 1; - if (load_k_stage == 2) { load_k_stage = 0; _phase_k_empty ^= 1; } - if ((tile & 1) != 0) { - load_k_index_stage += 1; - if (load_k_index_stage == 6) { load_k_index_stage = 0; _phase_index_full ^= 1; } - } - } - int prev_tile = tile - 1; - int global_prev_tile = first_tile_1 + prev_tile; - int v_index_stage_base = smem_sparse_indices_addr + load_v_index_stage * 1024; - { - mbarrier_wait(v_empty_addr + (load_v_stage) * 8, _phase_v_empty); - if (load_warp_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); - } - } - if (lane < 16) { - int v_dst = smem_v_full_addr + load_v_stage * 32768 + (unsigned int)(load_warp_rank / 2 * 16384); - int v_col = v_chunk * 256 + load_warp_rank / 2 * 128; - int group_2 = (load_warp_rank & 1) * 16 + lane; - int group_offset_2 = (prev_tile & 1) * 128 + group_2 * 4; - int raw_rows_2[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_2[0])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 3])) - : "r"(v_index_stage_base + group_offset_2 * 4)); - int row0_2 = ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : 0); - int row1_2 = ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : 0); - int row2_2 = ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : 0); - int row3_2 = ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : 0); - if (global_prev_tile == 0) { - tma_gather4_gmem2smem_mc_cta2(v_dst + group_2 * 512, (&tmap_swa_kv), v_col, row0_2, row1_2, row2_2, row3_2, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } else { - tma_gather4_gmem2smem_mc_cta2(v_dst + group_2 * 512, (&tmap_compressed_kv), v_col, row0_2, row1_2, row2_2, row3_2, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } - } - load_v_stage += 1; - if (load_v_stage == 2) { load_v_stage = 0; _phase_v_empty ^= 1; } - if ((prev_tile & 1) != 0) { - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(index_empty_addr + (load_v_index_stage) * 8); - load_v_index_stage = load_v_index_stage + 1; - if (load_v_index_stage == 6) { - load_v_index_stage = 0; - } - } - } - } - int last_tile = num_kv_tiles - 1; - int global_last_tile = first_tile_1 + last_tile; - { - mbarrier_wait(v_empty_addr + (load_v_stage) * 8, _phase_v_empty); - if (load_warp_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); - } - } - if (lane < 16) { - int v_dst_1 = smem_v_full_addr + load_v_stage * 32768 + (unsigned int)(load_warp_rank / 2 * 16384); - int v_col_1 = v_chunk * 256 + load_warp_rank / 2 * 128; - int group_3 = (load_warp_rank & 1) * 16 + lane; - int group_offset_3 = (last_tile & 1) * 128 + group_3 * 4; - int raw_rows_3[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_3[0])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 3])) - : "r"(smem_sparse_indices_addr + load_v_index_stage * 1024 + (unsigned int)(group_offset_3 * 4))); - int row0_3 = ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : 0); - int row1_3 = ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : 0); - int row2_3 = ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : 0); - int row3_3 = ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : 0); - if (global_last_tile == 0) { - tma_gather4_gmem2smem_mc_cta2(v_dst_1 + group_3 * 512, (&tmap_swa_kv), v_col_1, row0_3, row1_3, row2_3, row3_3, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } else { - tma_gather4_gmem2smem_mc_cta2(v_dst_1 + group_3 * 512, (&tmap_compressed_kv), v_col_1, row0_3, row1_3, row2_3, row3_3, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } - } - load_v_stage += 1; - if (load_v_stage == 2) { load_v_stage = 0; _phase_v_empty ^= 1; } - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(index_empty_addr + (load_v_index_stage) * 8); - load_v_index_stage = load_v_index_stage + 1; - if (load_v_index_stage == 6) { - load_v_index_stage = 0; - } - } - } - mbarrier_arrive(tmem_dealloc_addr); - } - } - // ---- Role: softmax_wg ---- - if (warp == 0 || warp == 2) { - { // softmax_wg_main - float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; - const int wg_dummy_inc = 0; - int all_num_kv_tiles_2 = (sparse_topk + 128 - 1) / 128; - const int tmem_row_base = ((1) ? warp % 2 * 32 : warp % 4 * 32); - const int tmem_score_row_base = ((1) ? (int)(warp % 4 * 32) : tmem_row_base); - const int n_half = ((1) ? (int)(warp % 4 / 2) : 0); - const int my_row = tmem_row_base + lane; - const int stats_row = n_half * 64 + my_row; - int softmax_tile_cursor = 0; - unsigned int softmax_index_stage = 0; - unsigned int _phase_q_full_0_1 = 0; - #pragma unroll - for (unsigned int work_idx_1 = 0; work_idx_1 < 1; work_idx_1++) { - int split_idx_2 = 0; - int query_idx_2 = bid / 2 >> 1; - { - split_idx_2 = 0; - query_idx_2 = bid / 2; - } - int tiles_per_split_2 = all_num_kv_tiles_2 + 1 - 1; - int first_tile_2 = split_idx_2 * tiles_per_split_2; - int num_kv_tiles_1 = tiles_per_split_2; - int _max_0 = ((sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) : (0)); - int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); - int active_topk = _min_0; - float row_max_val = -CAKE_INF; - float row_sum_val = 0.0f; - int sink_head = ((1) ? bid % 2 * 64 + my_row : my_row); - if (has_sinks != 0 && sink_head < num_heads && split_idx_2 == 0) { - row_max_val = sinks[sink_head] * 1.4426950408889634f / softmax_scale_log2; - row_sum_val = 1.0f; - } - #pragma unroll 1 - for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { - int pipeline_tile = softmax_tile_cursor + tile_1; - int phase = pipeline_tile & 1; - int s_wait_phase = pipeline_tile >> 1 & 1; - mbarrier_wait(s_full_addr + (phase) * 8, s_wait_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - int s_off = ((phase != 0) ? 128 : 0); - int s_base = taddr + (unsigned int)s_off + (unsigned int)(tmem_score_row_base << 16); - float new_max = row_max_val; - int valid_sparse_cols = ((active_topk < sparse_topk) ? active_topk : sparse_topk); - valid_sparse_cols = valid_sparse_cols - (first_tile_2 + tile_1) * 128 - n_half * 64; - if (valid_sparse_cols < 0) { - valid_sparse_cols = 0; - } - if (valid_sparse_cols > 64) { - valid_sparse_cols = 64; - } - int index_tile_addr = smem_sparse_indices_addr + softmax_index_stage * 1024 + (unsigned int)((tile_1 & 1) * 512); - int index_lane_addr = index_tile_addr + (n_half * 64 + lane) * 4; - int staged_index[2]; - unsigned int invalid_cols[2]; - unsigned int any_invalid = 0; - #pragma unroll - for (int w = 0; w < 2; w++) { - asm volatile("ld.shared.b32 %0, [%1];" : "=r"(*reinterpret_cast(&staged_index[w])) : "r"(index_lane_addr + w * 128)); - unsigned int _vote_0 = __ballot_sync(0xFFFFFFFF, staged_index[w] < 0); - invalid_cols[w] = _vote_0; - any_invalid = any_invalid | invalid_cols[w]; - } - if ((tile_1 & 1) != 0 || tile_1 + 1 == num_kv_tiles_1) { - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(index_empty_addr + (softmax_index_stage) * 8); - softmax_index_stage = softmax_index_stage + 1; - if (softmax_index_stage == 6) { - softmax_index_stage = 0; - } - } - float _tmem_load_0[4]; - tmem_ld_x4(&_tmem_load_0[0], s_base); - float _tmem_load_1[64]; - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]) - : "r"(s_base)); - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]) - : "r"(s_base + 32)); - { - uint32_t _slice_lo_mask_0; - { - int _lim_0 = valid_sparse_cols; - if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } - else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); - } - } - if (!(_slice_lo_mask_0 & (1u << 0))) _tmem_load_1[0] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 1))) _tmem_load_1[1] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 2))) _tmem_load_1[2] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 3))) _tmem_load_1[3] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 4))) _tmem_load_1[4] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 5))) _tmem_load_1[5] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 6))) _tmem_load_1[6] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 7))) _tmem_load_1[7] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 8))) _tmem_load_1[8] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 9))) _tmem_load_1[9] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 10))) _tmem_load_1[10] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 11))) _tmem_load_1[11] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 12))) _tmem_load_1[12] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 13))) _tmem_load_1[13] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 14))) _tmem_load_1[14] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 15))) _tmem_load_1[15] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 16))) _tmem_load_1[16] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 17))) _tmem_load_1[17] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 18))) _tmem_load_1[18] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 19))) _tmem_load_1[19] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 20))) _tmem_load_1[20] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 21))) _tmem_load_1[21] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 22))) _tmem_load_1[22] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 23))) _tmem_load_1[23] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 24))) _tmem_load_1[24] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 25))) _tmem_load_1[25] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 26))) _tmem_load_1[26] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 27))) _tmem_load_1[27] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 28))) _tmem_load_1[28] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 29))) _tmem_load_1[29] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 30))) _tmem_load_1[30] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 31))) _tmem_load_1[31] = -CAKE_INF; - uint32_t _slice_lo_mask_1; - { - int _lim_1 = valid_sparse_cols - 32; - if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } - else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); - } - } - if (!(_slice_lo_mask_1 & (1u << 0))) _tmem_load_1[32] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 1))) _tmem_load_1[33] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 2))) _tmem_load_1[34] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 3))) _tmem_load_1[35] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 4))) _tmem_load_1[36] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 5))) _tmem_load_1[37] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 6))) _tmem_load_1[38] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 7))) _tmem_load_1[39] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 8))) _tmem_load_1[40] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 9))) _tmem_load_1[41] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 10))) _tmem_load_1[42] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 11))) _tmem_load_1[43] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 12))) _tmem_load_1[44] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 13))) _tmem_load_1[45] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 14))) _tmem_load_1[46] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 15))) _tmem_load_1[47] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 16))) _tmem_load_1[48] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 17))) _tmem_load_1[49] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 18))) _tmem_load_1[50] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 19))) _tmem_load_1[51] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 20))) _tmem_load_1[52] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 21))) _tmem_load_1[53] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 22))) _tmem_load_1[54] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 23))) _tmem_load_1[55] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 24))) _tmem_load_1[56] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 25))) _tmem_load_1[57] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 26))) _tmem_load_1[58] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 27))) _tmem_load_1[59] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 28))) _tmem_load_1[60] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 29))) _tmem_load_1[61] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 30))) _tmem_load_1[62] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 31))) _tmem_load_1[63] = -CAKE_INF; - if (any_invalid != 0) { - #pragma unroll - for (int w_1 = 0; w_1 < 2; w_1++) { - unsigned int mask_word = invalid_cols[w_1]; - _tmem_load_1[w_1 * 32] = (((mask_word & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32]); - _tmem_load_1[w_1 * 32 + 1] = (((mask_word >> 1 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 1]); - _tmem_load_1[w_1 * 32 + 2] = (((mask_word >> 2 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 2]); - _tmem_load_1[w_1 * 32 + 3] = (((mask_word >> 3 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 3]); - _tmem_load_1[w_1 * 32 + 4] = (((mask_word >> 4 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 4]); - _tmem_load_1[w_1 * 32 + 5] = (((mask_word >> 5 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 5]); - _tmem_load_1[w_1 * 32 + 6] = (((mask_word >> 6 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 6]); - _tmem_load_1[w_1 * 32 + 7] = (((mask_word >> 7 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 7]); - _tmem_load_1[w_1 * 32 + 8] = (((mask_word >> 8 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 8]); - _tmem_load_1[w_1 * 32 + 9] = (((mask_word >> 9 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 9]); - _tmem_load_1[w_1 * 32 + 10] = (((mask_word >> 10 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 10]); - _tmem_load_1[w_1 * 32 + 11] = (((mask_word >> 11 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 11]); - _tmem_load_1[w_1 * 32 + 12] = (((mask_word >> 12 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 12]); - _tmem_load_1[w_1 * 32 + 13] = (((mask_word >> 13 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 13]); - _tmem_load_1[w_1 * 32 + 14] = (((mask_word >> 14 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 14]); - _tmem_load_1[w_1 * 32 + 15] = (((mask_word >> 15 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 15]); - _tmem_load_1[w_1 * 32 + 16] = (((mask_word >> 16 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 16]); - _tmem_load_1[w_1 * 32 + 17] = (((mask_word >> 17 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 17]); - _tmem_load_1[w_1 * 32 + 18] = (((mask_word >> 18 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 18]); - _tmem_load_1[w_1 * 32 + 19] = (((mask_word >> 19 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 19]); - _tmem_load_1[w_1 * 32 + 20] = (((mask_word >> 20 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 20]); - _tmem_load_1[w_1 * 32 + 21] = (((mask_word >> 21 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 21]); - _tmem_load_1[w_1 * 32 + 22] = (((mask_word >> 22 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 22]); - _tmem_load_1[w_1 * 32 + 23] = (((mask_word >> 23 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 23]); - _tmem_load_1[w_1 * 32 + 24] = (((mask_word >> 24 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 24]); - _tmem_load_1[w_1 * 32 + 25] = (((mask_word >> 25 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 25]); - _tmem_load_1[w_1 * 32 + 26] = (((mask_word >> 26 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 26]); - _tmem_load_1[w_1 * 32 + 27] = (((mask_word >> 27 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 27]); - _tmem_load_1[w_1 * 32 + 28] = (((mask_word >> 28 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 28]); - _tmem_load_1[w_1 * 32 + 29] = (((mask_word >> 29 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 29]); - _tmem_load_1[w_1 * 32 + 30] = (((mask_word >> 30 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 30]); - _tmem_load_1[w_1 * 32 + 31] = (((mask_word >> 31 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 31]); - } - } - float2 _reg_reduce_max2_2 = {-CAKE_INF, -CAKE_INF}; - row_max_x32_accum(&_tmem_load_1[0], _reg_reduce_max2_2); - row_max_x32_accum(&_tmem_load_1[32], _reg_reduce_max2_2); - float _tmem_load_1_max = row_max_reduce(_reg_reduce_max2_2); - float _max_1 = max_noftz(new_max, _tmem_load_1_max); - new_max = _max_1; - smem_softmax_warp_pair_exchange[phase * 128 + stats_row] = new_max; - asm volatile("barrier.sync %0, 64;" :: "r"(2 + warp % 2) : "memory"); - float _max_2 = max_noftz(new_max, smem_softmax_warp_pair_exchange[phase * 128 + (stats_row ^ 64)]); - new_max = _max_2; - } - float _fma_0 = __fmaf_rn(row_max_val, softmax_scale_log2, (-new_max) * softmax_scale_log2); - float delta = _fma_0; - float _exp2_0 = approx_exp2(delta); - float exp_delta = _exp2_0; - float acc_scale = ((row_max_val > -CAKE_INF) ? exp_delta : 1.0f); - smem_stats_max[phase * 128 + stats_row] = acc_scale; - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(stats_addr + (phase) * 8); - row_max_val = new_max; - float safe_max = ((new_max == -CAKE_INF) ? 0.0f : new_max); - float max_scaled = safe_max * softmax_scale_log2; - float block_sum = 0.0f; - { - { - const float2 _fma_b2_3 = {softmax_scale_log2, softmax_scale_log2}; - const float2 _fma_c2_4 = {-max_scaled, -max_scaled}; - #pragma unroll - for (int _lf = 0; _lf < 32; _lf++) - fma_f32x2_inplace(&reinterpret_cast(_tmem_load_1)[_lf], _fma_b2_3, _fma_c2_4); - #pragma unroll - for (int _le = 0; _le < 64; _le++) { - _tmem_load_1[_le] = approx_exp2(_tmem_load_1[_le]); - } - { - const float2* _reg_reduce_src2_5 = reinterpret_cast(&_tmem_load_1[0]); - float2 _reg_reduce_sum2_5_0 = make_float2(0.0f, 0.0f); - float2 _reg_reduce_sum2_5_1 = make_float2(0.0f, 0.0f); - float2 _reg_reduce_sum2_5_2 = make_float2(0.0f, 0.0f); - float2 _reg_reduce_sum2_5_3 = make_float2(0.0f, 0.0f); - #pragma unroll - for (int _rr = 0; _rr < 8; _rr++) { - add_f32x2_inplace(&_reg_reduce_sum2_5_0, _reg_reduce_src2_5[_rr * 4 + 0]); - add_f32x2_inplace(&_reg_reduce_sum2_5_1, _reg_reduce_src2_5[_rr * 4 + 1]); - add_f32x2_inplace(&_reg_reduce_sum2_5_2, _reg_reduce_src2_5[_rr * 4 + 2]); - add_f32x2_inplace(&_reg_reduce_sum2_5_3, _reg_reduce_src2_5[_rr * 4 + 3]); - } - add_f32x2_inplace(&_reg_reduce_sum2_5_0, _reg_reduce_sum2_5_1); - add_f32x2_inplace(&_reg_reduce_sum2_5_2, _reg_reduce_sum2_5_3); - add_f32x2_inplace(&_reg_reduce_sum2_5_0, _reg_reduce_sum2_5_2); - float _tmem_load_1_sum = _reg_reduce_sum2_5_0.x + _reg_reduce_sum2_5_0.y; - block_sum = _tmem_load_1_sum; - } - uint32_t _fp8_0[16]; - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[0]), "f"(_tmem_load_1[1]), - "f"(_tmem_load_1[2]), "f"(_tmem_load_1[3])); - _fp8_0[0] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[4]), "f"(_tmem_load_1[5]), - "f"(_tmem_load_1[6]), "f"(_tmem_load_1[7])); - _fp8_0[1] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[8]), "f"(_tmem_load_1[9]), - "f"(_tmem_load_1[10]), "f"(_tmem_load_1[11])); - _fp8_0[2] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[12]), "f"(_tmem_load_1[13]), - "f"(_tmem_load_1[14]), "f"(_tmem_load_1[15])); - _fp8_0[3] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[16]), "f"(_tmem_load_1[17]), - "f"(_tmem_load_1[18]), "f"(_tmem_load_1[19])); - _fp8_0[4] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[20]), "f"(_tmem_load_1[21]), - "f"(_tmem_load_1[22]), "f"(_tmem_load_1[23])); - _fp8_0[5] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[24]), "f"(_tmem_load_1[25]), - "f"(_tmem_load_1[26]), "f"(_tmem_load_1[27])); - _fp8_0[6] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[28]), "f"(_tmem_load_1[29]), - "f"(_tmem_load_1[30]), "f"(_tmem_load_1[31])); - _fp8_0[7] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[32]), "f"(_tmem_load_1[33]), - "f"(_tmem_load_1[34]), "f"(_tmem_load_1[35])); - _fp8_0[8] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[36]), "f"(_tmem_load_1[37]), - "f"(_tmem_load_1[38]), "f"(_tmem_load_1[39])); - _fp8_0[9] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[40]), "f"(_tmem_load_1[41]), - "f"(_tmem_load_1[42]), "f"(_tmem_load_1[43])); - _fp8_0[10] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[44]), "f"(_tmem_load_1[45]), - "f"(_tmem_load_1[46]), "f"(_tmem_load_1[47])); - _fp8_0[11] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[48]), "f"(_tmem_load_1[49]), - "f"(_tmem_load_1[50]), "f"(_tmem_load_1[51])); - _fp8_0[12] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[52]), "f"(_tmem_load_1[53]), - "f"(_tmem_load_1[54]), "f"(_tmem_load_1[55])); - _fp8_0[13] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[56]), "f"(_tmem_load_1[57]), - "f"(_tmem_load_1[58]), "f"(_tmem_load_1[59])); - _fp8_0[14] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[60]), "f"(_tmem_load_1[61]), - "f"(_tmem_load_1[62]), "f"(_tmem_load_1[63])); - _fp8_0[15] = _packed; - } - int p_stage_base = smem_p_fp8_addr + (unsigned int)(phase * 8192); - #pragma unroll - for (int p_vec = 0; p_vec < 4; p_vec++) { - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_stage_base + (my_row * 128 + (n_half * 4 + p_vec) * 16 ^ (my_row * 128 + (n_half * 4 + p_vec) * 16 >> 7 & 7) << 4))), "r"(_fp8_0[p_vec * 4]), "r"(_fp8_0[p_vec * 4 + 1]), "r"(_fp8_0[p_vec * 4 + 2]), "r"(_fp8_0[p_vec * 4 + 3]) : "memory"); - } - } - } - { - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - } - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - { - asm volatile("" ::: "memory"); - } - mbarrier_arrive(p_full_addr + (phase) * 8); - { - asm volatile("" ::: "memory"); - float _fma_3 = __fmaf_rn(row_sum_val, acc_scale, block_sum); - row_sum_val = _fma_3; - } - } - smem_stats_sum[stats_row] = row_sum_val; - smem_stats_final_max[stats_row] = row_max_val; - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(sum_ready_addr); - softmax_tile_cursor = softmax_tile_cursor + num_kv_tiles_1; - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(tmem_dealloc_addr); - } - } - // ---- Role: softmax_padding ---- - if (warp == 1 || warp == 3) { - { // softmax_padding_main - const int lowhead_softmax_padding = 0; - } - } - // ---- Role: correction_wg ---- - if (warp == 4 || warp == 6) { - { // correction_wg_main - float softmax_scale_log2_1 = bmm1_scale[0] * 1.4426950408889634f; - float output_scale = bmm2_scale[0]; - const int wg_dummy_inc_1 = 0; - int all_num_kv_tiles_3 = (sparse_topk + 128 - 1) / 128; - const int tmem_row_base_1 = ((1) ? warp % 2 * 32 : warp % 4 * 32); - const int n_half_1 = ((1) ? (int)(warp % 4 / 2) : 0); - const int my_row_1 = tmem_row_base_1 + lane; - const int stats_row_1 = n_half_1 * 64 + my_row_1; - const int corr_row = tmem_row_base_1 << 16; - int correction_tile_cursor = 0; - unsigned int _phase_q_full_0_2 = 0; - unsigned int _phase_pv_done_0 = 0; - unsigned int _phase_o_done_0 = 0; - unsigned int _phase_sum_ready_0 = 0; - #pragma unroll - for (unsigned int work_idx_2 = 0; work_idx_2 < 1; work_idx_2++) { - int split_idx_3 = 0; - int query_idx_3 = bid / 2 >> 1; - int v_chunk_1 = bid / 2 & 1; - { - split_idx_3 = 0; - query_idx_3 = bid / 2; - v_chunk_1 = bid % 2; - } - int tiles_per_split_3 = all_num_kv_tiles_3 + 1 - 1; - int first_tile_3 = split_idx_3 * tiles_per_split_3; - int num_kv_tiles_2 = tiles_per_split_3; - #pragma unroll 1 - for (int tile_2 = 0; tile_2 < num_kv_tiles_2; tile_2++) { - int pipeline_tile_1 = correction_tile_cursor + tile_2; - int phase_1 = pipeline_tile_1 & 1; - int stats_wait_phase = pipeline_tile_1 >> 1 & 1; - mbarrier_wait(stats_addr + (phase_1) * 8, stats_wait_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - float acc_scale_1 = smem_stats_max[phase_1 * 128 + stats_row_1]; - if (tile_2 > 0) { - mbarrier_wait(pv_done_addr, _phase_pv_done_0); - _phase_pv_done_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - } - if (tile_2 > 0) { - int _vote_1 = __any_sync(0xFFFFFFFF, acc_scale_1 < 1.0f); - int any_rescale = _vote_1; - if (any_rescale != 0) { - #pragma unroll - for (int vs = 0; vs < 2; vs++) { - int o_base = taddr + 256 + (unsigned int)(vs * 128) + (unsigned int)corr_row; - #pragma unroll - for (int c = 0; c < 128; c += 16) { - float _tmem_load_4[16]; - tmem_ld_x16(&_tmem_load_4[0], o_base + c); - const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; - #pragma unroll - for (int _ls = 0; _ls < 8; _ls++) - mul_f32x2_inplace(&reinterpret_cast(_tmem_load_4)[_ls], _scale2_0); - tmem_st_x16_f32(o_base + c, _tmem_load_4); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - } - } - } - } - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(corr_done_addr + (phase_1) * 8); - } - mbarrier_wait(o_done_addr, _phase_o_done_0); - _phase_o_done_0 ^= 1; - mbarrier_wait(sum_ready_addr, _phase_sum_ready_0); - _phase_sum_ready_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - float total_sum = smem_stats_sum[stats_row_1]; - float final_max = smem_stats_final_max[stats_row_1]; - { - smem_corr_warp_pair_exchange[stats_row_1] = total_sum; - asm volatile("barrier.sync %0, 64;" :: "r"(4 + warp % 2) : "memory"); - total_sum = total_sum + smem_corr_warp_pair_exchange[stats_row_1 ^ 64]; - } - float _rcp_0 = approx_rcp(total_sum); - float inv_sum = ((total_sum > 0.0f) ? _rcp_0 : 0.0f); - int head_idx = my_row_1; - { - head_idx = bid % 2 * 64 + my_row_1; - } - int direct_o_offset = (query_idx_3 * num_heads + head_idx) * 512 + v_chunk_1 * 256; - int partial_o_offset = (query_idx_3 * num_heads + head_idx + split_idx_3) * 512 + ((0) ? v_chunk_1 : n_half_1) * 256; - int o_offset = partial_o_offset; - #pragma unroll - for (int vs_1 = 0; vs_1 < 2; vs_1++) { - int o_base_epi = taddr + 256 + (unsigned int)(vs_1 * 128) + (unsigned int)corr_row; - #pragma unroll - for (int c_1 = 0; c_1 < 128; c_1 += 32) { - float _tmem_load_5[32]; - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(_tmem_load_5[0]), "=f"(_tmem_load_5[1]), "=f"(_tmem_load_5[2]), "=f"(_tmem_load_5[3]), "=f"(_tmem_load_5[4]), "=f"(_tmem_load_5[5]), "=f"(_tmem_load_5[6]), "=f"(_tmem_load_5[7]), "=f"(_tmem_load_5[8]), "=f"(_tmem_load_5[9]), "=f"(_tmem_load_5[10]), "=f"(_tmem_load_5[11]), "=f"(_tmem_load_5[12]), "=f"(_tmem_load_5[13]), "=f"(_tmem_load_5[14]), "=f"(_tmem_load_5[15]), "=f"(_tmem_load_5[16]), "=f"(_tmem_load_5[17]), "=f"(_tmem_load_5[18]), "=f"(_tmem_load_5[19]), "=f"(_tmem_load_5[20]), "=f"(_tmem_load_5[21]), "=f"(_tmem_load_5[22]), "=f"(_tmem_load_5[23]), "=f"(_tmem_load_5[24]), "=f"(_tmem_load_5[25]), "=f"(_tmem_load_5[26]), "=f"(_tmem_load_5[27]), "=f"(_tmem_load_5[28]), "=f"(_tmem_load_5[29]), "=f"(_tmem_load_5[30]), "=f"(_tmem_load_5[31]) - : "r"(o_base_epi + c_1)); - int gmem_base = o_offset + vs_1 * 128 + c_1; - if (head_idx < num_heads) { - #pragma unroll - for (int j = 0; j < 32; j += 8) { - { - const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; - #if __CUDA_ARCH__ >= 1000 - #pragma unroll - for (int _ps = 0; _ps < 4; _ps++) - mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_5[j])[_ps], _prescale2_1); - #else - #pragma unroll - for (int _ps = 0; _ps < 8; _ps++) - _tmem_load_5[j + _ps] *= inv_sum * output_scale; - #endif - __nv_bfloat162 _pk[4]; - _pk[0] = __floats2bfloat162_rn(_tmem_load_5[j + 0], _tmem_load_5[j + 1]); - _pk[1] = __floats2bfloat162_rn(_tmem_load_5[j + 2], _tmem_load_5[j + 3]); - _pk[2] = __floats2bfloat162_rn(_tmem_load_5[j + 4], _tmem_load_5[j + 5]); - _pk[3] = __floats2bfloat162_rn(_tmem_load_5[j + 6], _tmem_load_5[j + 7]); - *reinterpret_cast(&((__nv_bfloat16*)(O + (gmem_base + j)))[0]) = *reinterpret_cast(&_pk[0]); - } - } - } - } - } - correction_tile_cursor = correction_tile_cursor + num_kv_tiles_2; - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(tmem_dealloc_addr); - } - } - // ---- Role: correction_padding ---- - if (warp == 5 || warp == 7) { - { // correction_padding_main - const int lowhead_correction_padding = 0; - } - } - // ---- Role: mma_warp ---- - if (warp == 8) { - { // mma_warp_main - const int wg2_dummy_1 = 0; - int all_num_kv_tiles_4 = (sparse_topk + 128 - 1) / 128; - unsigned int mma_k_stage = 0; - unsigned int mma_v_stage = 0; - unsigned int mma_kv_stage = 0; - int mma_tile_cursor = 0; - unsigned int _phase_s_seeded_0 = 0; - unsigned int _phase_q_full_0_3 = 0; - unsigned int _phase_q_pair_ready_0 = 0; - unsigned int _phase_k_full = 0; - unsigned int _phase_k_full_1 = 0; - unsigned int _phase_kv_pair_ready = 0; - unsigned int _phase_v_full = 0; - #pragma unroll - for (unsigned int work_idx_3 = 0; work_idx_3 < 1; work_idx_3++) { - int split_idx_4 = 0; - { - split_idx_4 = 0; - } - int tiles_per_split_4 = all_num_kv_tiles_4 + 1 - 1; - int first_tile_4 = split_idx_4 * tiles_per_split_4; - int num_kv_tiles_3 = tiles_per_split_4; - { - if (cta_rank == 0) { - mbarrier_wait(q_full_addr, _phase_q_full_0_3); - _phase_q_full_0_3 ^= 1; - } - } - int first_pv = 1; - #pragma unroll 1 - for (int tile_3 = 0; tile_3 < num_kv_tiles_3; tile_3++) { - int pipeline_tile_2 = mma_tile_cursor + tile_3; - int phase_2 = pipeline_tile_2 & 1; - int score_col = ((phase_2 != 0) ? 128 : 0); - { - if (cta_rank == 0) { - mbarrier_wait(k_full_addr + (mma_k_stage) * 8, _phase_k_full); - int _mma_a_lo_0 = (((smem_q_full_addr) >> 4) & 0x3FFF) + (0) * 2048; - int _mma_b_lo_0 = (((smem_k_full_addr) >> 4) & 0x3FFF) + (mma_k_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 136314896;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (score_col))), "r"(0)); - int _mma_a_lo_1 = (((smem_q_full_addr + 8192) >> 4) & 0x3FFF) + (0) * 2048; - int _mma_b_lo_1 = (((smem_k_full_addr + 8192) >> 4) & 0x3FFF) + (mma_k_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 136314896;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); - int _mma_a_lo_2 = (((smem_q_full_addr + 16384) >> 4) & 0x3FFF) + (0) * 2048; - int _mma_b_lo_2 = (((smem_k_full_addr + 16384) >> 4) & 0x3FFF) + (mma_k_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 136314896;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); - int _mma_a_lo_3 = (((smem_q_full_addr + 24576) >> 4) & 0x3FFF) + (0) * 2048; - int _mma_b_lo_3 = (((smem_k_full_addr + 24576) >> 4) & 0x3FFF) + (mma_k_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 136314896;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); - if (cta_rank == 0) { - elect_commit_cg2_multicast(s_full_addr + (phase_2) * 8, (uint16_t)(3)); - } - if (cta_rank == 0) { - elect_commit_cg2_multicast(k_empty_addr + (mma_k_stage) * 8, (uint16_t)(3)); - } - } - mma_k_stage += 1; - if (mma_k_stage == 2) { mma_k_stage = 0; _phase_k_full ^= 1; } - } - if (tile_3 > 0) { - int prev_pipeline_tile = pipeline_tile_2 - 1; - int prev_phase = prev_pipeline_tile & 1; - int pv_wait_phase = prev_pipeline_tile >> 1 & 1; - mbarrier_wait(p_full_addr + (prev_phase) * 8, pv_wait_phase); - mbarrier_wait(corr_done_addr + (prev_phase) * 8, pv_wait_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - { - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cluster.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(pv_pair_ready_addr + (unsigned int)(prev_phase * 8)), "r"(0) : "memory"); - if (cta_rank == 0) { - mbarrier_wait_cluster_hint(pv_pair_ready_addr + (prev_phase) * 8, pv_wait_phase, 10000000); - } - } - int p_col = ((prev_phase != 0) ? 224 : 96); - { - if (cta_rank == 0) { - mbarrier_wait(v_full_addr + (mma_v_stage) * 8, _phase_v_full); - int _mma_a_lo_6 = (((smem_p_fp8_addr) >> 4) & 0x3FFF) + (prev_phase) * 512; - int _mma_b_lo_6 = ((((smem_v_full_addr) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 138477584;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_6), "r"(_mma_b_lo_6), "r"((tmem_tmem_scratch + (256))), "r"(((first_pv) ? 0 : 1))); - int _mma_a_lo_7 = (((smem_p_fp8_addr) >> 4) & 0x3FFF) + (prev_phase) * 512; - int _mma_b_lo_7 = ((((smem_v_full_addr + 16384) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 138477584;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_7), "r"(_mma_b_lo_7), "r"((tmem_tmem_scratch + (384))), "r"(((first_pv) ? 0 : 1))); - if (cta_rank == 0) { - elect_commit_cg2_multicast(v_empty_addr + (mma_v_stage) * 8, (uint16_t)(3)); - } - } - mma_v_stage += 1; - if (mma_v_stage == 2) { mma_v_stage = 0; _phase_v_full ^= 1; } - } - first_pv = 0; - if (cta_rank == 0) { - elect_commit_cg2_multicast(pv_done_addr, (uint16_t)(3)); - } - } - } - int last_pipeline_tile = mma_tile_cursor + num_kv_tiles_3 - 1; - int last_phase = last_pipeline_tile & 1; - int drain_wait_phase = last_pipeline_tile >> 1 & 1; - mbarrier_wait(p_full_addr + (last_phase) * 8, drain_wait_phase); - mbarrier_wait(corr_done_addr + (last_phase) * 8, drain_wait_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - { - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cluster.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(pv_pair_ready_addr + (unsigned int)(last_phase * 8)), "r"(0) : "memory"); - if (cta_rank == 0) { - mbarrier_wait_cluster_hint(pv_pair_ready_addr + (last_phase) * 8, drain_wait_phase, 10000000); - } - } - int p_col_last = ((last_phase != 0) ? 224 : 96); - { - if (cta_rank == 0) { - mbarrier_wait(v_full_addr + (mma_v_stage) * 8, _phase_v_full); - int _mma_a_lo_10 = (((smem_p_fp8_addr) >> 4) & 0x3FFF) + (last_phase) * 512; - int _mma_b_lo_10 = ((((smem_v_full_addr) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 138477584;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_10), "r"(_mma_b_lo_10), "r"((tmem_tmem_scratch + (256))), "r"(((first_pv) ? 0 : 1))); - int _mma_a_lo_11 = (((smem_p_fp8_addr) >> 4) & 0x3FFF) + (last_phase) * 512; - int _mma_b_lo_11 = ((((smem_v_full_addr + 16384) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 138477584;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_11), "r"(_mma_b_lo_11), "r"((tmem_tmem_scratch + (384))), "r"(((first_pv) ? 0 : 1))); - if (cta_rank == 0) { - elect_commit_cg2_multicast(v_empty_addr + (mma_v_stage) * 8, (uint16_t)(3)); - } - } - mma_v_stage += 1; - if (mma_v_stage == 2) { mma_v_stage = 0; _phase_v_full ^= 1; } - } - if (cta_rank == 0) { - elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); - } - if (cta_rank == 0) { - elect_commit_cg2_multicast(o_done_addr, (uint16_t)(3)); - } - mma_tile_cursor = mma_tile_cursor + num_kv_tiles_3; - } - mbarrier_arrive(tmem_dealloc_addr); - unsigned int _phase_tmem_dealloc_0 = 0; - mbarrier_wait(tmem_dealloc_addr, _phase_tmem_dealloc_0); - _phase_tmem_dealloc_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - { - int tmem_dealloc_peer_rank = bid % 2 ^ 1; - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cluster.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(tmem_dealloc_peer_addr), "r"(tmem_dealloc_peer_rank) : "memory"); - mbarrier_wait(tmem_dealloc_peer_addr, 0); - int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); - asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); - } - } - } - // ---- Role: empty1 ---- - if (warp >= 10 && warp <= 11) { - { // empty1_main - const int wg2_dummy_2 = 0; - unsigned int _phase_q_full_0_4 = 0; - } - } - - // Cleanup -} - -} // extern "C" diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_fbf138c989dfde9f764f_binding.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_fbf138c989dfde9f764f_binding.cu deleted file mode 100644 index e832234c20a..00000000000 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_fbf138c989dfde9f764f_binding.cu +++ /dev/null @@ -1,446 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -// Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_fbf138c989dfde9f764f'. -#include -#include -#include -#include - -#include "tvm_ffi_utils.h" - -#include -#include -#include -#include -#include -#include -#include - -extern "C" __global__ void kernel_cake_dsv4_fbf138c989dfde9f764f(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int max_q_len, int batch_size, int has_sinks); - - -namespace cake_host_shim_e92a6da8d9e51169 { - -using tvm::ffi::Optional; -using tvm::ffi::TensorView; - -class ScopedCudaDevice { - public: - explicit ScopedCudaDevice(int device_id) { - cudaError_t error = cudaGetDevice(&previous_device_); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed before host-shim launch: cudaError=" - << static_cast(error); - if (previous_device_ != device_id) { - error = cudaSetDevice(device_id); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaSetDevice failed before host-shim launch for cuda:" - << device_id << ": cudaError=" << static_cast(error); - restore_ = true; - } - } - - ScopedCudaDevice(const ScopedCudaDevice&) = delete; - ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; - - ~ScopedCudaDevice() noexcept { - if (restore_) { - (void)cudaSetDevice(previous_device_); - } - } - - private: - int previous_device_ = -1; - bool restore_ = false; -}; - -inline int64_t CakeDeviceMultiprocessorCount(int device_id) { - constexpr int kMaxCachedCudaDevices = 64; - TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) - << "physical-SM-count cache does not cover cuda:" << device_id; - static std::atomic count_by_device[kMaxCachedCudaDevices]{}; - int cached = count_by_device[device_id].load(std::memory_order_acquire); - if (cached > 0) return cached; - - int count = 0; - cudaError_t error = cudaDeviceGetAttribute( - &count, cudaDevAttrMultiProcessorCount, device_id); - TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) - << "querying multiProcessorCount failed for cuda:" << device_id - << ": cudaError=" << static_cast(error) << ", count=" << count; - // Concurrent first launches may repeat the immutable device query, but all - // publication is atomic and every later hot-path lookup is one acquire load. - count_by_device[device_id].store(count, std::memory_order_release); - return count; -} - -inline void CheckCudaTensor(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) - << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; -} - -inline void CheckSameCudaDevice( - const TensorView& t, - const TensorView& reference, - const char* name, - const char* reference_name) { - TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) - << name << " must be on the same CUDA device as " << reference_name - << ": got cuda:" << t.device().device_id - << " versus cuda:" << reference.device().device_id; -} - -inline void CheckCurrentCudaDevice( - const TensorView& reference, - const char* reference_name) { - int current_device = -1; - cudaError_t error = cudaGetDevice(¤t_device); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed while validating " << reference_name - << ": cudaError=" << static_cast(error); - TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) - << "current CUDA device must match " << reference_name - << ": current=cuda:" << current_device - << ", tensor=cuda:" << reference.device().device_id; -} - -inline void CheckContiguous(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; -} - -inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { - DLDataType d = t.dtype(); - TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) - << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits - << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits - << ", lanes=" << (int)d.lanes << ")"; -} - -// A logical axis.outer(trailing) folds every source dim above the trailing -// dimensions. Shape products are independent of physical strides, so verify -// the leading dimensions form one dense row-major chain instead of inventing -// a "folded stride". The descriptor reads its exact adjacent physical step -// separately through stride[-(trailing + 1)]. -inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { - TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) - << name << " cannot fold leading dimensions above " << trailing - << " trailing dims from ndim=" << t.ndim(); - int outer_last = t.ndim() - trailing - 1; - if (outer_last <= 0) { - return; - } - int64_t step = t.stride(outer_last); - TVM_FFI_CHECK(step > 0, ValueError) - << name << " physical strides must be positive"; - int64_t expected = step; - for (int axis = outer_last - 1; axis >= 0; --axis) { - expected *= t.size(axis + 1); - if (t.size(axis) > 1) { - TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) - << name << " leading dims are not physically foldable above " << trailing - << " trailing dims: stride(" << axis << ")=" << t.stride(axis) - << ", expected " << expected; - } - } -} - -// 4D TMA descriptor for buffer 'tmap_q' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 3, ValueError) - << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - int64_t d2 = t.size(t.ndim() - 2); - int64_t d3 = t.size(t.ndim() - 3); - TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) - << "TMA source 'tmap_q' trailing dims must be positive"; - TVM_FFI_CHECK(d1 % 128 == 0, ValueError) - << "TMA source 'tmap_q' extent " << d1 - << " must divide exactly by " << 128; - uint64_t global_dim[4] = {(uint64_t)(128), (uint64_t)(d2), (uint64_t)((d1 / 128)), (uint64_t)(d3)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; - TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) - << "TMA box (128, 8, 1, 1) exceeds resolved global dims for 'tmap_q'"; - int64_t carrier_stride_0 = d1; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; - int64_t carrier_stride_1 = 128; - TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; - TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; - TVM_FFI_CHECK((carrier_stride_1 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; - int64_t carrier_stride_2 = (d2 * d1); - TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; - TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; - TVM_FFI_CHECK((carrier_stride_2 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; - uint64_t global_strides[3] = { - (uint64_t)((carrier_stride_0 * 8) / 8), - (uint64_t)((carrier_stride_1 * 8) / 8), - (uint64_t)((carrier_stride_2 * 8) / 8), - }; - uint32_t box_dim[4] = {128u, 8u, 1u, 1u}; - uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_swa_kv' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_swa_kv' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; - TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (128, 1) exceeds resolved global dims for 'tmap_swa_kv'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 8) / 8), - }; - uint32_t box_dim[2] = {128u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_compressed_kv' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_compressed_kv(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_compressed_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_compressed_kv' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_compressed_kv' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_compressed_kv"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_compressed_kv' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved a non-positive global dim"; - TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (128, 1) exceeds resolved global dims for 'tmap_compressed_kv'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 8) / 8), - }; - uint32_t box_dim[2] = {128u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_kv') failed: CUresult=" << (int)r; - return tm; -} - -void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_compressed_kv, TensorView arg_O, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_seq_lens, TensorView arg_cum_seq_lens_q, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_max_q_len, int64_t arg_batch_size, int64_t arg_has_sinks, int64_t grid_x, int64_t grid_y, int64_t grid_z) { - DLDevice dev = arg_tmap_q.device(); - ScopedCudaDevice device_guard(dev.device_id); - CheckCudaTensor(arg_tmap_q, "tmap_q"); - CheckDtype(arg_tmap_q, "tmap_q", 1, 8, 1); - CheckContiguous(arg_tmap_q, "tmap_q"); - CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); - CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 1, 8, 1); - CheckCudaTensor(arg_tmap_compressed_kv, "tmap_compressed_kv"); - CheckDtype(arg_tmap_compressed_kv, "tmap_compressed_kv", 1, 8, 1); - CheckCudaTensor(arg_O, "O"); - CheckDtype(arg_O, "O", 4, 16, 1); - CheckContiguous(arg_O, "O"); - CheckCudaTensor(arg_swa_indices, "swa_indices"); - CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); - CheckContiguous(arg_swa_indices, "swa_indices"); - CheckCudaTensor(arg_compressed_indices, "compressed_indices"); - CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); - CheckContiguous(arg_compressed_indices, "compressed_indices"); - CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); - CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckCudaTensor(arg_seq_lens, "seq_lens"); - CheckDtype(arg_seq_lens, "seq_lens", 0, 32, 1); - CheckContiguous(arg_seq_lens, "seq_lens"); - CheckCudaTensor(arg_cum_seq_lens_q, "cum_seq_lens_q"); - CheckDtype(arg_cum_seq_lens_q, "cum_seq_lens_q", 0, 32, 1); - CheckContiguous(arg_cum_seq_lens_q, "cum_seq_lens_q"); - CheckCudaTensor(arg_sinks, "sinks"); - CheckDtype(arg_sinks, "sinks", 2, 32, 1); - CheckContiguous(arg_sinks, "sinks"); - CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); - CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); - CheckContiguous(arg_bmm1_scale, "bmm1_scale"); - CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); - CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); - CheckContiguous(arg_bmm2_scale, "bmm2_scale"); - TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) - << "scalar 'swa_index_stride' value " << arg_swa_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) - << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) - << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) - << "scalar 'sparse_topk' value " << arg_sparse_topk - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_max_q_len >= -2147483648LL && arg_max_q_len <= 2147483647LL, ValueError) - << "scalar 'max_q_len' value " << arg_max_q_len - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_batch_size >= -2147483648LL && arg_batch_size <= 2147483647LL, ValueError) - << "scalar 'batch_size' value " << arg_batch_size - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) - << "scalar 'has_sinks' value " << arg_has_sinks - << " is outside i32 range [-2147483648, 2147483647]"; - CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); - CheckSameCudaDevice(arg_tmap_compressed_kv, arg_tmap_q, "tmap_compressed_kv", "tmap_q"); - CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); - CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); - CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); - CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); - CheckSameCudaDevice(arg_seq_lens, arg_tmap_q, "seq_lens", "tmap_q"); - CheckSameCudaDevice(arg_cum_seq_lens_q, arg_tmap_q, "cum_seq_lens_q", "tmap_q"); - CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); - CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); - CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); - TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) - << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y - << ", " << grid_z << ")"; - - cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); - CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); - CUtensorMap p_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); - CUtensorMap p_tmap_compressed_kv = EncodeTma_tmap_compressed_kv(arg_tmap_compressed_kv); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); - int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); - int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); - int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); - int* p_seq_lens = static_cast(arg_seq_lens.data_ptr()); - int* p_cum_seq_lens_q = static_cast(arg_cum_seq_lens_q.data_ptr()); - float* p_sinks = static_cast(arg_sinks.data_ptr()); - float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); - float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); - int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; - int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; - int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; - int32_t v_sparse_topk = (int32_t)arg_sparse_topk; - int32_t v_max_q_len = (int32_t)arg_max_q_len; - int32_t v_batch_size = (int32_t)arg_batch_size; - int32_t v_has_sinks = (int32_t)arg_has_sinks; - void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_tmap_compressed_kv, &p_O, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_seq_lens, &p_cum_seq_lens_q, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_max_q_len, &v_batch_size, &v_has_sinks}; - - static std::mutex smem_mu; - static auto* smem_status_by_device = - new std::unordered_map(); - cudaError_t smem_status = cudaSuccess; - { - std::lock_guard lock(smem_mu); - auto it = smem_status_by_device->find(dev.device_id); - if (it == smem_status_by_device->end()) { - int current_device = -1; - smem_status = cudaGetDevice(¤t_device); - if (smem_status == cudaSuccess && current_device != dev.device_id) { - smem_status = cudaErrorInvalidDevice; - } - if (smem_status == cudaSuccess) { - smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_fbf138c989dfde9f764f), - cudaFuncAttributeMaxDynamicSharedMemorySize, - 147072); - } - smem_status_by_device->emplace(dev.device_id, smem_status); - } else { - smem_status = it->second; - } - } - TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_fbf138c989dfde9f764f failed: " - << cudaGetErrorString(smem_status); - dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); - dim3 block(640u, 1u, 1u); - - cudaLaunchConfig_t config{}; - config.gridDim = grid; - config.blockDim = block; - config.dynamicSmemBytes = 147072u; - config.stream = stream; - cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_fbf138c989dfde9f764f), kargs); - TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_fbf138c989dfde9f764f failed: " - << cudaGetErrorString(launch_status); -} - -} // namespace cake_host_shim_e92a6da8d9e51169 - -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_e92a6da8d9e51169::Run); diff --git a/csrc/cake_dsv4/sm_100a/cake_dsv4_fbf138c989dfde9f764f_kernel.cu b/csrc/cake_dsv4/sm_100a/cake_dsv4_fbf138c989dfde9f764f_kernel.cu deleted file mode 100644 index f09dcbb858b..00000000000 --- a/csrc/cake_dsv4/sm_100a/cake_dsv4_fbf138c989dfde9f764f_kernel.cu +++ /dev/null @@ -1,2616 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -typedef signed char int8_t; -typedef unsigned char uint8_t; -typedef unsigned short uint16_t; -typedef unsigned int uint32_t; -#if defined(__CUDACC_RTC__) -typedef unsigned long long uint64_t; -#else -typedef unsigned long uint64_t; -#endif -static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); -typedef signed int int32_t; -typedef short int int16_t; -struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; -struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; -static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); -static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); - -#if defined(__CUDACC_RTC__) -typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; -#else -#include -#endif - -static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); -static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); -#include -#include - -#define CAKE_INF CUDART_INF_F -#define TMEM_NCOLS 512 -#define TMEM_SOURCE_EXACT_SCRATCH_OFFSET 0 -#define NUM_Q_PIPE_STAGES 1 -#define NUM_KV_PIPE_STAGES 8 -#define NUM_PAGE_PIPE_STAGES 6 -#define NUM_SCORE_PIPE_STAGES 2 -#define NUM_STATS_PIPE_STAGES 2 -#define NUM_P_PIPE_STAGES 2 -#define NUM_O_PIPE_STAGES 1 -#define NUM_WORK_PIPE_STAGES 2 -#define NUM_THROTTLE_PIPE_STAGES 2 -#define SMEM_SMEM_Q_OFF 1024 -#define SMEM_SMEM_Q_STAGE_BYTES 1024 -#define SMEM_SMEM_Q_STRIDE 1024 -#define SMEM_SMEM_KV_OFF 5120 -#define SMEM_SMEM_KV_STAGE_BYTES 16384 -#define SMEM_SMEM_KV_STRIDE 16384 -#define SMEM_SMEM_V_OFF 5120 -#define SMEM_SMEM_V_STAGE_BYTES 16384 -#define SMEM_SMEM_V_STRIDE 16384 -#define SMEM_SMEM_P_OFF 136192 -#define SMEM_SMEM_P_STAGE_BYTES 1024 -#define SMEM_SMEM_P_STRIDE 1024 -#define SMEM_SMEM_PAGE_OFFSETS_OFF 138240 -#define SMEM_SMEM_PAGE_OFFSETS_STAGE_BYTES 1024 -#define SMEM_SMEM_PAGE_OFFSETS_STRIDE 1024 -#define SMEM_SMEM_O_OFF 144384 -#define SMEM_SMEM_O_STAGE_BYTES 2048 -#define SMEM_SMEM_O_STRIDE 2048 -#define SMEM_SMEM_SOFTMAX_OFF 146432 -#define SMEM_SMEM_SOFTMAX_STAGE_BYTES 128 -#define SMEM_SMEM_SOFTMAX_STRIDE 128 -#define SMEM_SMEM_SOFTMAX_U32_OFF 146432 -#define SMEM_SMEM_SOFTMAX_U32_STAGE_BYTES 128 -#define SMEM_SMEM_SOFTMAX_U32_STRIDE 128 -#define SMEM_SMEM_SOFTMAX_UNUSED_OFF 146560 -#define SMEM_SMEM_SOFTMAX_UNUSED_STAGE_BYTES 128 -#define SMEM_SMEM_SOFTMAX_UNUSED_STRIDE 128 -#define SMEM_SMEM_CORR_OFF 146688 -#define SMEM_SMEM_CORR_STAGE_BYTES 128 -#define SMEM_SMEM_CORR_STRIDE 128 -#define SMEM_SMEM_STATS_OFF 146816 -#define SMEM_SMEM_STATS_STAGE_BYTES 96 -#define SMEM_SMEM_STATS_STRIDE 96 -#define SMEM_WORK_RESPONSE_OFF 147008 -#define SMEM_WORK_RESPONSE_STAGE_BYTES 16 -#define SMEM_WORK_RESPONSE_STRIDE 16 -#define SMEM_TOTAL 147072 -#define THREADS 640 - -#include - -__device__ __forceinline__ uint32_t elect_sync() { - uint32_t pred = 0; - asm volatile( - "{\n\t" - ".reg .pred %%px;\n\t" - "elect.sync _|%%px, %1;\n\t" - "@%%px mov.s32 %0, 1;\n\t" - "}\n" - : "+r"(pred) - : "r"(0xFFFFFFFF)); - return pred; -} - - -__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { - asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" - :: "r"(mbar_addr), "r"(count) : "memory"); -} - -__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { - asm volatile("mbarrier.init.b64 [%0], %1;" - :: "l"(mbar_addr), "r"(count) : "memory"); -} - - -__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - - -// CTA-local pipelines have short, resident producer/consumer edges. Omitting -// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the -// explicit loop still makes this helper blocking until acquire succeeds. -__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE;\n\t" - "bra.uni LAB_WAIT;\n\t" - "DONE:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Source-faithful relaxed CTA wait used only by a typed protocol that does -// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. -__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, 10000000;\n\t" - "@P1 bra.uni DONE_RELAXED;\n\t" - "bra.uni LAB_WAIT_RELAXED;\n\t" - "DONE_RELAXED:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Exact source ports may request the PTX suspendTimeHint operand explicitly. -// The hint is expressed in nanoseconds and is kept separate from the canonical -// no-hint CTA helper so unrelated schedules retain their existing retry path. -__device__ __forceinline__ void mbarrier_wait_suspend( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_SUSPEND:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_SUSPEND;\n\t" - "bra.uni LAB_WAIT_SUSPEND;\n\t" - "DONE_SUSPEND:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE_CLUSTER;\n\t" - "bra.uni LAB_WAIT_CLUSTER;\n\t" - "DONE_CLUSTER:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - ".reg .u32 WAIT_ADDR;\n\t" - "mov.u32 WAIT_ADDR, %0;\n\t" - "LAB_WAIT_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [WAIT_ADDR], %1, %2;\n\t" - "@P1 bra.uni DONE_HINT;\n\t" - "bra.uni LAB_WAIT_HINT;\n\t" - "DONE_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -// Exact unqualified CTA wait used by source schedules whose PTX intentionally -// omits the acquire qualifier while retaining a typed suspendTimeHint operand. -__device__ __forceinline__ void mbarrier_wait_relaxed_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED_HINT:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra DONE_RELAXED_HINT;\n\t" - "bra LAB_WAIT_RELAXED_HINT;\n\t" - "DONE_RELAXED_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); -} - -__device__ __forceinline__ void mbarrier_wait_cluster_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" - "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" - "DONE_CLUSTER_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_suspend( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait_cluster(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); - } -} - - -__device__ __forceinline__ void tcgen05_mma_f8f6f4( - int taddr, uint64_t a_desc, uint64_t b_desc, - uint32_t i_desc, int enable_input_d) { - asm volatile( - "{\n\t" - ".reg .pred p;\n\t" - "setp.ne.b32 p, %4, 0;\n\t" - "tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], %1, %2, %3, p;\n\t" - "}\n" - :: "r"(taddr), "l"(a_desc), "l"(b_desc), - "r"(i_desc), "r"(enable_input_d) - : "memory"); -} - - -__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { - return (x & 0x3FFFFULL) >> 4ULL; -} - - -union MmaSmemDesc { - uint64_t u64; - uint32_t u32[2]; -}; - -__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { - MmaSmemDesc tmp; - tmp.u64 = smem_desc; - tmp.u32[0] += offset; - smem_desc = tmp.u64; -} - - -__device__ __forceinline__ void elect_commit(int mbar_addr) { - asm volatile( - "{\n\t" - ".reg .pred leader;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" - ".shared::cluster.b64 [%0];\n\t" - "}\n" - :: "r"(mbar_addr)); -} - - -__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { - asm volatile( - "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" - :: "r"(mbar_addr), "r"(bytes) : "memory"); -} - - -__device__ __forceinline__ void tmem_ld_x4(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x4.b32" - " {%0, %1, %2, %3}, [%4];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { - asm volatile( - "tcgen05.st.sync.aligned.32x32b.x4.b32" - " [%0], {%1, %2, %3, %4};" - :: "r"(tmem_addr), - "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); -} - - -__device__ __forceinline__ float approx_exp2(float x) { - float y; - asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float approx_rcp(float x) { - float y; - asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float max_noftz(float a, float b) { - float c; - asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); - return c; -} - - -__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { - asm("mul.rn.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { - asm("mul.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { - asm("add.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { - asm("add.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { - asm("sub.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { - asm("sub.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("add.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { - float2 r; - asm("sub.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("sub.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ void fma_scale_x32( - float* sv, const float2* scale2, const float2* neg_max2) -{ - float2* sv_2 = reinterpret_cast(sv); - #pragma unroll - for (int j = 0; j < 16; j++) - fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); -} - -__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("mul.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - - -// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) - -__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("add.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("add.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("add.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("add.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("add.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("add.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("add.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - - -__device__ __forceinline__ void fence_async_shared() { - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); -} - - -__device__ __forceinline__ uint64_t make_smem_desc(int addr) { - const int SBO = 1024; - return desc_encode(addr) - | (desc_encode(SBO) << 32ULL) - | (1ULL << 46ULL) - | (2ULL << 61ULL); -} - - -__device__ __forceinline__ void tma_2d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3}], [%4];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_4d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.4d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5}], [%6];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr) { - // Canonical .shared::cta form for non-multicast gather4, matching - // trtllm-gen / cuda_ptx and the PTX ISA qualifier order - // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, - // .shared::cluster is reserved for the multicast variant (ctaMask). - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { - asm volatile( - "tcgen05.commit.cta_group::1.mbarrier::arrive::one" - ".shared::cluster.b64 [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tmem_ld_x4_wait(float* dst, int addr) { - tmem_ld_x4(dst, addr); - asm volatile("tcgen05.wait::ld.sync.aligned;"); -} - - -__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { - uint32_t result; - asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" - : "=r"(result) : "r"(val)); - return result; -} - -extern "C" { - -__global__ __launch_bounds__(640, 1) void -kernel_cake_dsv4_fbf138c989dfde9f764f(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int max_q_len, int batch_size, int has_sinks) -{ - const int tid = threadIdx.x; - const int warp = make_warp_uniform(tid / 32); - const int lane = tid % 32; - - extern __shared__ __align__(1024) char smem_raw[]; - int smem; - asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); - smem = make_warp_uniform(smem); - - const int mbar_base = smem; - #define q_full_addr (mbar_base + 0) - #define q_empty_addr (mbar_base + 8) - #define kv_full_addr (mbar_base + 16) - #define kv_empty_addr (mbar_base + 80) - #define page_full_addr (mbar_base + 144) - #define page_empty_addr (mbar_base + 192) - #define score_full_addr (mbar_base + 240) - #define score_empty_addr (mbar_base + 256) - #define stats_full_addr (mbar_base + 272) - #define stats_empty_addr (mbar_base + 288) - #define p_full_addr (mbar_base + 304) - #define p_empty_addr (mbar_base + 320) - #define o_full_addr (mbar_base + 336) - #define o_empty_addr (mbar_base + 344) - #define work_full_addr (mbar_base + 352) - #define work_empty_addr (mbar_base + 368) - #define throttle_full_addr (mbar_base + 384) - #define throttle_empty_addr (mbar_base + 400) - - const int bid = blockIdx.x; - const int num_bids = gridDim.x; - - const int cta_rank = 0; - - // Kernel setup ops - uint8_t* smem_q = reinterpret_cast(smem_raw + 1024); - const int smem_q_addr = smem + 1024; - uint8_t* smem_kv = reinterpret_cast(smem_raw + 5120); - const int smem_kv_addr = smem + 5120; - uint8_t* smem_v = reinterpret_cast(smem_raw + 5120); - const int smem_v_addr = smem + 5120; - uint8_t* smem_p = reinterpret_cast(smem_raw + 136192); - const int smem_p_addr = smem + 136192; - int* smem_page_offsets = reinterpret_cast(smem_raw + 138240); - const int smem_page_offsets_addr = smem + 138240; - __nv_bfloat16* smem_o = reinterpret_cast<__nv_bfloat16*>(smem_raw + 144384); - const int smem_o_addr = smem + 144384; - float* smem_softmax = reinterpret_cast(smem_raw + 146432); - const int smem_softmax_addr = smem + 146432; - unsigned int* smem_softmax_u32 = reinterpret_cast(smem_raw + 146432); - const int smem_softmax_u32_addr = smem + 146432; - float* smem_softmax_unused = reinterpret_cast(smem_raw + 146560); - const int smem_softmax_unused_addr = smem + 146560; - float* smem_corr = reinterpret_cast(smem_raw + 146688); - const int smem_corr_addr = smem + 146688; - float* smem_stats = reinterpret_cast(smem_raw + 146816); - const int smem_stats_addr = smem + 146816; - unsigned int* work_response = reinterpret_cast(smem_raw + 147008); - const int work_response_addr = smem + 147008; - asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_q))) : "memory"); - asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_swa_kv))) : "memory"); - asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_compressed_kv))) : "memory"); - - // Mbarrier init (18 pipeline groups, 0 ordered-sequence groups, 52 barriers) - // Mbarriers at smem_raw[0..416) - - if (warp == 0) { - uint32_t leader = elect_sync(); - if (leader) { - // --- pipeline 'q_pipe' --- - // q_full: 1 barriers, init_count=1 - mbarrier_init(smem + 0, 1); - // q_empty: 1 barriers, init_count=1 - mbarrier_init(smem + 8, 1); - // --- pipeline 'kv_pipe' --- - // kv_full: 8 barriers, init_count=1 - mbarrier_init(smem + 16, 1); - mbarrier_init(smem + 24, 1); - mbarrier_init(smem + 32, 1); - mbarrier_init(smem + 40, 1); - mbarrier_init(smem + 48, 1); - mbarrier_init(smem + 56, 1); - mbarrier_init(smem + 64, 1); - mbarrier_init(smem + 72, 1); - // kv_empty: 8 barriers, init_count=1 - mbarrier_init(smem + 80, 1); - mbarrier_init(smem + 88, 1); - mbarrier_init(smem + 96, 1); - mbarrier_init(smem + 104, 1); - mbarrier_init(smem + 112, 1); - mbarrier_init(smem + 120, 1); - mbarrier_init(smem + 128, 1); - mbarrier_init(smem + 136, 1); - // --- pipeline 'page_pipe' --- - // page_full: 6 barriers, init_count=32 - mbarrier_init(smem + 144, 32); - mbarrier_init(smem + 152, 32); - mbarrier_init(smem + 160, 32); - mbarrier_init(smem + 168, 32); - mbarrier_init(smem + 176, 32); - mbarrier_init(smem + 184, 32); - // page_empty: 6 barriers, init_count=256 - mbarrier_init(smem + 192, 256); - mbarrier_init(smem + 200, 256); - mbarrier_init(smem + 208, 256); - mbarrier_init(smem + 216, 256); - mbarrier_init(smem + 224, 256); - mbarrier_init(smem + 232, 256); - // --- pipeline 'score_pipe' --- - // score_full: 2 barriers, init_count=1 - mbarrier_init(smem + 240, 1); - mbarrier_init(smem + 248, 1); - // score_empty: 2 barriers, init_count=128 - mbarrier_init(smem + 256, 128); - mbarrier_init(smem + 264, 128); - // --- pipeline 'stats_pipe' --- - // stats_full: 2 barriers, init_count=128 - mbarrier_init(smem + 272, 128); - mbarrier_init(smem + 280, 128); - // stats_empty: 2 barriers, init_count=128 - mbarrier_init(smem + 288, 128); - mbarrier_init(smem + 296, 128); - // --- pipeline 'p_pipe' --- - // p_full: 2 barriers, init_count=128 - mbarrier_init(smem + 304, 128); - mbarrier_init(smem + 312, 128); - // p_empty: 2 barriers, init_count=1 - mbarrier_init(smem + 320, 1); - mbarrier_init(smem + 328, 1); - // --- pipeline 'o_pipe' --- - // o_full: 1 barriers, init_count=1 - mbarrier_init(smem + 336, 1); - // o_empty: 1 barriers, init_count=128 - mbarrier_init(smem + 344, 128); - // --- pipeline 'work_pipe' --- - // work_full: 2 barriers, init_count=1 - mbarrier_init(smem + 352, 1); - mbarrier_init(smem + 360, 1); - // work_empty: 2 barriers, init_count=640 - mbarrier_init(smem + 368, 640); - mbarrier_init(smem + 376, 640); - // --- pipeline 'throttle_pipe' --- - // throttle_full: 2 barriers, init_count=256 - mbarrier_init(smem + 384, 256); - mbarrier_init(smem + 392, 256); - // throttle_empty: 2 barriers, init_count=32 - mbarrier_init(smem + 400, 32); - mbarrier_init(smem + 408, 32); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - - __syncwarp(); - - // TMEM alloc (512 columns, 512 used) - volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 416); - if (warp == 0) { - int _tmem_hold = smem + 416; - asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); - __syncwarp(); - asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); - } - - __syncthreads(); - asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); - - const int taddr = tmem_addr_storage[0]; - - // Kernel post-init ops - const int tmem_source_exact_scratch = taddr; - - // ---- Ordered hardware-WG register redistribution ---- - // Dec phase frees registers before any WG attempts inc. - if (warp >= 8 && warp <= 11) { - asm volatile("setmaxnreg.dec.sync.aligned.u32 48;"); - } - - // ---- Role: softmax ---- - if (warp <= 3) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 152;"); - { // softmax_main - unsigned int work_stage = 0; - unsigned int score_stage = 0; - unsigned int score_phase = 0; - unsigned int stats_stage = 0; - unsigned int stats_phase = 1; - unsigned int p_stage = 0; - unsigned int p_phase = 1; - unsigned int work_x = blockIdx.x; - unsigned int work_y = blockIdx.y; - unsigned int work_z = blockIdx.z; - const int warp_in_wg = warp; - const int wg_tid = warp_in_wg * 32 + lane; - const int col_pair_base = wg_tid % 4 * 2; - const int ldtm_row_base = warp_in_wg * 32 + lane / 4; - float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; - unsigned int _phase_work_full = 0; - #pragma unroll 1 - for (unsigned int _work = 0; _work < 120; _work++) { - int q_begin = cum_seq_lens_q[work_z]; - int q_end = cum_seq_lens_q[work_z + 1]; - int q_len = q_end - q_begin; - int query_idx = (unsigned int)q_begin + work_x; - int active_topk = 0; - if (work_x < (unsigned int)q_len) { - int _max_4 = ((sparse_topk_lens[query_idx] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx] + sparse_topk_lens_offset) : (0)); - int _min_3 = ((_max_4) < (sparse_topk) ? (_max_4) : (sparse_topk)); - active_topk = _min_3; - } - int num_steps = (active_topk + 128 - 1) / 128; - float row_max_pair[2]; - float row_sum_pair[2]; - row_max_pair[0] = -CAKE_INF; - row_max_pair[1] = -CAKE_INF; - row_sum_pair[0] = 0.0f; - row_sum_pair[1] = 0.0f; - asm volatile("barrier.sync 1, 128;" ::: "memory"); - uint32_t _amf_u_0 = __float_as_uint(-3.4028235e+38f); - uint32_t _amf_mask_0 = -int32_t(_amf_u_0 >> 31) | 0x80000000u; - unsigned int _amf_enc_0 = _amf_u_0 ^ _amf_mask_0; - if (wg_tid < 8) { - smem_softmax_u32[wg_tid] = _amf_enc_0; - } - asm volatile("barrier.sync 1, 128;" ::: "memory"); - #pragma unroll 1 - for (int tile = 0; tile < num_steps; tile++) { - mbarrier_wait(score_full_addr + (score_stage) * 8, score_phase); - int score_col = ((score_stage == 0) ? 0 : 8); - float scores[8]; - float score_lo[4]; - float score_hi[4]; - int tile_end = (tile + 1) * 128; - int is_full_tile = ((active_topk % 128 == 0) ? 1 : 0); - is_full_tile = is_full_tile | ((tile_end < active_topk) ? 1 : 0); - if (is_full_tile != 0) { - asm volatile( - "tcgen05.ld.sync.aligned.16x256b.x1.b32" - " {%0, %1, %2, %3}, [%4];" - : "=r"(*reinterpret_cast(&score_lo[0])), "=r"(*reinterpret_cast(&score_lo[1])), "=r"(*reinterpret_cast(&score_lo[2])), "=r"(*reinterpret_cast(&score_lo[3])) - : "r"(taddr + (unsigned int)score_col)); - asm volatile( - "tcgen05.ld.sync.aligned.16x256b.x1.b32" - " {%0, %1, %2, %3}, [%4];" - : "=r"(*reinterpret_cast(&score_hi[0])), "=r"(*reinterpret_cast(&score_hi[1])), "=r"(*reinterpret_cast(&score_hi[2])), "=r"(*reinterpret_cast(&score_hi[3])) - : "r"(taddr + (unsigned int)score_col + 1048576)); - #pragma unroll - for (int c = 0; c < 4; c++) { - scores[c] = score_lo[c]; - scores[c + 4] = score_hi[c]; - } - } else { - asm volatile("" ::: "memory"); - asm volatile( - "tcgen05.ld.sync.aligned.16x256b.x1.b32" - " {%0, %1, %2, %3}, [%4];" - : "=r"(*reinterpret_cast(&score_lo[0])), "=r"(*reinterpret_cast(&score_lo[1])), "=r"(*reinterpret_cast(&score_lo[2])), "=r"(*reinterpret_cast(&score_lo[3])) - : "r"(taddr + (unsigned int)score_col)); - asm volatile( - "tcgen05.ld.sync.aligned.16x256b.x1.b32" - " {%0, %1, %2, %3}, [%4];" - : "=r"(*reinterpret_cast(&score_hi[0])), "=r"(*reinterpret_cast(&score_hi[1])), "=r"(*reinterpret_cast(&score_hi[2])), "=r"(*reinterpret_cast(&score_hi[3])) - : "r"(taddr + (unsigned int)score_col + 1048576)); - #pragma unroll - for (int c_1 = 0; c_1 < 4; c_1++) { - scores[c_1] = score_lo[c_1]; - scores[c_1 + 4] = score_hi[c_1]; - } - if (active_topk <= tile * 128 + ldtm_row_base) { - scores[0] = -CAKE_INF; - scores[1] = -CAKE_INF; - } - if (active_topk <= tile * 128 + ldtm_row_base + 8) { - scores[2] = -CAKE_INF; - scores[3] = -CAKE_INF; - } - if (active_topk <= tile * 128 + ldtm_row_base + 16) { - scores[4] = -CAKE_INF; - scores[5] = -CAKE_INF; - } - if (active_topk <= tile * 128 + ldtm_row_base + 24) { - scores[6] = -CAKE_INF; - scores[7] = -CAKE_INF; - } - asm volatile("" ::: "memory"); - } - float pair_max[2]; - float _max_5 = max_noftz(scores[0], scores[2]); - float _max_6 = max_noftz(scores[4], scores[6]); - float _max_7 = max_noftz(_max_5, _max_6); - pair_max[0] = _max_7; - float _max_8 = max_noftz(scores[1], scores[3]); - float _max_9 = max_noftz(scores[5], scores[7]); - float _max_10 = max_noftz(_max_8, _max_9); - pair_max[1] = _max_10; - #pragma unroll - for (int c_2 = 0; c_2 < 2; c_2++) { - float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, pair_max[c_2], 16); - float _max_11 = max_noftz(pair_max[c_2], _shfl_xor_0); - pair_max[c_2] = _max_11; - float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, pair_max[c_2], 8); - float _max_12 = max_noftz(pair_max[c_2], _shfl_xor_1); - pair_max[c_2] = _max_12; - } - float old_max_pair[2]; - float new_max_pair[2]; - #pragma unroll - for (int c_3 = 0; c_3 < 2; c_3++) { - old_max_pair[c_3] = row_max_pair[c_3]; - float _max_13 = max_noftz(row_max_pair[c_3], pair_max[c_3]); - new_max_pair[c_3] = _max_13; - } - if (lane < 8) { - uint32_t _amf_u_1 = __float_as_uint(new_max_pair[0]); - uint32_t _amf_mask_1 = -int32_t(_amf_u_1 >> 31) | 0x80000000u; - unsigned int _amf_enc_1 = _amf_u_1 ^ _amf_mask_1; - uint32_t _amf_u_2 = __float_as_uint(new_max_pair[1]); - uint32_t _amf_mask_2 = -int32_t(_amf_u_2 >> 31) | 0x80000000u; - unsigned int _amf_enc_2 = _amf_u_2 ^ _amf_mask_2; - atomicMax(&smem_softmax_u32[col_pair_base], _amf_enc_1); - atomicMax(&smem_softmax_u32[col_pair_base + 1], _amf_enc_2); - } - asm volatile("barrier.sync 1, 128;" ::: "memory"); - uint32_t _amf_u_3 = smem_softmax_u32[col_pair_base]; - uint32_t _amf_mask_3 = ((_amf_u_3 >> 31) - 1u) | 0x80000000u; - float _amf_dec_0 = __uint_as_float(_amf_u_3 ^ _amf_mask_3); - new_max_pair[0] = _amf_dec_0; - uint32_t _amf_u_4 = smem_softmax_u32[col_pair_base + 1]; - uint32_t _amf_mask_4 = ((_amf_u_4 >> 31) - 1u) | 0x80000000u; - float _amf_dec_1 = __uint_as_float(_amf_u_4 ^ _amf_mask_4); - new_max_pair[1] = _amf_dec_1; - mbarrier_wait(stats_empty_addr + (stats_stage) * 8, stats_phase); - float stats_pair[4]; - stats_pair[0] = old_max_pair[0]; - stats_pair[1] = old_max_pair[1]; - stats_pair[2] = new_max_pair[0]; - stats_pair[3] = new_max_pair[1]; - int stats_col = ((stats_stage == 0) ? 16 : 48); - tmem_st_x4_f32(taddr + (unsigned int)stats_col, stats_pair); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(stats_full_addr + (stats_stage) * 8); - stats_stage += 1; - if (stats_stage == 2) { stats_stage = 0; stats_phase ^= 1; } - float acc_scale_pair[2]; - #pragma unroll - for (int c_4 = 0; c_4 < 2; c_4++) { - float delta = softmax_scale_log2 * (old_max_pair[c_4] - new_max_pair[c_4]); - float _exp2_0 = approx_exp2(delta); - acc_scale_pair[c_4] = ((old_max_pair[c_4] > -CAKE_INF) ? _exp2_0 : 1.0f); - row_max_pair[c_4] = new_max_pair[c_4]; - } - asm volatile("" ::: "memory"); - float exp_values[8]; - #pragma unroll - for (int c_5 = 0; c_5 < 8; c_5++) { - float safe_max = ((new_max_pair[c_5 % 2] == -CAKE_INF) ? 0.0f : new_max_pair[c_5 % 2]); - float _exp2_1 = approx_exp2(scores[c_5] * softmax_scale_log2 - safe_max * softmax_scale_log2); - exp_values[c_5] = _exp2_1; - } - mbarrier_wait(p_empty_addr + (p_stage) * 8, p_phase); - float p_values[8]; - #pragma unroll - for (int c_6 = 0; c_6 < 8; c_6++) { - p_values[c_6] = exp_values[c_6] * 448.0f; - } - unsigned int regs_p[2]; - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(p_values[0]), "f"(p_values[1]), - "f"(p_values[2]), "f"(p_values[3])); - regs_p[0] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(p_values[4]), "f"(p_values[5]), - "f"(p_values[6]), "f"(p_values[7])); - regs_p[1] = _packed; - } - int mtx_idx = lane / 8; - int thr_row_idx = lane % 8; - int seg_col_idx = warp_in_wg * 2 + mtx_idx ^ thr_row_idx; - int stsm_offset = thr_row_idx * 128 + seg_col_idx * 16; - const void* _stmatrix_b8_ptr_5 = reinterpret_cast(reinterpret_cast(smem_p) + (p_stage * 1024 + (unsigned int)stsm_offset)); - uint64_t _stmatrix_b8_addr64_5; - asm volatile("cvta.to.shared.u64 %0, %1;" : "=l"(_stmatrix_b8_addr64_5) : "l"(_stmatrix_b8_ptr_5)); - uint32_t _stmatrix_b8_addr_5; - asm volatile("cvt.u32.u64 %0, %1;" : "=r"(_stmatrix_b8_addr_5) : "l"(_stmatrix_b8_addr64_5)); - asm volatile("stmatrix.sync.aligned.m16n8.x2.trans.shared.b8 [%0], {%1, %2};\n" - :: "r"(_stmatrix_b8_addr_5), "r"(regs_p[0]), "r"(regs_p[1]) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(p_full_addr + (p_stage) * 8); - p_stage += 1; - if (p_stage == 2) { p_stage = 0; p_phase ^= 1; } - float2 _f2_0 = make_float2(exp_values[0], exp_values[1]); - float2 _f2_1 = make_float2(exp_values[2], exp_values[3]); - float2 _f2_2 = make_float2(exp_values[4], exp_values[5]); - float2 _f2_3 = make_float2(exp_values[6], exp_values[7]); - float2 _f2_4 = make_float2(row_sum_pair[0], row_sum_pair[1]); - float2 _f2_5 = make_float2(acc_scale_pair[0], acc_scale_pair[1]); - float2 sum_p0_pair = fma_f32x2_rn_ftz(_f2_4, _f2_5, _f2_0); - float2 sum_p1_pair = add_f32x2(sum_p0_pair, _f2_1); - float2 sum_p2_pair = add_f32x2(sum_p1_pair, _f2_2); - float2 next_sum_pair = add_f32x2(sum_p2_pair, _f2_3); - row_sum_pair[0] = next_sum_pair.x; - row_sum_pair[1] = next_sum_pair.y; - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - mbarrier_arrive(score_empty_addr + (score_stage) * 8); - score_stage += 1; - if (score_stage == 2) { score_stage = 0; score_phase ^= 1; } - } - if (num_steps > 0) { - mbarrier_wait(stats_empty_addr + (stats_stage) * 8, stats_phase); - float final_stats_pair[4]; - final_stats_pair[0] = row_sum_pair[0]; - final_stats_pair[1] = row_sum_pair[1]; - final_stats_pair[2] = row_max_pair[0]; - final_stats_pair[3] = row_max_pair[1]; - int final_stats_col = ((stats_stage == 0) ? 16 : 48); - tmem_st_x4_f32(taddr + (unsigned int)final_stats_col, final_stats_pair); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(stats_full_addr + (stats_stage) * 8); - stats_stage += 1; - if (stats_stage == 2) { stats_stage = 0; stats_phase ^= 1; } - } - mbarrier_wait(work_full_addr + (work_stage) * 8, _phase_work_full); - uint32_t _clc_valid_3 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_3) - : "r"(work_response_addr + work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_9 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_9) - : "r"(work_response_addr + work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_10 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_10) - : "r"(work_response_addr + work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_11 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_11) - : "r"(work_response_addr + work_stage * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(work_empty_addr + work_stage * 8), "r"(0) : "memory"); - work_stage += 1; - if (work_stage == 2) { work_stage = 0; _phase_work_full ^= 1; } - if (_clc_valid_3 == 0) { - break; - } - work_x = _clc_ctaid_9; - work_y = _clc_ctaid_10; - work_z = _clc_ctaid_11; - } - } - } - // ---- Role: correction ---- - if (warp >= 4 && warp <= 7) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 136;"); - { // correction_main - unsigned int work_stage_1 = 0; - unsigned int stats_stage_1 = 0; - unsigned int stats_phase_1 = 0; - unsigned int work_x_1 = blockIdx.x; - unsigned int work_y_1 = blockIdx.y; - unsigned int work_z_1 = blockIdx.z; - const int corr_rank = warp - 4; - const int corr_tid = corr_rank * 32 + lane; - const int col_pair_base_1 = corr_tid % 4 * 2; - float softmax_scale_log2_1 = bmm1_scale[0] * 1.4426950408889634f; - float output_scale = bmm2_scale[0]; - unsigned int _phase_o_full_0 = 0; - unsigned int _phase_work_full_1 = 0; - #pragma unroll 1 - for (unsigned int _work_1 = 0; _work_1 < 120; _work_1++) { - int q_begin_1 = cum_seq_lens_q[work_z_1]; - int q_end_1 = cum_seq_lens_q[work_z_1 + 1]; - int q_len_1 = q_end_1 - q_begin_1; - int query_idx_1 = (unsigned int)q_begin_1 + work_x_1; - int active_topk_1 = 0; - if (work_x_1 < (unsigned int)q_len_1) { - int _max_14 = ((sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) : (0)); - int _min_4 = ((_max_14) < (sparse_topk) ? (_max_14) : (sparse_topk)); - active_topk_1 = _min_4; - } - int num_steps_1 = (active_topk_1 + 128 - 1) / 128; - int head_group = work_y_1 / 4; - int value_quarter = work_y_1 % 4; - if (num_steps_1 > 0) { - mbarrier_wait(stats_full_addr + (stats_stage_1) * 8, stats_phase_1); - mbarrier_arrive(stats_empty_addr + (stats_stage_1) * 8); - stats_stage_1 += 1; - if (stats_stage_1 == 2) { stats_stage_1 = 0; stats_phase_1 ^= 1; } - } - #pragma unroll 1 - for (int tile_1 = 1; tile_1 < num_steps_1; tile_1++) { - mbarrier_wait(stats_full_addr + (stats_stage_1) * 8, stats_phase_1); - int stats_col_1 = ((stats_stage_1 == 0) ? 16 : 48); - float _tmem_load_0[4]; - tmem_ld_x4(&_tmem_load_0[0], taddr + (unsigned int)stats_col_1); - asm volatile("tcgen05.wait::ld.sync.aligned;"); - float scale_pair[2]; - #pragma unroll - for (int c_7 = 0; c_7 < 2; c_7++) { - float max_diff = _tmem_load_0[c_7] - _tmem_load_0[c_7 + 2]; - float _exp2_2 = approx_exp2(softmax_scale_log2_1 * max_diff); - scale_pair[c_7] = ((max_diff != 0.0f) ? _exp2_2 : 1.0f); - } - mbarrier_arrive(stats_empty_addr + (stats_stage_1) * 8); - stats_stage_1 += 1; - if (stats_stage_1 == 2) { stats_stage_1 = 0; stats_phase_1 ^= 1; } - mbarrier_wait(o_full_addr, _phase_o_full_0); - _phase_o_full_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - int needs_scale = ((scale_pair[0] != 1.0f) ? 1 : 0); - needs_scale = needs_scale | ((scale_pair[1] != 1.0f) ? 1 : 0); - int _vote_0 = __any_sync(0xFFFFFFFF, needs_scale != 0); - needs_scale = _vote_0; - if (needs_scale != 0) { - float values_lo[4]; - float values_hi[4]; - asm volatile( - "tcgen05.ld.sync.aligned.16x256b.x1.b32" - " {%0, %1, %2, %3}, [%4];" - : "=r"(*reinterpret_cast(&values_lo[0])), "=r"(*reinterpret_cast(&values_lo[1])), "=r"(*reinterpret_cast(&values_lo[2])), "=r"(*reinterpret_cast(&values_lo[3])) - : "r"(taddr + 80)); - asm volatile( - "tcgen05.ld.sync.aligned.16x256b.x1.b32" - " {%0, %1, %2, %3}, [%4];" - : "=r"(*reinterpret_cast(&values_hi[0])), "=r"(*reinterpret_cast(&values_hi[1])), "=r"(*reinterpret_cast(&values_hi[2])), "=r"(*reinterpret_cast(&values_hi[3])) - : "r"(taddr + 80 + 1048576)); - #pragma unroll - for (int pair = 0; pair < 4; pair++) { - asm volatile("" ::: "memory"); - values_lo[pair] = values_lo[pair] * scale_pair[pair % 2]; - values_hi[pair] = values_hi[pair] * scale_pair[pair % 2]; - } - asm volatile( - "tcgen05.st.sync.aligned.16x256b.x1.b32" - " [%0], {%1, %2, %3, %4};" - :: "r"(taddr + 80), "r"(*reinterpret_cast(&values_lo[0])), "r"(*reinterpret_cast(&values_lo[1])), "r"(*reinterpret_cast(&values_lo[2])), "r"(*reinterpret_cast(&values_lo[3]))); - asm volatile( - "tcgen05.st.sync.aligned.16x256b.x1.b32" - " [%0], {%1, %2, %3, %4};" - :: "r"(taddr + 80 + 1048576), "r"(*reinterpret_cast(&values_hi[0])), "r"(*reinterpret_cast(&values_hi[1])), "r"(*reinterpret_cast(&values_hi[2])), "r"(*reinterpret_cast(&values_hi[3]))); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - } - mbarrier_arrive(o_empty_addr); - } - if (num_steps_1 > 0) { - mbarrier_wait(stats_full_addr + (stats_stage_1) * 8, stats_phase_1); - int final_stats_col_1 = ((stats_stage_1 == 0) ? 16 : 48); - float _tmem_load_1[4]; - tmem_ld_x4(&_tmem_load_1[0], taddr + (unsigned int)final_stats_col_1); - asm volatile("tcgen05.wait::ld.sync.aligned;"); - mbarrier_arrive(stats_empty_addr + (stats_stage_1) * 8); - stats_stage_1 += 1; - if (stats_stage_1 == 2) { stats_stage_1 = 0; stats_phase_1 ^= 1; } - float pair_sum[2]; - pair_sum[0] = _tmem_load_1[0]; - pair_sum[1] = _tmem_load_1[1]; - float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, pair_sum[0], 16); - pair_sum[0] = pair_sum[0] + _shfl_xor_2; - float _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, pair_sum[1], 16); - pair_sum[1] = pair_sum[1] + _shfl_xor_3; - float _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, pair_sum[0], 8); - pair_sum[0] = pair_sum[0] + _shfl_xor_4; - float _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, pair_sum[1], 8); - pair_sum[1] = pair_sum[1] + _shfl_xor_5; - float _shfl_xor_6 = __shfl_xor_sync(0xFFFFFFFF, pair_sum[0], 4); - pair_sum[0] = pair_sum[0] + _shfl_xor_6; - float _shfl_xor_7 = __shfl_xor_sync(0xFFFFFFFF, pair_sum[1], 4); - pair_sum[1] = pair_sum[1] + _shfl_xor_7; - if (lane < 4) { - smem_corr[corr_rank * 8 + col_pair_base_1] = pair_sum[0]; - smem_corr[corr_rank * 8 + col_pair_base_1 + 1] = pair_sum[1]; - } - asm volatile("barrier.sync 2, 128;" ::: "memory"); - float sum0 = smem_corr[col_pair_base_1] + smem_corr[8 + col_pair_base_1] + smem_corr[16 + col_pair_base_1] + smem_corr[24 + col_pair_base_1]; - float sum1 = smem_corr[col_pair_base_1 + 1] + smem_corr[8 + col_pair_base_1 + 1] + smem_corr[16 + col_pair_base_1 + 1] + smem_corr[24 + col_pair_base_1 + 1]; - if (has_sinks != 0) { - float _exp2_3 = approx_exp2(sinks[head_group * 8 + col_pair_base_1] * 1.4426950408889634f - _tmem_load_1[2] * softmax_scale_log2_1); - sum0 = sum0 + _exp2_3; - float _exp2_4 = approx_exp2(sinks[head_group * 8 + col_pair_base_1 + 1] * 1.4426950408889634f - _tmem_load_1[3] * softmax_scale_log2_1); - sum1 = sum1 + _exp2_4; - } - float _rcp_0 = approx_rcp(sum0); - float scale0 = ((sum0 > 0.0f) ? output_scale * 0.002232142857142857f * _rcp_0 : 0.0f); - float _rcp_1 = approx_rcp(sum1); - float scale1 = ((sum1 > 0.0f) ? output_scale * 0.002232142857142857f * _rcp_1 : 0.0f); - mbarrier_wait(o_full_addr, _phase_o_full_0); - _phase_o_full_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - float output_lo[4]; - float output_hi[4]; - asm volatile( - "tcgen05.ld.sync.aligned.16x256b.x1.b32" - " {%0, %1, %2, %3}, [%4];" - : "=r"(*reinterpret_cast(&output_lo[0])), "=r"(*reinterpret_cast(&output_lo[1])), "=r"(*reinterpret_cast(&output_lo[2])), "=r"(*reinterpret_cast(&output_lo[3])) - : "r"(taddr + 80)); - asm volatile( - "tcgen05.ld.sync.aligned.16x256b.x1.b32" - " {%0, %1, %2, %3}, [%4];" - : "=r"(*reinterpret_cast(&output_hi[0])), "=r"(*reinterpret_cast(&output_hi[1])), "=r"(*reinterpret_cast(&output_hi[2])), "=r"(*reinterpret_cast(&output_hi[3])) - : "r"(taddr + 80 + 1048576)); - unsigned int packed_o[4]; - #pragma unroll - for (int pair_1 = 0; pair_1 < 4; pair_1++) { - asm volatile("" ::: "memory"); - output_lo[pair_1] = output_lo[pair_1] * ((pair_1 % 2 == 0) ? scale0 : scale1); - output_hi[pair_1] = output_hi[pair_1] * ((pair_1 % 2 == 0) ? scale0 : scale1); - } - float output_values[8]; - #pragma unroll - for (int c_8 = 0; c_8 < 4; c_8++) { - output_values[c_8] = output_lo[c_8]; - output_values[c_8 + 4] = output_hi[c_8]; - } - #pragma unroll - for (int _lp = 0; _lp < 4; _lp++) { - __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(output_values[_lp*2 + 0], output_values[_lp*2+1 + 0])); - packed_o[_lp] = *(uint32_t*)&_bf2; - } - int slice_idx = corr_rank / 2; - int warp_idx_in_slice = corr_rank % 2; - int mtx_idx_1 = lane / 8; - int thr_row_idx_1 = lane % 8; - int seg_col_idx_1 = warp_idx_in_slice * 4 + mtx_idx_1 ^ thr_row_idx_1; - int stsm_offset_1 = slice_idx * 8 * 128 + thr_row_idx_1 * 128 + seg_col_idx_1 * 16; - asm volatile("stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 [%0], {%1, %2, %3, %4};\n" - :: "l"(__cvta_generic_to_shared(reinterpret_cast(smem_o) + stsm_offset_1)), "r"(*reinterpret_cast(&packed_o[0])), "r"(*reinterpret_cast(&packed_o[1])), "r"(*reinterpret_cast(&packed_o[2])), "r"(*reinterpret_cast(&packed_o[3])) - : "memory"); - asm volatile("barrier.sync 2, 128;" ::: "memory"); - int copy_base = corr_tid * 16; - int smem_row = copy_base / 128; - int copy_smem_offset = copy_base ^ smem_row % 8 * 16; - int copy_row = smem_row % 8; - int copy_col = (smem_row / 8 * 128 + copy_base % 128) / 2; - unsigned int copy_vec[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(©_vec[0])), "=r"(*reinterpret_cast(©_vec[(0) + 1])), "=r"(*reinterpret_cast(©_vec[(0) + 2])), "=r"(*reinterpret_cast(©_vec[(0) + 3])) - : "r"(smem_o_addr + (unsigned int)copy_smem_offset)); - int head_groups = gridDim.y / 4; - int output_base = query_idx_1 * head_groups * 8 * 512 + head_group * 8 * 512 + value_quarter * 128; - reinterpret_cast(O + (output_base + copy_row * 512 + copy_col))[0] = reinterpret_cast(copy_vec)[0]; - mbarrier_arrive(o_empty_addr); - } - mbarrier_wait(work_full_addr + (work_stage_1) * 8, _phase_work_full_1); - uint32_t _clc_valid_4 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_4) - : "r"(work_response_addr + work_stage_1 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_12 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_12) - : "r"(work_response_addr + work_stage_1 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_13 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_13) - : "r"(work_response_addr + work_stage_1 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_14 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_14) - : "r"(work_response_addr + work_stage_1 * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(work_empty_addr + work_stage_1 * 8), "r"(0) : "memory"); - work_stage_1 += 1; - if (work_stage_1 == 2) { work_stage_1 = 0; _phase_work_full_1 ^= 1; } - if (_clc_valid_4 == 0) { - break; - } - work_x_1 = _clc_ctaid_12; - work_y_1 = _clc_ctaid_13; - work_z_1 = _clc_ctaid_14; - } - asm volatile("barrier.sync 11, 128;" ::: "memory"); - if (corr_rank == 0) { - int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); - asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); - } - } - } - // ---- Role: mma ---- - if (warp == 8) { - { // mma_main - unsigned int work_stage_2 = 0; - unsigned int q_stage = 0; - unsigned int q_phase = 0; - unsigned int kv_stage = 0; - unsigned int kv_phase = 0; - unsigned int score_stage_1 = 0; - unsigned int score_phase_1 = 1; - unsigned int p_stage_1 = 0; - unsigned int p_phase_1 = 0; - unsigned int work_x_2 = blockIdx.x; - unsigned int work_y_2 = blockIdx.y; - unsigned int work_z_2 = blockIdx.z; - unsigned int _phase_o_empty_0 = 1; - unsigned int _phase_work_full_2 = 0; - #pragma unroll 1 - for (unsigned int _work_2 = 0; _work_2 < 120; _work_2++) { - int q_begin_2 = cum_seq_lens_q[work_z_2]; - int q_end_2 = cum_seq_lens_q[work_z_2 + 1]; - int q_len_2 = q_end_2 - q_begin_2; - int query_idx_2 = (unsigned int)q_begin_2 + work_x_2; - int active_topk_2 = 0; - if (work_x_2 < (unsigned int)q_len_2) { - int _max_3 = ((sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) : (0)); - int _min_2 = ((_max_3) < (sparse_topk) ? (_max_3) : (sparse_topk)); - active_topk_2 = _min_2; - } - int num_steps_2 = (active_topk_2 + 128 - 1) / 128; - int value_quarter_1 = work_y_2 % 4; - if (num_steps_2 > 0) { - mbarrier_wait(q_full_addr + (q_stage) * 8, q_phase); - int first_output = 1; - mbarrier_wait(score_empty_addr + (score_stage_1) * 8, score_phase_1); - int score_col_1 = ((score_stage_1 == 0) ? 0 : 8); - #pragma unroll - for (int hs = 0; hs < 3; hs++) { - int hs_sel = ((value_quarter_1 > hs) ? hs : hs + 1); - mbarrier_wait(kv_full_addr + (kv_stage) * 8, kv_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - int _mma_a_lo_0 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage) * 1024); - int _mma_b_lo_0 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (hs_sel) * 64); - { - uint64_t _mma_ss_a_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_0); - uint64_t _mma_ss_b_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_0); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134348816, ((hs == 0) ? 0 : 1)); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134348816, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134348816, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134348816, 1); - } - } - elect_commit(kv_empty_addr + (kv_stage) * 8); - kv_stage += 1; - if (kv_stage == 8) { kv_stage = 0; kv_phase ^= 1; } - } - mbarrier_wait(kv_full_addr + (kv_stage) * 8, kv_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - int _mma_a_lo_1 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage) * 1024); - int _mma_b_lo_1 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (value_quarter_1) * 64); - { - uint64_t _mma_ss_a_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_1); - uint64_t _mma_ss_b_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_1); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134348816, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134348816, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134348816, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134348816, 1); - } - } - kv_stage += 1; - if (kv_stage == 8) { kv_stage = 0; kv_phase ^= 1; } - elect_commit(score_full_addr + (score_stage_1) * 8); - score_stage_1 += 1; - if (score_stage_1 == 2) { score_stage_1 = 0; score_phase_1 ^= 1; } - #pragma unroll 1 - for (int _tile = 0; _tile < num_steps_2 - 1; _tile++) { - mbarrier_wait(score_empty_addr + (score_stage_1) * 8, score_phase_1); - int score_col_0 = ((score_stage_1 == 0) ? 0 : 8); - #pragma unroll - for (int hs_1 = 0; hs_1 < 3; hs_1++) { - int hs_sel_1 = ((value_quarter_1 > hs_1) ? hs_1 : hs_1 + 1); - mbarrier_wait(kv_full_addr + (kv_stage) * 8, kv_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - int _mma_a_lo_2 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage) * 1024); - int _mma_b_lo_2 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (hs_sel_1) * 64); - { - uint64_t _mma_ss_a_desc_2 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_2); - uint64_t _mma_ss_b_desc_2 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_2); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134348816, ((hs_1 == 0) ? 0 : 1)); - } - incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134348816, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134348816, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134348816, 1); - } - } - elect_commit(kv_empty_addr + (kv_stage) * 8); - kv_stage += 1; - if (kv_stage == 8) { kv_stage = 0; kv_phase ^= 1; } - } - mbarrier_wait(kv_full_addr + (kv_stage) * 8, kv_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - int _mma_a_lo_3 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage) * 1024); - int _mma_b_lo_3 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (value_quarter_1) * 64); - { - uint64_t _mma_ss_a_desc_3 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_3); - uint64_t _mma_ss_b_desc_3 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_3); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134348816, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134348816, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134348816, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134348816, 1); - } - } - kv_stage += 1; - if (kv_stage == 8) { kv_stage = 0; kv_phase ^= 1; } - elect_commit(score_full_addr + (score_stage_1) * 8); - score_stage_1 += 1; - if (score_stage_1 == 2) { score_stage_1 = 0; score_phase_1 ^= 1; } - unsigned int v_cur = (kv_stage + 3) % 8; - mbarrier_wait(p_full_addr + (p_stage_1) * 8, p_phase_1); - mbarrier_wait(o_empty_addr, _phase_o_empty_0); - _phase_o_empty_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - int _mma_a_lo_4 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (v_cur) * 1024); - int _mma_b_lo_4 = make_warp_uniform(((((smem_p_addr) >> 4) & 0x3FFF) | 0x400000) + (p_stage_1) * 64); - { - uint64_t _mma_ss_a_desc_4 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_4); - uint64_t _mma_ss_b_desc_4 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_4); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134381584, ((first_output) ? 0 : 1)); - } - incr_smem_desc_lo(_mma_ss_a_desc_4, 256U); - incr_smem_desc_lo(_mma_ss_b_desc_4, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134381584, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_4, 256U); - incr_smem_desc_lo(_mma_ss_b_desc_4, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134381584, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_4, 256U); - incr_smem_desc_lo(_mma_ss_b_desc_4, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134381584, 1); - } - } - elect_commit(o_full_addr); - elect_commit(kv_empty_addr + (v_cur) * 8); - elect_commit(p_empty_addr + (p_stage_1) * 8); - p_stage_1 += 1; - if (p_stage_1 == 2) { p_stage_1 = 0; p_phase_1 ^= 1; } - first_output = 0; - } - unsigned int v_tail = (kv_stage + 7) % 8; - mbarrier_wait(p_full_addr + (p_stage_1) * 8, p_phase_1); - mbarrier_wait(o_empty_addr, _phase_o_empty_0); - _phase_o_empty_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - int _mma_a_lo_5 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (v_tail) * 1024); - int _mma_b_lo_5 = make_warp_uniform(((((smem_p_addr) >> 4) & 0x3FFF) | 0x400000) + (p_stage_1) * 64); - { - uint64_t _mma_ss_a_desc_5 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_5); - uint64_t _mma_ss_b_desc_5 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_5); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134381584, ((first_output) ? 0 : 1)); - } - incr_smem_desc_lo(_mma_ss_a_desc_5, 256U); - incr_smem_desc_lo(_mma_ss_b_desc_5, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134381584, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_5, 256U); - incr_smem_desc_lo(_mma_ss_b_desc_5, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134381584, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_5, 256U); - incr_smem_desc_lo(_mma_ss_b_desc_5, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134381584, 1); - } - } - elect_commit(o_full_addr); - elect_commit(kv_empty_addr + (v_tail) * 8); - elect_commit(p_empty_addr + (p_stage_1) * 8); - p_stage_1 += 1; - if (p_stage_1 == 2) { p_stage_1 = 0; p_phase_1 ^= 1; } - elect_commit(q_empty_addr + (q_stage) * 8); - q_phase ^= 1; - } - mbarrier_wait(work_full_addr + (work_stage_2) * 8, _phase_work_full_2); - uint32_t _clc_valid_2 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_2) - : "r"(work_response_addr + work_stage_2 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_6 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_6) - : "r"(work_response_addr + work_stage_2 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_7 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_7) - : "r"(work_response_addr + work_stage_2 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_8 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_8) - : "r"(work_response_addr + work_stage_2 * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(work_empty_addr + work_stage_2 * 8), "r"(0) : "memory"); - work_stage_2 += 1; - if (work_stage_2 == 2) { work_stage_2 = 0; _phase_work_full_2 ^= 1; } - if (_clc_valid_2 == 0) { - break; - } - work_x_2 = _clc_ctaid_6; - work_y_2 = _clc_ctaid_7; - work_z_2 = _clc_ctaid_8; - } - #pragma unroll - for (int tail = 0; tail < 2; tail++) { - elect_commit(score_full_addr + (tail) * 8); - } - } - } - // ---- Role: page_loader ---- - if (warp == 9) { - { // page_loader_main - unsigned int work_stage_3 = 0; - unsigned int page_stage = 0; - unsigned int page_phase = 1; - unsigned int work_x_3 = blockIdx.x; - unsigned int work_y_3 = blockIdx.y; - unsigned int work_z_3 = blockIdx.z; - unsigned int _phase_work_full_3 = 0; - #pragma unroll 1 - for (unsigned int _work_3 = 0; _work_3 < 120; _work_3++) { - int q_begin_3 = cum_seq_lens_q[work_z_3]; - int q_end_3 = cum_seq_lens_q[work_z_3 + 1]; - int q_len_3 = q_end_3 - q_begin_3; - int query_idx_3 = (unsigned int)q_begin_3 + work_x_3; - int active_topk_3 = 0; - if (work_x_3 < (unsigned int)q_len_3) { - int _max_0 = ((sparse_topk_lens[query_idx_3] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_3] + sparse_topk_lens_offset) : (0)); - int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); - active_topk_3 = _min_0; - } - int num_steps_3 = (active_topk_3 + 128 - 1) / 128; - int num_pairs = (num_steps_3 + 1) / 2; - int _max_1 = ((active_topk_3 - 1 & -4) > (0) ? (active_topk_3 - 1 & -4) : (0)); - int last_aligned = _max_1; - int* swa_row = swa_indices + (query_idx_3 * swa_index_stride); - int* compressed_row = compressed_indices + (query_idx_3 * compressed_index_stride); - #pragma unroll 1 - for (int pair_2 = 0; pair_2 < num_pairs; pair_2++) { - int lane_base = lane * 4; - #pragma unroll - for (int duplicate = 0; duplicate < 2; duplicate++) { - mbarrier_wait(page_empty_addr + (page_stage) * 8, page_phase); - int page0 = pair_2 * 2 * 128 + lane_base; - int page1 = page0 + 128; - int safe0 = ((page0 <= last_aligned) ? page0 : last_aligned); - int safe1 = ((page1 <= last_aligned) ? page1 : last_aligned); - int* src0 = ((safe0 < 128) ? swa_row + safe0 : compressed_row + (safe0 - 128)); - int* src1 = ((safe1 < 128) ? swa_row + safe1 : compressed_row + (safe1 - 128)); - asm volatile("cp.async.ca.shared::cta.global [%0], [%1], 16;" - :: "r"(smem_page_offsets_addr + page_stage * 1024 + (unsigned int)(lane_base * 4)), "l"(src0)); - asm volatile("cp.async.ca.shared::cta.global [%0], [%1], 16;" - :: "r"(smem_page_offsets_addr + page_stage * 1024 + (unsigned int)((128 + lane_base) * 4)), "l"(src1)); - asm volatile( - "{\n\t" - "cp.async.mbarrier.arrive.noinc.shared::cta.b64 [%0];\n\t" - "}" - :: "r"(page_full_addr + (page_stage) * 8) : "memory"); - page_stage += 1; - if (page_stage == 6) { page_stage = 0; page_phase ^= 1; } - } - } - mbarrier_wait(work_full_addr + (work_stage_3) * 8, _phase_work_full_3); - uint32_t _clc_valid_0 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_0) - : "r"(work_response_addr + work_stage_3 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_0 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_0) - : "r"(work_response_addr + work_stage_3 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_1 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_1) - : "r"(work_response_addr + work_stage_3 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_2 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_2) - : "r"(work_response_addr + work_stage_3 * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(work_empty_addr + work_stage_3 * 8), "r"(0) : "memory"); - work_stage_3 += 1; - if (work_stage_3 == 2) { work_stage_3 = 0; _phase_work_full_3 ^= 1; } - if (_clc_valid_0 == 0) { - break; - } - work_x_3 = _clc_ctaid_0; - work_y_3 = _clc_ctaid_1; - work_z_3 = _clc_ctaid_2; - } - } - } - // ---- Role: scheduler ---- - if (warp == 10) { - { // scheduler_main - asm volatile("griddepcontrol.wait;" ::: "memory"); - unsigned int work_stage_4 = 0; - unsigned int throttle_stage = 0; - unsigned int throttle_phase = 0; - unsigned int _phase_work_empty = 1; - unsigned int _phase_work_full_4 = 0; - #pragma unroll 1 - for (unsigned int _work_4 = 0; _work_4 < 120; _work_4++) { - mbarrier_wait(throttle_full_addr + (throttle_stage) * 8, throttle_phase); - mbarrier_arrive(throttle_empty_addr + (throttle_stage) * 8); - throttle_stage += 1; - if (throttle_stage == 2) { throttle_stage = 0; throttle_phase ^= 1; } - mbarrier_wait(work_empty_addr + (work_stage_4) * 8, _phase_work_empty); - if (elect_sync()) { - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [remAddr32], %2;\n\t" - "}" - :: "r"(work_full_addr + work_stage_4 * 8), "r"(0), "r"((uint32_t)(16)) : "memory"); - asm volatile( - "fence.proxy.async.shared::cta;\n\t" - "clusterlaunchcontrol.try_cancel.async.shared::cta" - ".mbarrier::complete_tx::bytes.b128" - " [%0], [%1];" - :: "r"(work_response_addr + work_stage_4 * 16 + 0 * 16), "r"(work_full_addr + work_stage_4 * 8) - : "memory"); - } - mbarrier_wait(work_full_addr + (work_stage_4) * 8, _phase_work_full_4); - uint32_t _clc_valid_5 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_5) - : "r"(work_response_addr + work_stage_4 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_15 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_15) - : "r"(work_response_addr + work_stage_4 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_16 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_16) - : "r"(work_response_addr + work_stage_4 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_17 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_17) - : "r"(work_response_addr + work_stage_4 * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(work_empty_addr + work_stage_4 * 8), "r"(0) : "memory"); - work_stage_4 += 1; - if (work_stage_4 == 2) { work_stage_4 = 0; _phase_work_empty ^= 1; _phase_work_full_4 ^= 1; } - if (_clc_valid_5 == 0) { - break; - } - } - } - } - // ---- Role: padding ---- - if (warp == 11) { - { // padding_main - unsigned int work_stage_5 = 0; - unsigned int work_x_4 = blockIdx.x; - unsigned int work_y_4 = blockIdx.y; - unsigned int work_z_4 = blockIdx.z; - unsigned int _phase_work_full_5 = 0; - #pragma unroll 1 - for (unsigned int _work_5 = 0; _work_5 < 120; _work_5++) { - mbarrier_wait(work_full_addr + (work_stage_5) * 8, _phase_work_full_5); - uint32_t _clc_valid_6 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_6) - : "r"(work_response_addr + work_stage_5 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_18 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_18) - : "r"(work_response_addr + work_stage_5 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_19 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_19) - : "r"(work_response_addr + work_stage_5 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_20 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_20) - : "r"(work_response_addr + work_stage_5 * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(work_empty_addr + work_stage_5 * 8), "r"(0) : "memory"); - work_stage_5 += 1; - if (work_stage_5 == 2) { work_stage_5 = 0; _phase_work_full_5 ^= 1; } - if (_clc_valid_6 == 0) { - break; - } - work_x_4 = _clc_ctaid_18; - work_y_4 = _clc_ctaid_19; - work_z_4 = _clc_ctaid_20; - } - } - } - // ---- Role: loader ---- - if (warp >= 12 && warp <= 19) { - asm volatile("setmaxnreg.dec.sync.aligned.u32 72;"); - { // loader_main - asm volatile("griddepcontrol.wait;" ::: "memory"); - const int load_rank = warp - 12; - unsigned int work_stage_6 = 0; - unsigned int throttle_stage_1 = 0; - unsigned int throttle_phase_1 = 1; - unsigned int q_stage_1 = 0; - unsigned int q_phase_1 = 1; - unsigned int kv_stage_1 = 0; - unsigned int kv_phase_1 = 1; - unsigned int page_k_stage = 0; - unsigned int page_k_phase = 0; - unsigned int page_v_stage = 1; - unsigned int page_v_phase = 0; - unsigned int work_x_5 = blockIdx.x; - unsigned int work_y_5 = blockIdx.y; - unsigned int work_z_5 = blockIdx.z; - unsigned int _phase_work_full_6 = 0; - #pragma unroll 1 - for (unsigned int _work_6 = 0; _work_6 < 120; _work_6++) { - int q_begin_4 = cum_seq_lens_q[work_z_5]; - int q_end_4 = cum_seq_lens_q[work_z_5 + 1]; - int q_len_4 = q_end_4 - q_begin_4; - int query_idx_4 = (unsigned int)q_begin_4 + work_x_5; - int active_topk_4 = 0; - if (work_x_5 < (unsigned int)q_len_4) { - int _max_2 = ((sparse_topk_lens[query_idx_4] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_4] + sparse_topk_lens_offset) : (0)); - int _min_1 = ((_max_2) < (sparse_topk) ? (_max_2) : (sparse_topk)); - active_topk_4 = _min_1; - } - int num_steps_4 = (active_topk_4 + 128 - 1) / 128; - int head_group_1 = work_y_5 / 4; - int value_quarter_2 = work_y_5 % 4; - mbarrier_wait(throttle_empty_addr + (throttle_stage_1) * 8, throttle_phase_1); - mbarrier_arrive(throttle_full_addr + (throttle_stage_1) * 8); - throttle_stage_1 += 1; - if (throttle_stage_1 == 2) { throttle_stage_1 = 0; throttle_phase_1 ^= 1; } - if (num_steps_4 > 0) { - mbarrier_wait(q_empty_addr + (q_stage_1) * 8, q_phase_1); - if (load_rank == 0) { - if (elect_sync()) { - mbarrier_arrive_expect_tx(q_full_addr + (q_stage_1) * 8, 4096); - #pragma unroll - for (int hs_2 = 0; hs_2 < 4; hs_2++) { - tma_4d_gmem2smem(smem_q_addr + (unsigned int)(hs_2 * 1024), (&tmap_q), 0, head_group_1 * 8, hs_2, query_idx_4, q_full_addr + (q_stage_1) * 8); - } - } - } - mbarrier_wait(page_full_addr + (page_k_stage) * 8, page_k_phase); - int page_base = 0; - #pragma unroll - for (int hs_3 = 0; hs_3 < 4; hs_3++) { - int hs_sel_2 = ((hs_3 == 3) ? value_quarter_2 : ((value_quarter_2 > hs_3) ? hs_3 : hs_3 + 1)); - mbarrier_wait(kv_empty_addr + (kv_stage_1) * 8, kv_phase_1); - if (load_rank == 0) { - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_1) * 8, 16384); - } - } - asm volatile("barrier.sync 9, 256;" ::: "memory"); - int rows_all[16]; - #pragma unroll - for (int group = 0; group < 4; group++) { - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&rows_all[group * 4])), "=r"(*reinterpret_cast(&rows_all[(group * 4) + 1])), "=r"(*reinterpret_cast(&rows_all[(group * 4) + 2])), "=r"(*reinterpret_cast(&rows_all[(group * 4) + 3])) - : "r"(smem_page_offsets_addr + page_k_stage * 1024 + (unsigned int)((page_base + load_rank * 4 + group * 32) * 4))); - } - if (elect_sync()) { - { - #pragma unroll - for (int group_1 = 0; group_1 < 4; group_1++) { - const int row_base = load_rank * 4 + group_1 * 32; - int dst_k = smem_kv_addr + kv_stage_1 * 16384 + (unsigned int)(row_base * 128); - tma_gather4_gmem2smem(dst_k, (&tmap_swa_kv), hs_sel_2 * 128, rows_all[group_1 * 4], rows_all[group_1 * 4 + 1], rows_all[group_1 * 4 + 2], rows_all[group_1 * 4 + 3], kv_full_addr + (kv_stage_1) * 8); - } - } - } - kv_stage_1 += 1; - if (kv_stage_1 == 8) { kv_stage_1 = 0; kv_phase_1 ^= 1; } - } - #pragma unroll 1 - for (int tile_2 = 0; tile_2 < num_steps_4 - 1; tile_2++) { - int next_tile = tile_2 + 1; - if ((next_tile & 1) == 0) { - mbarrier_wait(page_full_addr + (page_k_stage) * 8, page_k_phase); - } - int page_base_0 = (next_tile & 1) * 128; - #pragma unroll - for (int hs_4 = 0; hs_4 < 4; hs_4++) { - int hs_sel_3 = ((hs_4 == 3) ? value_quarter_2 : ((value_quarter_2 > hs_4) ? hs_4 : hs_4 + 1)); - mbarrier_wait(kv_empty_addr + (kv_stage_1) * 8, kv_phase_1); - if (load_rank == 0) { - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_1) * 8, 16384); - } - } - asm volatile("barrier.sync 9, 256;" ::: "memory"); - int rows_all_1[16]; - #pragma unroll - for (int group_2 = 0; group_2 < 4; group_2++) { - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&rows_all_1[group_2 * 4])), "=r"(*reinterpret_cast(&rows_all_1[(group_2 * 4) + 1])), "=r"(*reinterpret_cast(&rows_all_1[(group_2 * 4) + 2])), "=r"(*reinterpret_cast(&rows_all_1[(group_2 * 4) + 3])) - : "r"(smem_page_offsets_addr + page_k_stage * 1024 + (unsigned int)((page_base_0 + load_rank * 4 + group_2 * 32) * 4))); - } - if (elect_sync()) { - if (next_tile == 0) { - #pragma unroll - for (int group_3 = 0; group_3 < 4; group_3++) { - const int row_base_1 = load_rank * 4 + group_3 * 32; - int dst_k_1 = smem_kv_addr + kv_stage_1 * 16384 + (unsigned int)(row_base_1 * 128); - tma_gather4_gmem2smem(dst_k_1, (&tmap_swa_kv), hs_sel_3 * 128, rows_all_1[group_3 * 4], rows_all_1[group_3 * 4 + 1], rows_all_1[group_3 * 4 + 2], rows_all_1[group_3 * 4 + 3], kv_full_addr + (kv_stage_1) * 8); - } - } else { - #pragma unroll - for (int group_4 = 0; group_4 < 4; group_4++) { - const int row_base_2 = load_rank * 4 + group_4 * 32; - int dst_k_2 = smem_kv_addr + kv_stage_1 * 16384 + (unsigned int)(row_base_2 * 128); - tma_gather4_gmem2smem(dst_k_2, (&tmap_compressed_kv), hs_sel_3 * 128, rows_all_1[group_4 * 4], rows_all_1[group_4 * 4 + 1], rows_all_1[group_4 * 4 + 2], rows_all_1[group_4 * 4 + 3], kv_full_addr + (kv_stage_1) * 8); - } - } - } - kv_stage_1 += 1; - if (kv_stage_1 == 8) { kv_stage_1 = 0; kv_phase_1 ^= 1; } - } - if ((next_tile & 1) != 0 && next_tile < num_steps_4 - 1) { - mbarrier_arrive(page_empty_addr + (page_k_stage) * 8); - page_k_stage += 1; - if (page_k_stage == 6) { page_k_stage = 0; page_k_phase ^= 1; } - page_k_stage += 1; - if (page_k_stage == 6) { page_k_stage = 0; page_k_phase ^= 1; } - } - if ((tile_2 & 1) == 0) { - mbarrier_wait(page_full_addr + (page_v_stage) * 8, page_v_phase); - } - if ((tile_2 & 1) != 0) { - mbarrier_arrive(page_empty_addr + (page_v_stage) * 8); - page_v_stage += 1; - if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } - page_v_stage += 1; - if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } - } - } - int last_tile = num_steps_4 - 1; - mbarrier_arrive(page_empty_addr + (page_k_stage) * 8); - page_k_stage += 1; - if (page_k_stage == 6) { page_k_stage = 0; page_k_phase ^= 1; } - page_k_stage += 1; - if (page_k_stage == 6) { page_k_stage = 0; page_k_phase ^= 1; } - if ((last_tile & 1) == 0) { - mbarrier_wait(page_full_addr + (page_v_stage) * 8, page_v_phase); - } - mbarrier_arrive(page_empty_addr + (page_v_stage) * 8); - page_v_stage += 1; - if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } - page_v_stage += 1; - if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } - } - mbarrier_wait(work_full_addr + (work_stage_6) * 8, _phase_work_full_6); - uint32_t _clc_valid_1 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_1) - : "r"(work_response_addr + work_stage_6 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_3 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_3) - : "r"(work_response_addr + work_stage_6 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_4 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_4) - : "r"(work_response_addr + work_stage_6 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_5 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_5) - : "r"(work_response_addr + work_stage_6 * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(work_empty_addr + work_stage_6 * 8), "r"(0) : "memory"); - work_stage_6 += 1; - if (work_stage_6 == 2) { work_stage_6 = 0; _phase_work_full_6 ^= 1; } - if (_clc_valid_1 == 0) { - break; - } - work_x_5 = _clc_ctaid_3; - work_y_5 = _clc_ctaid_4; - work_z_5 = _clc_ctaid_5; - } - } - } - - // Cleanup -} - -} // extern "C" diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_125332c2066d89bd7610_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_125332c2066d89bd7610_binding.cu deleted file mode 100644 index e463bc58de8..00000000000 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_125332c2066d89bd7610_binding.cu +++ /dev/null @@ -1,555 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -// Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_125332c2066d89bd7610'. -#include -#include -#include -#include - -#include "tvm_ffi_utils.h" - -#include -#include -#include -#include -#include -#include -#include - -extern "C" __global__ void kernel_cake_dsv4_125332c2066d89bd7610(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); - - -namespace cake_host_shim_0f71e276da509b5c { - -using tvm::ffi::Optional; -using tvm::ffi::TensorView; - -class ScopedCudaDevice { - public: - explicit ScopedCudaDevice(int device_id) { - cudaError_t error = cudaGetDevice(&previous_device_); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed before host-shim launch: cudaError=" - << static_cast(error); - if (previous_device_ != device_id) { - error = cudaSetDevice(device_id); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaSetDevice failed before host-shim launch for cuda:" - << device_id << ": cudaError=" << static_cast(error); - restore_ = true; - } - } - - ScopedCudaDevice(const ScopedCudaDevice&) = delete; - ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; - - ~ScopedCudaDevice() noexcept { - if (restore_) { - (void)cudaSetDevice(previous_device_); - } - } - - private: - int previous_device_ = -1; - bool restore_ = false; -}; - -inline int64_t CakeDeviceMultiprocessorCount(int device_id) { - constexpr int kMaxCachedCudaDevices = 64; - TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) - << "physical-SM-count cache does not cover cuda:" << device_id; - static std::atomic count_by_device[kMaxCachedCudaDevices]{}; - int cached = count_by_device[device_id].load(std::memory_order_acquire); - if (cached > 0) return cached; - - int count = 0; - cudaError_t error = cudaDeviceGetAttribute( - &count, cudaDevAttrMultiProcessorCount, device_id); - TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) - << "querying multiProcessorCount failed for cuda:" << device_id - << ": cudaError=" << static_cast(error) << ", count=" << count; - // Concurrent first launches may repeat the immutable device query, but all - // publication is atomic and every later hot-path lookup is one acquire load. - count_by_device[device_id].store(count, std::memory_order_release); - return count; -} - -inline void CheckCudaTensor(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) - << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; -} - -inline void CheckSameCudaDevice( - const TensorView& t, - const TensorView& reference, - const char* name, - const char* reference_name) { - TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) - << name << " must be on the same CUDA device as " << reference_name - << ": got cuda:" << t.device().device_id - << " versus cuda:" << reference.device().device_id; -} - -inline void CheckCurrentCudaDevice( - const TensorView& reference, - const char* reference_name) { - int current_device = -1; - cudaError_t error = cudaGetDevice(¤t_device); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed while validating " << reference_name - << ": cudaError=" << static_cast(error); - TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) - << "current CUDA device must match " << reference_name - << ": current=cuda:" << current_device - << ", tensor=cuda:" << reference.device().device_id; -} - -inline void CheckContiguous(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; -} - -inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { - DLDataType d = t.dtype(); - TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) - << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits - << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits - << ", lanes=" << (int)d.lanes << ")"; -} - -// A logical axis.outer(trailing) folds every source dim above the trailing -// dimensions. Shape products are independent of physical strides, so verify -// the leading dimensions form one dense row-major chain instead of inventing -// a "folded stride". The descriptor reads its exact adjacent physical step -// separately through stride[-(trailing + 1)]. -inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { - TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) - << name << " cannot fold leading dimensions above " << trailing - << " trailing dims from ndim=" << t.ndim(); - int outer_last = t.ndim() - trailing - 1; - if (outer_last <= 0) { - return; - } - int64_t step = t.stride(outer_last); - TVM_FFI_CHECK(step > 0, ValueError) - << name << " physical strides must be positive"; - int64_t expected = step; - for (int axis = outer_last - 1; axis >= 0; --axis) { - expected *= t.size(axis + 1); - if (t.size(axis) > 1) { - TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) - << name << " leading dims are not physically foldable above " << trailing - << " trailing dims: stride(" << axis << ")=" << t.stride(axis) - << ", expected " << expected; - } - } -} - -// 4D TMA descriptor for buffer 'tmap_q' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 3, ValueError) - << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - int64_t d2 = t.size(t.ndim() - 2); - int64_t d3 = t.size(t.ndim() - 3); - TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) - << "TMA source 'tmap_q' trailing dims must be positive"; - TVM_FFI_CHECK(d1 % 64 == 0, ValueError) - << "TMA source 'tmap_q' extent " << d1 - << " must divide exactly by " << 64; - uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(d3)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) - << "TMA box (64, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; - int64_t carrier_stride_0 = d1; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; - int64_t carrier_stride_1 = 64; - TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; - TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; - TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; - int64_t carrier_stride_2 = (d2 * d1); - TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; - TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; - TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; - uint64_t global_strides[3] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - (uint64_t)((carrier_stride_1 * 16) / 8), - (uint64_t)((carrier_stride_2 * 16) / 8), - }; - uint32_t box_dim[4] = {64u, 64u, 1u, 1u}; - uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_swa_k' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_swa_k(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_swa_k' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_swa_k' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_swa_k' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_swa_k"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_swa_k' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_swa_k' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_k'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_swa_k') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_compressed_k' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_compressed_k(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_compressed_k' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_compressed_k' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_compressed_k' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_compressed_k"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_compressed_k' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_compressed_k' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_k'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_k') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_swa_v' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_swa_v(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_swa_v' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_swa_v' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_swa_v' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_swa_v"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_swa_v' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_swa_v' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_v'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_swa_v') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_compressed_v' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_compressed_v(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_compressed_v' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_compressed_v' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_compressed_v' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_compressed_v"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_compressed_v' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_compressed_v' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_v'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_v') failed: CUresult=" << (int)r; - return tm; -} - -void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_compressed_k, TensorView arg_tmap_swa_v, TensorView arg_tmap_compressed_v, TensorView arg_O, TensorView arg_partial_lse, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_num_query_tokens, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, int64_t grid_x, int64_t grid_y, int64_t grid_z) { - DLDevice dev = arg_tmap_q.device(); - ScopedCudaDevice device_guard(dev.device_id); - CheckCudaTensor(arg_tmap_q, "tmap_q"); - CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); - CheckContiguous(arg_tmap_q, "tmap_q"); - CheckCudaTensor(arg_tmap_swa_k, "tmap_swa_k"); - CheckDtype(arg_tmap_swa_k, "tmap_swa_k", 4, 16, 1); - CheckCudaTensor(arg_tmap_compressed_k, "tmap_compressed_k"); - CheckDtype(arg_tmap_compressed_k, "tmap_compressed_k", 4, 16, 1); - CheckCudaTensor(arg_tmap_swa_v, "tmap_swa_v"); - CheckDtype(arg_tmap_swa_v, "tmap_swa_v", 4, 16, 1); - CheckCudaTensor(arg_tmap_compressed_v, "tmap_compressed_v"); - CheckDtype(arg_tmap_compressed_v, "tmap_compressed_v", 4, 16, 1); - CheckCudaTensor(arg_O, "O"); - CheckDtype(arg_O, "O", 4, 16, 1); - CheckContiguous(arg_O, "O"); - CheckCudaTensor(arg_partial_lse, "partial_lse"); - CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); - CheckContiguous(arg_partial_lse, "partial_lse"); - CheckCudaTensor(arg_swa_indices, "swa_indices"); - CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); - CheckContiguous(arg_swa_indices, "swa_indices"); - CheckCudaTensor(arg_compressed_indices, "compressed_indices"); - CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); - CheckContiguous(arg_compressed_indices, "compressed_indices"); - CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); - CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckCudaTensor(arg_sinks, "sinks"); - CheckDtype(arg_sinks, "sinks", 2, 32, 1); - CheckContiguous(arg_sinks, "sinks"); - CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); - CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); - CheckContiguous(arg_bmm1_scale, "bmm1_scale"); - CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); - CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); - CheckContiguous(arg_bmm2_scale, "bmm2_scale"); - TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) - << "scalar 'num_heads' value " << arg_num_heads - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_num_query_tokens >= -2147483648LL && arg_num_query_tokens <= 2147483647LL, ValueError) - << "scalar 'num_query_tokens' value " << arg_num_query_tokens - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) - << "scalar 'swa_index_stride' value " << arg_swa_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) - << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) - << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) - << "scalar 'sparse_topk' value " << arg_sparse_topk - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) - << "scalar 'has_sinks' value " << arg_has_sinks - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) - << "scalar 'total_work_items' value " << arg_total_work_items - << " is outside i32 range [-2147483648, 2147483647]"; - CheckSameCudaDevice(arg_tmap_swa_k, arg_tmap_q, "tmap_swa_k", "tmap_q"); - CheckSameCudaDevice(arg_tmap_compressed_k, arg_tmap_q, "tmap_compressed_k", "tmap_q"); - CheckSameCudaDevice(arg_tmap_swa_v, arg_tmap_q, "tmap_swa_v", "tmap_q"); - CheckSameCudaDevice(arg_tmap_compressed_v, arg_tmap_q, "tmap_compressed_v", "tmap_q"); - CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); - CheckSameCudaDevice(arg_partial_lse, arg_tmap_q, "partial_lse", "tmap_q"); - CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); - CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); - CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); - CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); - CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); - CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); - TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) - << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y - << ", " << grid_z << ")"; - TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) - << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z - << ") must be divisible by cluster dims (2, 1, 1)"; - - cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); - CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); - CUtensorMap p_tmap_swa_k = EncodeTma_tmap_swa_k(arg_tmap_swa_k); - CUtensorMap p_tmap_compressed_k = EncodeTma_tmap_compressed_k(arg_tmap_compressed_k); - CUtensorMap p_tmap_swa_v = EncodeTma_tmap_swa_v(arg_tmap_swa_v); - CUtensorMap p_tmap_compressed_v = EncodeTma_tmap_compressed_v(arg_tmap_compressed_v); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); - float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); - int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); - int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); - int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); - float* p_sinks = static_cast(arg_sinks.data_ptr()); - float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); - float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); - int32_t v_num_heads = (int32_t)arg_num_heads; - int32_t v_num_query_tokens = (int32_t)arg_num_query_tokens; - int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; - int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; - int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; - int32_t v_sparse_topk = (int32_t)arg_sparse_topk; - int32_t v_has_sinks = (int32_t)arg_has_sinks; - int32_t v_total_work_items = (int32_t)arg_total_work_items; - void* kargs[] = {&p_tmap_q, &p_tmap_swa_k, &p_tmap_compressed_k, &p_tmap_swa_v, &p_tmap_compressed_v, &p_O, &p_partial_lse, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_num_query_tokens, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_has_sinks, &v_total_work_items}; - - static std::mutex smem_mu; - static auto* smem_status_by_device = - new std::unordered_map(); - cudaError_t smem_status = cudaSuccess; - { - std::lock_guard lock(smem_mu); - auto it = smem_status_by_device->find(dev.device_id); - if (it == smem_status_by_device->end()) { - int current_device = -1; - smem_status = cudaGetDevice(¤t_device); - if (smem_status == cudaSuccess && current_device != dev.device_id) { - smem_status = cudaErrorInvalidDevice; - } - if (smem_status == cudaSuccess) { - smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_125332c2066d89bd7610), - cudaFuncAttributeMaxDynamicSharedMemorySize, - 231936); - } - smem_status_by_device->emplace(dev.device_id, smem_status); - } else { - smem_status = it->second; - } - } - TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_125332c2066d89bd7610 failed: " - << cudaGetErrorString(smem_status); - dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); - dim3 block(512u, 1u, 1u); - - // Extended direct-source launch for a linked CUDA kernel symbol. - cudaLaunchAttribute attrs[2]{}; - int n = 0; - attrs[n].id = cudaLaunchAttributeClusterDimension; - attrs[n].val.clusterDim.x = 2u; - attrs[n].val.clusterDim.y = 1u; - attrs[n].val.clusterDim.z = 1u; - ++n; - attrs[n].id = cudaLaunchAttributeClusterSchedulingPolicyPreference; - attrs[n].val.clusterSchedulingPolicyPreference = cudaClusterSchedulingPolicySpread; - ++n; - cudaLaunchConfig_t config{}; - config.gridDim = grid; - config.blockDim = block; - config.dynamicSmemBytes = 231936u; - config.stream = stream; - config.attrs = attrs; - config.numAttrs = n; - cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_125332c2066d89bd7610), kargs); - TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_125332c2066d89bd7610 failed: " - << cudaGetErrorString(launch_status); -} - -} // namespace cake_host_shim_0f71e276da509b5c - -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_0f71e276da509b5c::Run); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_125332c2066d89bd7610_kernel.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_125332c2066d89bd7610_kernel.cu deleted file mode 100644 index d4c22d86df8..00000000000 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_125332c2066d89bd7610_kernel.cu +++ /dev/null @@ -1,2407 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -typedef signed char int8_t; -typedef unsigned char uint8_t; -typedef unsigned short uint16_t; -typedef unsigned int uint32_t; -#if defined(__CUDACC_RTC__) -typedef unsigned long long uint64_t; -#else -typedef unsigned long uint64_t; -#endif -static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); -typedef signed int int32_t; -typedef short int int16_t; -struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; -struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; -static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); -static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); - -#if defined(__CUDACC_RTC__) -typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; -#else -#include -#endif - -static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); -static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); -#include -#include - -__device__ __forceinline__ int make_warp_uniform(int x) { - int result; - asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" - : "=r"(result) : "r"(x)); - return result; -} - -#define CAKE_INF CUDART_INF_F -#define TMEM_NCOLS 512 -#define TMEM_TMEM_SCRATCH_OFFSET 0 -#define NUM_K_PIPE_STAGES 8 -#define NUM_V_PIPE_STAGES 2 -#define SMEM_SMEM_Q_OFF 0 -#define SMEM_SMEM_Q_STAGE_BYTES 8192 -#define SMEM_SMEM_Q_STRIDE 8192 -#define SMEM_SMEM_K_OFF 65536 -#define SMEM_SMEM_K_STAGE_BYTES 16384 -#define SMEM_SMEM_K_STRIDE 16384 -#define SMEM_SMEM_V_OFF 65536 -#define SMEM_SMEM_V_STAGE_BYTES 16384 -#define SMEM_SMEM_V_STRIDE 16384 -#define SMEM_SMEM_P_OFF 196608 -#define SMEM_SMEM_P_STAGE_BYTES 8192 -#define SMEM_SMEM_P_STRIDE 8192 -#define SMEM_SMEM_INDICES_OFF 229376 -#define SMEM_SMEM_INDICES_STAGE_BYTES 1024 -#define SMEM_SMEM_INDICES_STRIDE 1024 -#define SMEM_SMEM_SOFTMAX_EXCHANGE_OFF 230400 -#define SMEM_SMEM_SOFTMAX_EXCHANGE_STAGE_BYTES 512 -#define SMEM_SMEM_SOFTMAX_EXCHANGE_STRIDE 512 -#define SMEM_SMEM_EPILOGUE_EXCHANGE_OFF 230912 -#define SMEM_SMEM_EPILOGUE_EXCHANGE_STAGE_BYTES 512 -#define SMEM_SMEM_EPILOGUE_EXCHANGE_STRIDE 512 -#define SMEM_SMEM_INDEX_VALID_OFF 231424 -#define SMEM_SMEM_INDEX_VALID_STAGE_BYTES 64 -#define SMEM_SMEM_INDEX_VALID_STRIDE 64 -#define SMEM_TOTAL 231936 -#define THREADS 512 - -#include - -__device__ __forceinline__ uint32_t elect_sync() { - uint32_t pred = 0; - asm volatile( - "{\n\t" - ".reg .pred %%px;\n\t" - "elect.sync _|%%px, %1;\n\t" - "@%%px mov.s32 %0, 1;\n\t" - "}\n" - : "+r"(pred) - : "r"(0xFFFFFFFF)); - return pred; -} - - -__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { - asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" - :: "r"(mbar_addr), "r"(count) : "memory"); -} - -__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { - asm volatile("mbarrier.init.b64 [%0], %1;" - :: "l"(mbar_addr), "r"(count) : "memory"); -} - - -__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - - -// CTA-local pipelines have short, resident producer/consumer edges. Omitting -// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the -// explicit loop still makes this helper blocking until acquire succeeds. -__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE;\n\t" - "bra.uni LAB_WAIT;\n\t" - "DONE:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Source-faithful relaxed CTA wait used only by a typed protocol that does -// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. -__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, 10000000;\n\t" - "@P1 bra.uni DONE_RELAXED;\n\t" - "bra.uni LAB_WAIT_RELAXED;\n\t" - "DONE_RELAXED:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Exact source ports may request the PTX suspendTimeHint operand explicitly. -// The hint is expressed in nanoseconds and is kept separate from the canonical -// no-hint CTA helper so unrelated schedules retain their existing retry path. -__device__ __forceinline__ void mbarrier_wait_suspend( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_SUSPEND:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_SUSPEND;\n\t" - "bra.uni LAB_WAIT_SUSPEND;\n\t" - "DONE_SUSPEND:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE_CLUSTER;\n\t" - "bra.uni LAB_WAIT_CLUSTER;\n\t" - "DONE_CLUSTER:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - ".reg .u32 WAIT_ADDR;\n\t" - "mov.u32 WAIT_ADDR, %0;\n\t" - "LAB_WAIT_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [WAIT_ADDR], %1, %2;\n\t" - "@P1 bra.uni DONE_HINT;\n\t" - "bra.uni LAB_WAIT_HINT;\n\t" - "DONE_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -// Exact unqualified CTA wait used by source schedules whose PTX intentionally -// omits the acquire qualifier while retaining a typed suspendTimeHint operand. -__device__ __forceinline__ void mbarrier_wait_relaxed_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED_HINT:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra DONE_RELAXED_HINT;\n\t" - "bra LAB_WAIT_RELAXED_HINT;\n\t" - "DONE_RELAXED_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); -} - -__device__ __forceinline__ void mbarrier_wait_cluster_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" - "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" - "DONE_CLUSTER_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_suspend( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait_cluster(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); - } -} - - -__device__ __forceinline__ void tcgen05_mma_f16_cta2( - int taddr, uint64_t a_desc, uint64_t b_desc, - uint32_t i_desc, int enable_input_d) { - asm volatile( - "{\n\t" - ".reg .pred p;\n\t" - ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - "setp.ne.b32 p, %4, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" - "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "tcgen05.mma.cta_group::2.kind::f16 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" - "}\n" - :: "r"(taddr), "l"(a_desc), "l"(b_desc), - "r"(i_desc), "r"(enable_input_d) - : "memory"); -} - - -__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { - return (x & 0x3FFFFULL) >> 4ULL; -} - - -union MmaSmemDesc { - uint64_t u64; - uint32_t u32[2]; -}; - -__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { - MmaSmemDesc tmp; - tmp.u64 = smem_desc; - tmp.u32[0] += offset; - smem_desc = tmp.u64; -} - - -__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { - asm volatile( - "{\n\t" - ".reg .pred leader;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" - ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" - "}\n" - :: "r"(mbar_addr), "h"(cta_mask) : "memory"); -} - - -__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { - asm volatile( - "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" - :: "r"(mbar_addr), "r"(bytes) : "memory"); -} - - -__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7," - " %8, %9, %10, %11, %12, %13, %14, %15," - " %16, %17, %18, %19, %20, %21, %22, %23," - " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), - "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), - "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), - "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), - "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), - "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), - "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x16.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7," - " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), - "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), - "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ void tmem_ld_x4(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x4.b32" - " {%0, %1, %2, %3}, [%4];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { - asm volatile( - "tcgen05.st.sync.aligned.32x32b.x16.b32" - " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," - " %9, %10, %11, %12, %13, %14, %15, %16};" - :: "r"(tmem_addr), - "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), - "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), - "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), - "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); -} - - -__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { - asm volatile( - "tcgen05.st.sync.aligned.32x32b.x4.b32" - " [%0], {%1, %2, %3, %4};" - :: "r"(tmem_addr), - "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); -} - - -__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { - uint32_t addr; - asm("{\n\t" - ".reg .u64 u64addr;\n\t" - "cvta.to.shared.u64 u64addr, %1;\n\t" - "cvt.u32.u64 %0, u64addr;\n\t" - "}\n" : "=r"(addr) : "l"(ptr)); - return addr; -} - - -__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { - uint32_t remote; - asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" - : "=r"(remote) : "r"(local_addr), "r"(rank)); - return remote; -} - - -__device__ __forceinline__ float approx_exp2(float x) { - float y; - asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float approx_rcp(float x) { - float y; - asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float max_noftz(float a, float b) { - float c; - asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); - return c; -} - - -__device__ __forceinline__ float warp_reduce_max(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); - return val; -} - - -__device__ __forceinline__ float warp_reduce_sum(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val += __shfl_xor_sync(0xFFFFFFFF, val, offset); - return val; -} - - -__device__ __forceinline__ float row_max_reduce(float2 acc) { - return max_noftz(acc.x, acc.y); -} - - -__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { - #pragma unroll - for (int j = 0; j < 16; j++) { - if (j % 2 == 0) - acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); - else - acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); - } -} - - -__device__ __forceinline__ float2 ex2_emulation_f32x2_value(float2 value) { - const float c0 = 1.0f, c1 = 0.695146143436431884765625f; - const float c2 = 0.227564394474029541015625f, c3 = 0.077119089663028717041015625f; - const float magic = 12582912.0f; - float x0 = max_noftz(value.x, -127.0f), x1 = max_noftz(value.y, -127.0f); - float2 xc2 = make_float2(x0, x1), magic2 = make_float2(magic, magic); - float2 xr2; - asm("add.rm.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xr2) - : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&magic2)); - float2 c3_2 = make_float2(c3, c3), c2_2 = make_float2(c2, c2); - float2 c1_2 = make_float2(c1, c1), c0_2 = make_float2(c0, c0); - float2 xrb2, xfrac2; - asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xrb2) - : "l"(*(unsigned long long*)&xr2), "l"(*(unsigned long long*)&magic2)); - asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xfrac2) - : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&xrb2)); - float2 poly2; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&c3_2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c2_2)); - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c1_2)); - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c0_2)); - int x0r_i, x1r_i, p0_i, p1_i; - asm("mov.b64 {%0, %1}, %2;" : "=r"(x0r_i), "=r"(x1r_i) : "l"(*(unsigned long long*)&xr2)); - asm("mov.b64 {%0, %1}, %2;" : "=r"(p0_i), "=r"(p1_i) : "l"(*(unsigned long long*)&poly2)); - float r0, r1; - asm("mov.b32 %0, %1;" : "=f"(r0) : "r"((x0r_i << 23) + p0_i)); - asm("mov.b32 %0, %1;" : "=f"(r1) : "r"((x1r_i << 23) + p1_i)); - return make_float2(r0, r1); -} - -__device__ __forceinline__ void ex2_emulation_f32x2(float* x0_ptr, float* x1_ptr) { - float2 result = ex2_emulation_f32x2_value(make_float2(*x0_ptr, *x1_ptr)); - *x0_ptr = result.x; *x1_ptr = result.y; -} - -__device__ __forceinline__ void softmax_frag_exp2_cast( - float* sv, uint32_t* pv, int use_emu) -{ - #pragma unroll - for (int j = 0; j < 16; j++) { - if (use_emu && j >= 12) - ex2_emulation_f32x2(&sv[j*2], &sv[j*2+1]); - else { - sv[j*2] = approx_exp2(sv[j*2]); - sv[j*2+1] = approx_exp2(sv[j*2+1]); - } - } - #pragma unroll - for (int j = 0; j < 16; j++) { - __nv_bfloat162 bf = __float22bfloat162_rn({sv[j*2], sv[j*2+1]}); - pv[j] = reinterpret_cast(bf); - } -} - - - -__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { - const float2* sv2 = reinterpret_cast(sv); - #pragma unroll - for (int j = 0; j < 16; j++) { - asm("add.f32x2 %0, %1, %2;" - : "+l"(reinterpret_cast(*acc)) - : "l"(reinterpret_cast(*acc)), - "l"(reinterpret_cast(sv2[j]))); - } -} - - -__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { - asm("mul.rn.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { - asm("mul.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { - asm("add.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { - asm("add.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { - asm("sub.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { - asm("sub.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("add.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { - float2 r; - asm("sub.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("sub.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ void fma_scale_x32( - float* sv, const float2* scale2, const float2* neg_max2) -{ - float2* sv_2 = reinterpret_cast(sv); - #pragma unroll - for (int j = 0; j < 16; j++) - fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); -} - -__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("mul.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - - -// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) - -__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("add.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("add.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("add.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("add.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("add.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("add.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("add.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - - -__device__ __forceinline__ void fence_async_shared() { - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); -} - - -__device__ __forceinline__ uint64_t make_smem_desc(int addr) { - const int SBO = 1024; - return desc_encode(addr) - | (desc_encode(SBO) << 32ULL) - | (1ULL << 46ULL) - | (2ULL << 61ULL); -} - - -__device__ __forceinline__ void tma_2d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3}], [%4];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_4d_gmem2smem_cta2( - int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.4d.shared::cluster.global" - ".mbarrier::complete_tx::bytes.cta_group::2" - " [%0], [%1, {%2, %3, %4, %5}], [%6];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr) { - // Canonical .shared::cta form for non-multicast gather4, matching - // trtllm-gen / cuda_ptx and the PTX ISA qualifier order - // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, - // .shared::cluster is reserved for the multicast variant (ctaMask). - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem_cta2( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr) { - // Canonical .shared::cta form; see tma_gather4_gmem2smem above. - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" - ".mbarrier::complete_tx::bytes.cta_group::2" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem_mc( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr, unsigned short cta_mask) { - // Multicast variant: the PTX grammar ties the .shared::cluster - // destination to .multicast::cluster + ctaMask (cf. cuda_ptx / - // SM100_TMA_LOAD_MULTICAST_2D_GATHER4 in CUTLASS). - asm volatile( - "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" - ".mbarrier::complete_tx::bytes.multicast::cluster" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr), "h"(cta_mask) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem_mc_cta2( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr, unsigned short cta_mask) { - // Multicast + cta_group::2 variant; see tma_gather4_gmem2smem_mc. - asm volatile( - "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" - ".mbarrier::complete_tx::bytes.multicast::cluster.cta_group::2" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr), "h"(cta_mask) : "memory"); -} - - -__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { - asm volatile( - "{\n\t" - ".reg .b16 lo, hi;\n\t" - "mov.b32 {lo, hi}, %1;\n\t" - "tcgen05.commit.cta_group::2.mbarrier::arrive::one" - ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" - "}\n" - :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); -} - -extern "C" { - -__global__ __launch_bounds__(512, 1) __cluster_dims__(2,1,1) void -kernel_cake_dsv4_125332c2066d89bd7610(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items) -{ - const int tid = threadIdx.x; - const int warp = make_warp_uniform(tid / 32); - const int lane = tid % 32; - - extern __shared__ __align__(1024) char smem_raw[]; - int smem; - smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); - - const int mbar_base = smem + 231488; - #define q_full_addr (mbar_base + 0) - #define q_empty_addr (mbar_base + 8) - #define k_full_addr (mbar_base + 16) - #define k_empty_addr (mbar_base + 80) - #define slot_empty_addr (mbar_base + 144) - #define v_full_addr (mbar_base + 208) - #define v_empty_addr (mbar_base + 224) - #define s_full_addr (mbar_base + 240) - #define s_empty_addr (mbar_base + 256) - #define p_full_addr (mbar_base + 272) - #define p_empty_addr (mbar_base + 288) - #define o_empty_addr (mbar_base + 304) - #define stats_addr (mbar_base + 312) - #define stats_empty_addr (mbar_base + 328) - #define o_full_addr (mbar_base + 344) - #define o_first_slice_seeded_addr (mbar_base + 352) - #define tmem_dealloc_peer_addr (mbar_base + 360) - #define index_valid_full_addr (mbar_base + 368) - #define index_valid_empty_addr (mbar_base + 384) - - const int bid = blockIdx.x; - const int num_bids = gridDim.x; - const unsigned int clusters_x = gridDim.x / 2; - const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; - const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; - - int cta_rank; - asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); - - // Kernel setup ops - __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 0); - const int smem_q_addr = smem + 0; - __nv_bfloat16* smem_k = reinterpret_cast<__nv_bfloat16*>(smem_raw + 65536); - const int smem_k_addr = smem + 65536; - __nv_bfloat16* smem_v = reinterpret_cast<__nv_bfloat16*>(smem_raw + 65536); - const int smem_v_addr = smem + 65536; - __nv_bfloat16* smem_p = reinterpret_cast<__nv_bfloat16*>(smem_raw + 196608); - const int smem_p_addr = smem + 196608; - int* smem_indices = reinterpret_cast(smem_raw + 229376); - const int smem_indices_addr = smem + 229376; - float* smem_softmax_exchange = reinterpret_cast(smem_raw + 230400); - const int smem_softmax_exchange_addr = smem + 230400; - float* smem_epilogue_exchange = reinterpret_cast(smem_raw + 230912); - const int smem_epilogue_exchange_addr = smem + 230912; - unsigned int* smem_index_valid = reinterpret_cast(smem_raw + 231424); - const int smem_index_valid_addr = smem + 231424; - asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); - asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); - - // Mbarrier init (19 pipeline groups, 0 ordered-sequence groups, 50 barriers) - // Mbarriers at smem_raw[231488..231888) - - if (warp == 0) { - uint32_t leader = elect_sync(); - if (leader) { - // q_full: 1 barriers, init_count=1 - mbarrier_init(smem + 231488, 1); - // q_empty: 1 barriers, init_count=1 - mbarrier_init(smem + 231496, 1); - // --- pipeline 'k_pipe' --- - // k_full: 8 barriers, init_count=1 - mbarrier_init(smem + 231504, 1); - mbarrier_init(smem + 231512, 1); - mbarrier_init(smem + 231520, 1); - mbarrier_init(smem + 231528, 1); - mbarrier_init(smem + 231536, 1); - mbarrier_init(smem + 231544, 1); - mbarrier_init(smem + 231552, 1); - mbarrier_init(smem + 231560, 1); - // k_empty: 8 barriers, init_count=1 - mbarrier_init(smem + 231568, 1); - mbarrier_init(smem + 231576, 1); - mbarrier_init(smem + 231584, 1); - mbarrier_init(smem + 231592, 1); - mbarrier_init(smem + 231600, 1); - mbarrier_init(smem + 231608, 1); - mbarrier_init(smem + 231616, 1); - mbarrier_init(smem + 231624, 1); - // slot_empty: 8 barriers, init_count=1 - mbarrier_init(smem + 231632, 1); - mbarrier_init(smem + 231640, 1); - mbarrier_init(smem + 231648, 1); - mbarrier_init(smem + 231656, 1); - mbarrier_init(smem + 231664, 1); - mbarrier_init(smem + 231672, 1); - mbarrier_init(smem + 231680, 1); - mbarrier_init(smem + 231688, 1); - // --- pipeline 'v_pipe' --- - // v_full: 2 barriers, init_count=1 - mbarrier_init(smem + 231696, 1); - mbarrier_init(smem + 231704, 1); - // v_empty: 2 barriers, init_count=1 - mbarrier_init(smem + 231712, 1); - mbarrier_init(smem + 231720, 1); - // s_full: 2 barriers, init_count=1 - mbarrier_init(smem + 231728, 1); - mbarrier_init(smem + 231736, 1); - // s_empty: 2 barriers, init_count=256 - mbarrier_init(smem + 231744, 256); - mbarrier_init(smem + 231752, 256); - // p_full: 2 barriers, init_count=256 - mbarrier_init(smem + 231760, 256); - mbarrier_init(smem + 231768, 256); - // p_empty: 2 barriers, init_count=1 - mbarrier_init(smem + 231776, 1); - mbarrier_init(smem + 231784, 1); - // o_empty: 1 barriers, init_count=256 - mbarrier_init(smem + 231792, 256); - // stats: 2 barriers, init_count=128 - mbarrier_init(smem + 231800, 128); - mbarrier_init(smem + 231808, 128); - // stats_empty: 2 barriers, init_count=128 - mbarrier_init(smem + 231816, 128); - mbarrier_init(smem + 231824, 128); - // o_full: 1 barriers, init_count=1 - mbarrier_init(smem + 231832, 1); - // o_first_slice_seeded: 1 barriers, init_count=256 - mbarrier_init(smem + 231840, 256); - // tmem_dealloc_peer: 1 barriers, init_count=32 - mbarrier_init(smem + 231848, 32); - // index_valid_full: 2 barriers, init_count=32 - mbarrier_init(smem + 231856, 32); - mbarrier_init(smem + 231864, 32); - // index_valid_empty: 2 barriers, init_count=128 - mbarrier_init(smem + 231872, 128); - mbarrier_init(smem + 231880, 128); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - - asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); - asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); - - // TMEM alloc (512 columns, 512 used) - volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 231888); - if (warp == 0) { - int _tmem_hold = smem + 231888; - asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); - __syncwarp(); - asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); - } - - __syncthreads(); - asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); - asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); - - const int taddr = tmem_addr_storage[0]; - - // Kernel post-init ops - const int tmem_tmem_scratch = taddr; - - // ---- Ordered hardware-WG register redistribution ---- - // Dec phase frees registers before any WG attempts inc. - if (warp >= 8 && warp <= 15) { - asm volatile("setmaxnreg.dec.sync.aligned.u32 56;"); - } - - // ---- Role: softmax_wg ---- - if (warp <= 3) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); - { // softmax_wg_main - const int softmax_dummy = 0; - float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; - const int head_row_base = warp % 2 * 32; - const int score_row_base = warp % 4 * 32; - const int n_half = warp % 4 / 2; - const int my_row = head_row_base + lane; - const int exchange_idx = n_half * 64 + my_row; - int tile_cursor = 0; - int valid_iter = 0; - float seed_zero[4]; - #pragma unroll - for (int seed_i = 0; seed_i < 4; seed_i++) { - seed_zero[seed_i] = 0.0f; - } - #pragma unroll - for (int seed_stage = 0; seed_stage < 2; seed_stage++) { - #pragma unroll - for (int seed_col = 128 + seed_stage * 128; seed_col < 128 + (seed_stage + 1) * 128; seed_col += 4) { - int seed_addr = taddr + (unsigned int)seed_col + (unsigned int)(score_row_base << 16); - tmem_st_x4_f32(seed_addr, seed_zero); - } - } - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(o_first_slice_seeded_addr); - int peer_rank = cta_rank ^ 1; - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(o_first_slice_seeded_addr), "r"(peer_rank) : "memory"); - #pragma unroll 1 - for (unsigned int work_idx = cluster_id; work_idx < total_work_items; work_idx += num_clusters) { - int split_idx = work_idx % (unsigned int)(((1) ? 5 : 1)); - int query_idx = work_idx / (unsigned int)(((1) ? 5 : 1)); - int _max_0 = ((sparse_topk_lens[query_idx] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx] + sparse_topk_lens_offset) : (0)); - int _min_2 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); - int active_topk = _min_2; - int owner_capacity = ((sparse_topk + 128 - 1) / 128 + ((1) ? 5 : 1) - 1) / ((1) ? 5 : 1); - int first_owner_tile = work_idx % (unsigned int)(((1) ? 5 : 1)) * (unsigned int)owner_capacity; - int _min_3 = ((owner_capacity) < ((sparse_topk + 128 - 1) / 128 - first_owner_tile) ? (owner_capacity) : ((sparse_topk + 128 - 1) / 128 - first_owner_tile)); - int num_kv_tiles = _min_3; - int first_kv_tile = first_owner_tile; - int head_idx = cta_rank * 64 + my_row; - float row_max_val = -CAKE_INF; - float row_sum_val = 0.0f; - if (has_sinks != 0 && head_idx < num_heads && split_idx == 0) { - row_max_val = sinks[head_idx] * 1.4426950408889634f / softmax_scale_log2; - if (n_half == 0) { - row_sum_val = 1.0f; - } - } - int valid_slot = valid_iter & 1; - int valid_full_phase = valid_iter >> 1 & 1; - mbarrier_wait(index_valid_full_addr + (valid_slot) * 8, valid_full_phase); - int valid_word_base = valid_slot * 8 + n_half * 2; - #pragma unroll 1 - for (int tile = 0; tile < num_kv_tiles; tile++) { - int pipeline_tile = tile_cursor + tile; - int score_phase = pipeline_tile & 1; - int s_full_phase = pipeline_tile >> 1 & 1; - mbarrier_wait(s_full_addr + (score_phase) * 8, s_full_phase); - int score_offset = ((score_phase != 0) ? 64 : 0); - int score_base = taddr + (unsigned int)score_offset + (unsigned int)(score_row_base << 16); - float _tmem_load_0[64]; - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]) - : "r"(score_base)); - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]) - : "r"(score_base + 32)); - int valid_cols = active_topk - (first_kv_tile + tile) * 128 - n_half * 64; - if (valid_cols < 0) { - valid_cols = 0; - } - if (valid_cols > 64) { - valid_cols = 64; - } - unsigned int keep_lo = smem_index_valid[valid_word_base + tile * 4]; - unsigned int keep_hi = smem_index_valid[valid_word_base + tile * 4 + 1]; - uint32_t _thresh_mask_0; - { - int _lim_0 = valid_cols; - if (_lim_0 <= 0) { _thresh_mask_0 = 0u; } - else if (_lim_0 >= 32) { _thresh_mask_0 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_thresh_mask_0) : "r"(_lim_0)); - } - } - uint32_t _thresh_mask_1; - { - int _lim_1 = valid_cols - 32; - if (_lim_1 <= 0) { _thresh_mask_1 = 0u; } - else if (_lim_1 >= 32) { _thresh_mask_1 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_thresh_mask_1) : "r"(_lim_1)); - } - } - if (!(_thresh_mask_0 & keep_lo & (1u << 0))) _tmem_load_0[0] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 1))) _tmem_load_0[1] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 2))) _tmem_load_0[2] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 3))) _tmem_load_0[3] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 4))) _tmem_load_0[4] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 5))) _tmem_load_0[5] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 6))) _tmem_load_0[6] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 7))) _tmem_load_0[7] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 8))) _tmem_load_0[8] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 9))) _tmem_load_0[9] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 10))) _tmem_load_0[10] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 11))) _tmem_load_0[11] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 12))) _tmem_load_0[12] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 13))) _tmem_load_0[13] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 14))) _tmem_load_0[14] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 15))) _tmem_load_0[15] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 16))) _tmem_load_0[16] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 17))) _tmem_load_0[17] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 18))) _tmem_load_0[18] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 19))) _tmem_load_0[19] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 20))) _tmem_load_0[20] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 21))) _tmem_load_0[21] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 22))) _tmem_load_0[22] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 23))) _tmem_load_0[23] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 24))) _tmem_load_0[24] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 25))) _tmem_load_0[25] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 26))) _tmem_load_0[26] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 27))) _tmem_load_0[27] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 28))) _tmem_load_0[28] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 29))) _tmem_load_0[29] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 30))) _tmem_load_0[30] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 31))) _tmem_load_0[31] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 0))) _tmem_load_0[32] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 1))) _tmem_load_0[33] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 2))) _tmem_load_0[34] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 3))) _tmem_load_0[35] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 4))) _tmem_load_0[36] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 5))) _tmem_load_0[37] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 6))) _tmem_load_0[38] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 7))) _tmem_load_0[39] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 8))) _tmem_load_0[40] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 9))) _tmem_load_0[41] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 10))) _tmem_load_0[42] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 11))) _tmem_load_0[43] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 12))) _tmem_load_0[44] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 13))) _tmem_load_0[45] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 14))) _tmem_load_0[46] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 15))) _tmem_load_0[47] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 16))) _tmem_load_0[48] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 17))) _tmem_load_0[49] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 18))) _tmem_load_0[50] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 19))) _tmem_load_0[51] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 20))) _tmem_load_0[52] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 21))) _tmem_load_0[53] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 22))) _tmem_load_0[54] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 23))) _tmem_load_0[55] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 24))) _tmem_load_0[56] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 25))) _tmem_load_0[57] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 26))) _tmem_load_0[58] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 27))) _tmem_load_0[59] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 28))) _tmem_load_0[60] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 29))) _tmem_load_0[61] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 30))) _tmem_load_0[62] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 31))) _tmem_load_0[63] = -CAKE_INF; - float2 _reg_reduce_max2_2 = {-CAKE_INF, -CAKE_INF}; - row_max_x32_accum(&_tmem_load_0[0], _reg_reduce_max2_2); - row_max_x32_accum(&_tmem_load_0[32], _reg_reduce_max2_2); - float _tmem_load_0_max = row_max_reduce(_reg_reduce_max2_2); - float _max_1 = max_noftz(_tmem_load_0_max, row_max_val); - float new_max = _max_1; - smem_softmax_exchange[exchange_idx] = new_max; - asm volatile("barrier.sync 3, 128;" ::: "memory"); - float _max_2 = max_noftz(new_max, smem_softmax_exchange[exchange_idx ^ 64]); - new_max = _max_2; - asm volatile("barrier.sync 3, 128;" ::: "memory"); - float no_correction = (((new_max - row_max_val) * softmax_scale_log2 <= 0.0f) ? 1.0f : 0.0f); - float delta = softmax_scale_log2 * (row_max_val - new_max); - float _exp2_0 = approx_exp2(delta); - float exp_delta = _exp2_0; - float acc_scale = ((row_max_val > -CAKE_INF) ? exp_delta : 1.0f); - row_max_val = new_max; - float safe_max = ((new_max == -CAKE_INF) ? 0.0f : new_max); - float max_scaled = safe_max * softmax_scale_log2; - const float2 _fma_b2_3 = {softmax_scale_log2, softmax_scale_log2}; - const float2 _fma_c2_4 = {-max_scaled, -max_scaled}; - #pragma unroll - for (int _lf = 0; _lf < 32; _lf++) - fma_f32x2_inplace(&reinterpret_cast(_tmem_load_0)[_lf], _fma_b2_3, _fma_c2_4); - #pragma unroll - for (int _le = 0; _le < 64; _le++) { - _tmem_load_0[_le] = approx_exp2(_tmem_load_0[_le]); - } - float2 _reg_reduce_sum2_5 = make_float2(0.0f, 0.0f); - softmax_block_sum(&_tmem_load_0[0], &_reg_reduce_sum2_5); - softmax_block_sum(&_tmem_load_0[32], &_reg_reduce_sum2_5); - float _tmem_load_0_sum = _reg_reduce_sum2_5.x + _reg_reduce_sum2_5.y; - row_sum_val = row_sum_val * acc_scale + _tmem_load_0_sum; - int stats_empty_phase = pipeline_tile >> 1 & 1 ^ 1; - mbarrier_wait(stats_empty_addr + (score_phase) * 8, stats_empty_phase); - float meta[4]; - meta[0] = row_sum_val; - meta[1] = row_max_val; - meta[2] = acc_scale; - meta[3] = no_correction; - int meta_addr = taddr + 384 + (unsigned int)(score_phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(head_row_base << 16); - tmem_st_x4_f32(meta_addr, meta); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - mbarrier_arrive(stats_addr + (score_phase) * 8); - int p_stage = pipeline_tile % 2; - int p_empty_phase = pipeline_tile / 2 & 1 ^ 1; - mbarrier_wait(p_empty_addr + (p_stage) * 8, p_empty_phase); - uint32_t _tmem_load_0_bf16[32]; - #pragma unroll - for (int _lp = 0; _lp < 32; _lp++) { - __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_0[_lp*2 + 0], _tmem_load_0[_lp*2+1 + 0])); - _tmem_load_0_bf16[_lp] = *(uint32_t*)&_bf2; - } - int smem_p_stage = p_stage * 2 + n_half; - int p_base_smem = smem_p_addr + (unsigned int)(smem_p_stage * 8192); - #pragma unroll - for (int vec = 0; vec < 8; vec++) { - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_base_smem + (my_row * 128 + vec * 16 ^ (my_row * 128 + vec * 16 >> 7 & 7) << 4))), "r"(_tmem_load_0_bf16[vec * 4]), "r"(_tmem_load_0_bf16[vec * 4 + 1]), "r"(_tmem_load_0_bf16[vec * 4 + 2]), "r"(_tmem_load_0_bf16[vec * 4 + 3]) : "memory"); - } - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(p_full_addr + (unsigned int)(p_stage * 8)), "r"(0) : "memory"); - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(s_empty_addr + (unsigned int)(score_phase * 8)), "r"(0) : "memory"); - } - mbarrier_arrive(index_valid_empty_addr + (valid_slot) * 8); - valid_iter = valid_iter + 1; - tile_cursor = tile_cursor + num_kv_tiles; - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - } - } - // ---- Role: correction_wg ---- - if (warp >= 4 && warp <= 7) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); - { // correction_wg_main - const int correction_dummy = 0; - float softmax_scale_log2_1 = bmm1_scale[0] * 1.4426950408889634f; - float output_scale = bmm2_scale[0]; - const int head_row_base_1 = warp % 2 * 32; - const int output_row_base = warp % 4 * 32; - const int n_half_1 = warp % 4 / 2; - const int my_row_1 = head_row_base_1 + lane; - const int corr_row = head_row_base_1 << 16; - int tile_cursor_1 = 0; - unsigned int _phase_o_full_0 = 0; - #pragma unroll 1 - for (unsigned int work_idx_1 = cluster_id; work_idx_1 < total_work_items; work_idx_1 += num_clusters) { - int split_idx_1 = work_idx_1 % (unsigned int)(((1) ? 5 : 1)); - int query_idx_1 = work_idx_1 / (unsigned int)(((1) ? 5 : 1)); - int owner_capacity_1 = ((sparse_topk + 128 - 1) / 128 + ((1) ? 5 : 1) - 1) / ((1) ? 5 : 1); - int first_owner_tile_1 = work_idx_1 % (unsigned int)(((1) ? 5 : 1)) * (unsigned int)owner_capacity_1; - int _min_4 = ((owner_capacity_1) < ((sparse_topk + 128 - 1) / 128 - first_owner_tile_1) ? (owner_capacity_1) : ((sparse_topk + 128 - 1) / 128 - first_owner_tile_1)); - int num_kv_tiles_1 = _min_4; - float final_sum = 0.0f; - float final_max = -CAKE_INF; - #pragma unroll 1 - for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { - int pipeline_tile_1 = tile_cursor_1 + tile_1; - int score_phase_1 = pipeline_tile_1 & 1; - int stats_phase = pipeline_tile_1 >> 1 & 1; - mbarrier_wait(stats_addr + (score_phase_1) * 8, stats_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - int meta_addr_1 = taddr + 384 + (unsigned int)(score_phase_1 * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row; - float _tmem_load_1[4]; - tmem_ld_x4(&_tmem_load_1[0], meta_addr_1); - final_sum = _tmem_load_1[0]; - final_max = _tmem_load_1[1]; - float acc_scale_1 = _tmem_load_1[2]; - float no_correction_1 = _tmem_load_1[3]; - if (tile_1 > 0) { - mbarrier_wait(o_full_addr, _phase_o_full_0); - _phase_o_full_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - int _vote_0 = __all_sync(0xFFFFFFFF, no_correction_1 == 1.0f); - int skip_correction = _vote_0; - if (skip_correction == 0) { - #pragma unroll - for (int local_slice = 0; local_slice < 2; local_slice++) { - #pragma unroll - for (int acc_stage = 0; acc_stage < 2; acc_stage++) { - int o_base = taddr + 128 + (unsigned int)(acc_stage * 128) + (unsigned int)(local_slice * 64) + (unsigned int)(output_row_base << 16); - #pragma unroll - for (int c = 0; c < 64; c += 16) { - float _tmem_load_2[16]; - tmem_ld_x16(&_tmem_load_2[0], o_base + c); - const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; - #pragma unroll - for (int _ls = 0; _ls < 8; _ls++) - mul_f32x2_inplace(&reinterpret_cast(_tmem_load_2)[_ls], _scale2_0); - tmem_st_x16_f32(o_base + c, _tmem_load_2); - } - } - } - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - } - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(o_empty_addr), "r"(0) : "memory"); - } - mbarrier_arrive(stats_empty_addr + (score_phase_1) * 8); - } - mbarrier_wait(o_full_addr, _phase_o_full_0); - _phase_o_full_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - const int exchange_idx_1 = n_half_1 * 64 + my_row_1; - smem_epilogue_exchange[exchange_idx_1] = final_sum; - asm volatile("barrier.sync 4, 128;" ::: "memory"); - final_sum = final_sum + smem_epilogue_exchange[exchange_idx_1 ^ 64]; - asm volatile("barrier.sync 4, 128;" ::: "memory"); - float _rcp_0 = approx_rcp(final_sum); - float inv_sum = ((final_sum > 0.0f) ? _rcp_0 : 0.0f); - int head_idx_1 = cta_rank * 64 + my_row_1; - int output_offset = ((1) ? ((query_idx_1 * num_heads + head_idx_1) * ((1) ? 5 : 1) + split_idx_1) * 512 : (query_idx_1 * num_heads + head_idx_1) * 512); - if (n_half_1 == 0 && head_idx_1 < num_heads) { - int stat_offset = (query_idx_1 * num_heads + head_idx_1) * ((1) ? 5 : 1) + split_idx_1; - float _log2_0; - asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(final_sum)); - partial_lse[stat_offset] = ((final_sum > 0.0f) ? final_max * softmax_scale_log2_1 + _log2_0 : -CAKE_INF); - } - #pragma unroll - for (int acc_stage_1 = 0; acc_stage_1 < 2; acc_stage_1++) { - #pragma unroll - for (int local_slice_1 = 0; local_slice_1 < 2; local_slice_1++) { - int logical_slice = acc_stage_1 * 2 * 2 + n_half_1 * 2 + local_slice_1; - int o_base_1 = taddr + 128 + (unsigned int)(acc_stage_1 * 128) + (unsigned int)(local_slice_1 * 64) + (unsigned int)(output_row_base << 16); - #pragma unroll - for (int c_1 = 0; c_1 < 64; c_1 += 32) { - float _tmem_load_3[32]; - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]) - : "r"(o_base_1 + c_1)); - int gmem_base = output_offset + logical_slice * 64 + c_1; - #pragma unroll - for (int j = 0; j < 32; j += 16) { - { - const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; - #if __CUDA_ARCH__ >= 1000 - #pragma unroll - for (int _ps = 0; _ps < 8; _ps++) - mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_3[j])[_ps], _prescale2_1); - #else - #pragma unroll - for (int _ps = 0; _ps < 16; _ps++) - _tmem_load_3[j + _ps] *= inv_sum * output_scale; - #endif - { - __nv_bfloat162 _pk0 = __floats2bfloat162_rn(_tmem_load_3[j + 0], _tmem_load_3[j + 1]); - unsigned _pk_u0 = *reinterpret_cast(&_pk0); - __nv_bfloat162 _pk1 = __floats2bfloat162_rn(_tmem_load_3[j + 2], _tmem_load_3[j + 3]); - unsigned _pk_u1 = *reinterpret_cast(&_pk1); - __nv_bfloat162 _pk2 = __floats2bfloat162_rn(_tmem_load_3[j + 4], _tmem_load_3[j + 5]); - unsigned _pk_u2 = *reinterpret_cast(&_pk2); - __nv_bfloat162 _pk3 = __floats2bfloat162_rn(_tmem_load_3[j + 6], _tmem_load_3[j + 7]); - unsigned _pk_u3 = *reinterpret_cast(&_pk3); - __nv_bfloat162 _pk4 = __floats2bfloat162_rn(_tmem_load_3[j + 8], _tmem_load_3[j + 9]); - unsigned _pk_u4 = *reinterpret_cast(&_pk4); - __nv_bfloat162 _pk5 = __floats2bfloat162_rn(_tmem_load_3[j + 10], _tmem_load_3[j + 11]); - unsigned _pk_u5 = *reinterpret_cast(&_pk5); - __nv_bfloat162 _pk6 = __floats2bfloat162_rn(_tmem_load_3[j + 12], _tmem_load_3[j + 13]); - unsigned _pk_u6 = *reinterpret_cast(&_pk6); - __nv_bfloat162 _pk7 = __floats2bfloat162_rn(_tmem_load_3[j + 14], _tmem_load_3[j + 15]); - unsigned _pk_u7 = *reinterpret_cast(&_pk7); - asm volatile( - "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" - :: "l"((void*)(&((__nv_bfloat16*)(O + (gmem_base + j)))[0])), "r"(_pk_u0), "r"(_pk_u1), "r"(_pk_u2), "r"(_pk_u3), "r"(_pk_u4), "r"(_pk_u5), "r"(_pk_u6), "r"(_pk_u7) : "memory"); - } - } - } - } - } - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(o_empty_addr), "r"(0) : "memory"); - tile_cursor_1 = tile_cursor_1 + num_kv_tiles_1; - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - } - } - // ---- Role: mma_warp ---- - if (warp == 8) { - { // mma_warp_main - const int mma_dummy = 0; - unsigned int k_stage = 0; - unsigned int v_stage = 0; - int tile_cursor_2 = 0; - unsigned int _phase_o_first_slice_seeded_0 = 0; - mbarrier_wait(o_first_slice_seeded_addr, _phase_o_first_slice_seeded_0); - _phase_o_first_slice_seeded_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - unsigned int _phase_q_full_0 = 0; - unsigned int _phase_k_full = 0; - unsigned int _phase_o_empty_0 = 1; - unsigned int _phase_v_full = 0; - #pragma unroll 1 - for (unsigned int work_idx_2 = cluster_id; work_idx_2 < total_work_items; work_idx_2 += num_clusters) { - int owner_capacity_2 = ((sparse_topk + 128 - 1) / 128 + ((1) ? 5 : 1) - 1) / ((1) ? 5 : 1); - int first_owner_tile_2 = work_idx_2 % (unsigned int)(((1) ? 5 : 1)) * (unsigned int)owner_capacity_2; - int _min_1 = ((owner_capacity_2) < ((sparse_topk + 128 - 1) / 128 - first_owner_tile_2) ? (owner_capacity_2) : ((sparse_topk + 128 - 1) / 128 - first_owner_tile_2)); - int num_kv_tiles_2 = _min_1; - if (cta_rank == 0) { - mbarrier_wait(q_full_addr, _phase_q_full_0); - _phase_q_full_0 ^= 1; - int first_pv = 1; - #pragma unroll 1 - for (int tile_2 = 0; tile_2 < num_kv_tiles_2; tile_2++) { - int pipeline_tile_2 = tile_cursor_2 + tile_2; - int score_phase_2 = pipeline_tile_2 & 1; - int s_empty_phase = pipeline_tile_2 >> 1 & 1 ^ 1; - mbarrier_wait(s_empty_addr + (score_phase_2) * 8, s_empty_phase); - int score_col = ((score_phase_2 != 0) ? 64 : 0); - #pragma unroll - for (int qk_stage = 0; qk_stage < 4; qk_stage++) { - mbarrier_wait(k_full_addr + (k_stage) * 8, _phase_k_full); - int _mma_a_lo_0 = (((smem_q_addr) >> 4) & 0x3FFF) + (qk_stage * 2) * 512; - int _mma_b_lo_0 = (((smem_k_addr) >> 4) & 0x3FFF) + (k_stage) * 1024; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 136316048;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (score_col))), "r"(((qk_stage == 0) ? 0 : 1))); - int _mma_a_lo_1 = (((smem_q_addr) >> 4) & 0x3FFF) + (qk_stage * 2 + 1) * 512; - int _mma_b_lo_1 = (((smem_k_addr + 8192) >> 4) & 0x3FFF) + (k_stage) * 1024; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 136316048;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); - elect_commit_cg2_multicast(k_empty_addr + (k_stage) * 8, (uint16_t)(3)); - k_stage += 1; - if (k_stage == 8) { k_stage = 0; _phase_k_full ^= 1; } - } - elect_commit_cg2_multicast(s_full_addr + (score_phase_2) * 8, (uint16_t)(3)); - if (tile_2 > 0) { - int p_stage_1 = (pipeline_tile_2 - 1) % 2; - int p_full_phase = (pipeline_tile_2 - 1) / 2 & 1; - mbarrier_wait(o_empty_addr, _phase_o_empty_0); - _phase_o_empty_0 ^= 1; - mbarrier_wait(p_full_addr + (p_stage_1) * 8, p_full_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - #pragma unroll - for (int ps = 0; ps < 2; ps++) { - int smem_p_stage_1 = p_stage_1 * 2 + ps; - mbarrier_wait(v_full_addr + (v_stage) * 8, _phase_v_full); - #pragma unroll - for (int acc_stage_2 = 0; acc_stage_2 < 2; acc_stage_2++) { - int output_col = 128 + acc_stage_2 * 128; - unsigned int v_slot = (k_stage + (unsigned int)(acc_stage_2 * 2 + ps)) % 8; - int _mma_a_lo_2 = (((smem_p_addr) >> 4) & 0x3FFF) + (smem_p_stage_1) * 512; - int _mma_b_lo_2 = ((((smem_v_addr) >> 4) & 0x3FFF) | 0x2000000) + (v_slot) * 1024; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 138478736;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (output_col))), "r"(((((ps == 0) ? first_pv : 0)) ? 0 : 1))); - } - elect_commit_cg2_multicast(v_empty_addr + (v_stage) * 8, (uint16_t)(3)); - #pragma unroll - for (int acc_stage_3 = 0; acc_stage_3 < 2; acc_stage_3++) { - unsigned int free_slot = (k_stage + (unsigned int)(acc_stage_3 * 2 + ps)) % 8; - elect_commit_cg2_multicast(slot_empty_addr + (free_slot) * 8, (uint16_t)(3)); - } - v_stage += 1; - if (v_stage == 2) { v_stage = 0; _phase_v_full ^= 1; } - } - elect_commit_cg2_multicast(p_empty_addr + (p_stage_1) * 8, (uint16_t)(3)); - elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); - first_pv = 0; - } - } - int p_stage_2 = (tile_cursor_2 + num_kv_tiles_2 - 1) % 2; - int p_full_phase_1 = (tile_cursor_2 + num_kv_tiles_2 - 1) / 2 & 1; - mbarrier_wait(o_empty_addr, _phase_o_empty_0); - _phase_o_empty_0 ^= 1; - mbarrier_wait(p_full_addr + (p_stage_2) * 8, p_full_phase_1); - asm volatile("tcgen05.fence::after_thread_sync;"); - #pragma unroll - for (int ps_1 = 0; ps_1 < 2; ps_1++) { - int smem_p_stage_2 = p_stage_2 * 2 + ps_1; - mbarrier_wait(v_full_addr + (v_stage) * 8, _phase_v_full); - #pragma unroll - for (int acc_stage_4 = 0; acc_stage_4 < 2; acc_stage_4++) { - int output_col_1 = 128 + acc_stage_4 * 128; - unsigned int v_slot_1 = (k_stage + (unsigned int)(4 + acc_stage_4 * 2 + ps_1)) % 8; - int _mma_a_lo_3 = (((smem_p_addr) >> 4) & 0x3FFF) + (smem_p_stage_2) * 512; - int _mma_b_lo_3 = ((((smem_v_addr) >> 4) & 0x3FFF) | 0x2000000) + (v_slot_1) * 1024; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 138478736;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"((tmem_tmem_scratch + (output_col_1))), "r"(((((ps_1 == 0) ? first_pv : 0)) ? 0 : 1))); - } - elect_commit_cg2_multicast(v_empty_addr + (v_stage) * 8, (uint16_t)(3)); - #pragma unroll - for (int acc_stage_5 = 0; acc_stage_5 < 2; acc_stage_5++) { - unsigned int free_slot_1 = (k_stage + (unsigned int)(4 + acc_stage_5 * 2 + ps_1)) % 8; - elect_commit_cg2_multicast(slot_empty_addr + (free_slot_1) * 8, (uint16_t)(3)); - } - v_stage += 1; - if (v_stage == 2) { v_stage = 0; _phase_v_full ^= 1; } - } - elect_commit_cg2_multicast(p_empty_addr + (p_stage_2) * 8, (uint16_t)(3)); - elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); - elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); - } - tile_cursor_2 = tile_cursor_2 + num_kv_tiles_2; - } - if (cta_rank == 0) { - #pragma unroll - for (int tail_offset = 0; tail_offset < 2; tail_offset++) { - int tail_tile = tile_cursor_2 + tail_offset; - int tail_stage = tail_tile & 1; - int tail_phase = tail_tile >> 1 & 1 ^ 1; - mbarrier_wait(s_empty_addr + (tail_stage) * 8, tail_phase); - } - mbarrier_wait(o_empty_addr, _phase_o_empty_0); - _phase_o_empty_0 ^= 1; - } - asm volatile("tcgen05.fence::before_thread_sync;"); - int peer_rank_1 = cta_rank ^ 1; - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(tmem_dealloc_peer_addr), "r"(peer_rank_1) : "memory"); - mbarrier_wait(tmem_dealloc_peer_addr, 0); - int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); - asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); - } - } - // ---- Role: load_wg ---- - if (warp >= 9 && warp <= 12) { - { // load_wg_main - const int load_dummy = 0; - const int load_warp_rank = warp - 9; - unsigned int k_stage_1 = 0; - unsigned int v_stage_1 = 0; - int load_tile_cursor = 0; - unsigned int _phase_q_empty_0 = 1; - unsigned int _phase_slot_empty = 1; - unsigned int _phase_v_empty = 1; - #pragma unroll 1 - for (unsigned int work_idx_3 = cluster_id; work_idx_3 < total_work_items; work_idx_3 += num_clusters) { - int split_idx_2 = work_idx_3 % (unsigned int)(((1) ? 5 : 1)); - int query_idx_2 = work_idx_3 / (unsigned int)(((1) ? 5 : 1)); - int owner_capacity_3 = ((sparse_topk + 128 - 1) / 128 + ((1) ? 5 : 1) - 1) / ((1) ? 5 : 1); - int first_owner_tile_3 = work_idx_3 % (unsigned int)(((1) ? 5 : 1)) * (unsigned int)owner_capacity_3; - int _min_0 = ((owner_capacity_3) < ((sparse_topk + 128 - 1) / 128 - first_owner_tile_3) ? (owner_capacity_3) : ((sparse_topk + 128 - 1) / 128 - first_owner_tile_3)); - int num_kv_tiles_3 = _min_0; - int first_kv_tile_1 = first_owner_tile_3; - int group = load_warp_rank * 4 + lane; - mbarrier_wait(q_empty_addr, _phase_q_empty_0); - _phase_q_empty_0 ^= 1; - if (load_warp_rank == 0) { - if (cta_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(131072)) : "memory"); - } - } - if (elect_sync()) { - #pragma unroll - for (int q_stage = 0; q_stage < 8; q_stage++) { - tma_4d_gmem2smem_cta2(smem_q_addr + (unsigned int)(q_stage * 8192), (&tmap_q), 0, cta_rank * 64, q_stage, query_idx_2, ((q_full_addr) & 0xFEFFFFFF)); - } - } - } - asm volatile("barrier.sync 2, 160;" ::: "memory"); - #pragma unroll - for (int qk_stage_1 = 0; qk_stage_1 < 4; qk_stage_1++) { - if (first_kv_tile_1 == 0) { - mbarrier_wait(slot_empty_addr + (k_stage_1) * 8, _phase_slot_empty); - if (load_warp_rank == 0) { - if (cta_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); - } - } - } - if (lane < 4) { - int index_offset = cta_rank * 64 + group * 4; - int raw_rows[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows[0])), "=r"(*reinterpret_cast(&raw_rows[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows[(0) + 3])) - : "r"(smem_indices_addr + (unsigned int)(index_offset * 4))); - tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group * 512), (&tmap_swa_k), qk_stage_1 * 128, ((raw_rows[0] >= 0) ? raw_rows[0] : 0), ((raw_rows[1] >= 0) ? raw_rows[1] : 0), ((raw_rows[2] >= 0) ? raw_rows[2] : 0), ((raw_rows[3] >= 0) ? raw_rows[3] : 0), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group * 512) + 8192, (&tmap_swa_k), qk_stage_1 * 128 + 64, ((raw_rows[0] >= 0) ? raw_rows[0] : 0), ((raw_rows[1] >= 0) ? raw_rows[1] : 0), ((raw_rows[2] >= 0) ? raw_rows[2] : 0), ((raw_rows[3] >= 0) ? raw_rows[3] : 0), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - } - k_stage_1 += 1; - if (k_stage_1 == 8) { k_stage_1 = 0; _phase_slot_empty ^= 1; } - } else { - mbarrier_wait(slot_empty_addr + (k_stage_1) * 8, _phase_slot_empty); - if (load_warp_rank == 0) { - if (cta_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); - } - } - } - if (lane < 4) { - int index_offset_1 = cta_rank * 64 + group * 4; - int raw_rows_1[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_1[0])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 3])) - : "r"(smem_indices_addr + (unsigned int)(index_offset_1 * 4))); - tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group * 512), (&tmap_compressed_k), qk_stage_1 * 128, ((raw_rows_1[0] >= 0) ? raw_rows_1[0] : 0), ((raw_rows_1[1] >= 0) ? raw_rows_1[1] : 0), ((raw_rows_1[2] >= 0) ? raw_rows_1[2] : 0), ((raw_rows_1[3] >= 0) ? raw_rows_1[3] : 0), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group * 512) + 8192, (&tmap_compressed_k), qk_stage_1 * 128 + 64, ((raw_rows_1[0] >= 0) ? raw_rows_1[0] : 0), ((raw_rows_1[1] >= 0) ? raw_rows_1[1] : 0), ((raw_rows_1[2] >= 0) ? raw_rows_1[2] : 0), ((raw_rows_1[3] >= 0) ? raw_rows_1[3] : 0), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - } - k_stage_1 += 1; - if (k_stage_1 == 8) { k_stage_1 = 0; _phase_slot_empty ^= 1; } - } - } - #pragma unroll 1 - for (int tile_3 = 1; tile_3 < num_kv_tiles_3; tile_3++) { - #pragma unroll - for (int qk_stage_2 = 0; qk_stage_2 < 4; qk_stage_2++) { - mbarrier_wait(slot_empty_addr + (k_stage_1) * 8, _phase_slot_empty); - if (load_warp_rank == 0) { - if (cta_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); - } - } - } - if (lane < 4) { - int index_offset_2 = tile_3 * 128 + cta_rank * 64 + group * 4; - int raw_rows_2[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_2[0])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 3])) - : "r"(smem_indices_addr + (unsigned int)(index_offset_2 * 4))); - tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group * 512), (&tmap_compressed_k), qk_stage_2 * 128, ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : 0), ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : 0), ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : 0), ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : 0), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group * 512) + 8192, (&tmap_compressed_k), qk_stage_2 * 128 + 64, ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : 0), ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : 0), ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : 0), ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : 0), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - } - k_stage_1 += 1; - if (k_stage_1 == 8) { k_stage_1 = 0; _phase_slot_empty ^= 1; } - } - int prev_tile = tile_3 - 1; - #pragma unroll - for (int ps_2 = 0; ps_2 < 2; ps_2++) { - if (first_kv_tile_1 + prev_tile == 0) { - mbarrier_wait(v_empty_addr + (v_stage_1) * 8, _phase_v_empty); - if (load_warp_rank == 0) { - if (cta_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); - } - } - } - if (cta_rank != ps_2) { - int fill_base = (load_tile_cursor + prev_tile) * 4; - #pragma unroll - for (int acc_stage_6 = 0; acc_stage_6 < 2; acc_stage_6++) { - int fill = fill_base + acc_stage_6 * 2 + ps_2; - mbarrier_wait(k_empty_addr + ((k_stage_1 + (unsigned int)(acc_stage_6 * 2 + ps_2)) % 8) * 8, fill / 8 & 1); - } - if (lane < 4) { - int index_offset_3 = prev_tile * 128 + ps_2 * 64 + group * 4; - int raw_rows_3[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_3[0])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 3])) - : "r"(smem_indices_addr + (unsigned int)(index_offset_3 * 4))); - #pragma unroll - for (int acc_stage_7 = 0; acc_stage_7 < 2; acc_stage_7++) { - unsigned int copy_slot = (k_stage_1 + (unsigned int)(acc_stage_7 * 2 + ps_2)) % 8; - tma_gather4_gmem2smem_mc_cta2(smem_k_addr + copy_slot * 16384 + (unsigned int)(group * 512), (&tmap_swa_v), (acc_stage_7 * 2 + (1 - ps_2)) * 128, ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : 0), ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : 0), ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : 0), ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : 0), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - tma_gather4_gmem2smem_mc_cta2(smem_k_addr + copy_slot * 16384 + (unsigned int)(group * 512) + 8192, (&tmap_swa_v), (acc_stage_7 * 2 + (1 - ps_2)) * 128 + 64, ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : 0), ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : 0), ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : 0), ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : 0), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - } - } - } - v_stage_1 += 1; - if (v_stage_1 == 2) { v_stage_1 = 0; _phase_v_empty ^= 1; } - } else { - mbarrier_wait(v_empty_addr + (v_stage_1) * 8, _phase_v_empty); - if (load_warp_rank == 0) { - if (cta_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); - } - } - } - if (cta_rank != ps_2) { - int fill_base_1 = (load_tile_cursor + prev_tile) * 4; - #pragma unroll - for (int acc_stage_8 = 0; acc_stage_8 < 2; acc_stage_8++) { - int fill_1 = fill_base_1 + acc_stage_8 * 2 + ps_2; - mbarrier_wait(k_empty_addr + ((k_stage_1 + (unsigned int)(acc_stage_8 * 2 + ps_2)) % 8) * 8, fill_1 / 8 & 1); - } - if (lane < 4) { - int index_offset_4 = prev_tile * 128 + ps_2 * 64 + group * 4; - int raw_rows_4[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_4[0])), "=r"(*reinterpret_cast(&raw_rows_4[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_4[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_4[(0) + 3])) - : "r"(smem_indices_addr + (unsigned int)(index_offset_4 * 4))); - #pragma unroll - for (int acc_stage_9 = 0; acc_stage_9 < 2; acc_stage_9++) { - unsigned int copy_slot_1 = (k_stage_1 + (unsigned int)(acc_stage_9 * 2 + ps_2)) % 8; - tma_gather4_gmem2smem_mc_cta2(smem_k_addr + copy_slot_1 * 16384 + (unsigned int)(group * 512), (&tmap_compressed_v), (acc_stage_9 * 2 + (1 - ps_2)) * 128, ((raw_rows_4[0] >= 0) ? raw_rows_4[0] : 0), ((raw_rows_4[1] >= 0) ? raw_rows_4[1] : 0), ((raw_rows_4[2] >= 0) ? raw_rows_4[2] : 0), ((raw_rows_4[3] >= 0) ? raw_rows_4[3] : 0), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - tma_gather4_gmem2smem_mc_cta2(smem_k_addr + copy_slot_1 * 16384 + (unsigned int)(group * 512) + 8192, (&tmap_compressed_v), (acc_stage_9 * 2 + (1 - ps_2)) * 128 + 64, ((raw_rows_4[0] >= 0) ? raw_rows_4[0] : 0), ((raw_rows_4[1] >= 0) ? raw_rows_4[1] : 0), ((raw_rows_4[2] >= 0) ? raw_rows_4[2] : 0), ((raw_rows_4[3] >= 0) ? raw_rows_4[3] : 0), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - } - } - } - v_stage_1 += 1; - if (v_stage_1 == 2) { v_stage_1 = 0; _phase_v_empty ^= 1; } - } - } - } - int last_tile = num_kv_tiles_3 - 1; - #pragma unroll - for (int ps_3 = 0; ps_3 < 2; ps_3++) { - if (first_kv_tile_1 + last_tile == 0) { - mbarrier_wait(v_empty_addr + (v_stage_1) * 8, _phase_v_empty); - if (load_warp_rank == 0) { - if (cta_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); - } - } - } - if (cta_rank != ps_3) { - int fill_base_2 = (load_tile_cursor + last_tile) * 4; - #pragma unroll - for (int acc_stage_10 = 0; acc_stage_10 < 2; acc_stage_10++) { - int fill_2 = fill_base_2 + acc_stage_10 * 2 + ps_3; - mbarrier_wait(k_empty_addr + ((k_stage_1 + (unsigned int)(4 + acc_stage_10 * 2 + ps_3)) % 8) * 8, fill_2 / 8 & 1); - } - if (lane < 4) { - int index_offset_5 = last_tile * 128 + ps_3 * 64 + group * 4; - int raw_rows_5[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_5[0])), "=r"(*reinterpret_cast(&raw_rows_5[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_5[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_5[(0) + 3])) - : "r"(smem_indices_addr + (unsigned int)(index_offset_5 * 4))); - #pragma unroll - for (int acc_stage_11 = 0; acc_stage_11 < 2; acc_stage_11++) { - unsigned int copy_slot_2 = (k_stage_1 + (unsigned int)(4 + acc_stage_11 * 2 + ps_3)) % 8; - tma_gather4_gmem2smem_mc_cta2(smem_k_addr + copy_slot_2 * 16384 + (unsigned int)(group * 512), (&tmap_swa_v), (acc_stage_11 * 2 + (1 - ps_3)) * 128, ((raw_rows_5[0] >= 0) ? raw_rows_5[0] : 0), ((raw_rows_5[1] >= 0) ? raw_rows_5[1] : 0), ((raw_rows_5[2] >= 0) ? raw_rows_5[2] : 0), ((raw_rows_5[3] >= 0) ? raw_rows_5[3] : 0), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - tma_gather4_gmem2smem_mc_cta2(smem_k_addr + copy_slot_2 * 16384 + (unsigned int)(group * 512) + 8192, (&tmap_swa_v), (acc_stage_11 * 2 + (1 - ps_3)) * 128 + 64, ((raw_rows_5[0] >= 0) ? raw_rows_5[0] : 0), ((raw_rows_5[1] >= 0) ? raw_rows_5[1] : 0), ((raw_rows_5[2] >= 0) ? raw_rows_5[2] : 0), ((raw_rows_5[3] >= 0) ? raw_rows_5[3] : 0), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - } - } - } - v_stage_1 += 1; - if (v_stage_1 == 2) { v_stage_1 = 0; _phase_v_empty ^= 1; } - } else { - mbarrier_wait(v_empty_addr + (v_stage_1) * 8, _phase_v_empty); - if (load_warp_rank == 0) { - if (cta_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); - } - } - } - if (cta_rank != ps_3) { - int fill_base_3 = (load_tile_cursor + last_tile) * 4; - #pragma unroll - for (int acc_stage_12 = 0; acc_stage_12 < 2; acc_stage_12++) { - int fill_3 = fill_base_3 + acc_stage_12 * 2 + ps_3; - mbarrier_wait(k_empty_addr + ((k_stage_1 + (unsigned int)(4 + acc_stage_12 * 2 + ps_3)) % 8) * 8, fill_3 / 8 & 1); - } - if (lane < 4) { - int index_offset_6 = last_tile * 128 + ps_3 * 64 + group * 4; - int raw_rows_6[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_6[0])), "=r"(*reinterpret_cast(&raw_rows_6[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_6[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_6[(0) + 3])) - : "r"(smem_indices_addr + (unsigned int)(index_offset_6 * 4))); - #pragma unroll - for (int acc_stage_13 = 0; acc_stage_13 < 2; acc_stage_13++) { - unsigned int copy_slot_3 = (k_stage_1 + (unsigned int)(4 + acc_stage_13 * 2 + ps_3)) % 8; - tma_gather4_gmem2smem_mc_cta2(smem_k_addr + copy_slot_3 * 16384 + (unsigned int)(group * 512), (&tmap_compressed_v), (acc_stage_13 * 2 + (1 - ps_3)) * 128, ((raw_rows_6[0] >= 0) ? raw_rows_6[0] : 0), ((raw_rows_6[1] >= 0) ? raw_rows_6[1] : 0), ((raw_rows_6[2] >= 0) ? raw_rows_6[2] : 0), ((raw_rows_6[3] >= 0) ? raw_rows_6[3] : 0), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - tma_gather4_gmem2smem_mc_cta2(smem_k_addr + copy_slot_3 * 16384 + (unsigned int)(group * 512) + 8192, (&tmap_compressed_v), (acc_stage_13 * 2 + (1 - ps_3)) * 128 + 64, ((raw_rows_6[0] >= 0) ? raw_rows_6[0] : 0), ((raw_rows_6[1] >= 0) ? raw_rows_6[1] : 0), ((raw_rows_6[2] >= 0) ? raw_rows_6[2] : 0), ((raw_rows_6[3] >= 0) ? raw_rows_6[3] : 0), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - } - } - } - v_stage_1 += 1; - if (v_stage_1 == 2) { v_stage_1 = 0; _phase_v_empty ^= 1; } - } - } - asm volatile("barrier.sync 2, 160;" ::: "memory"); - load_tile_cursor = load_tile_cursor + num_kv_tiles_3; - } - mbarrier_wait(q_empty_addr, _phase_q_empty_0); - _phase_q_empty_0 ^= 1; - } - } - // ---- Role: index_warp ---- - if (warp == 13) { - { // index_warp_main - const int index_dummy = 0; - int valid_iter_1 = 0; - #pragma unroll 1 - for (unsigned int work_idx_4 = cluster_id; work_idx_4 < total_work_items; work_idx_4 += num_clusters) { - int split_idx_3 = work_idx_4 % (unsigned int)(((1) ? 5 : 1)); - int query_idx_3 = work_idx_4 / (unsigned int)(((1) ? 5 : 1)); - int owner_col_base = split_idx_3 * ((9 + ((1) ? 5 : 1) - 1) / ((1) ? 5 : 1)) * 128; - int* swa_row = swa_indices + (query_idx_3 * swa_index_stride); - int* compressed_row = compressed_indices + (query_idx_3 * compressed_index_stride); - int valid_slot_1 = valid_iter_1 & 1; - int valid_empty_phase = valid_iter_1 >> 1 & 1 ^ 1; - mbarrier_wait(index_valid_empty_addr + (valid_slot_1) * 8, valid_empty_phase); - #pragma unroll - for (int index_pass = 0; index_pass < (9 + ((1) ? 5 : 1) - 1) / ((1) ? 5 : 1); index_pass++) { - int index_offset_7 = index_pass * 128 + lane * 4; - int combined_col = owner_col_base + index_offset_7; - int* index_src = ((combined_col < 128) ? swa_row + combined_col : compressed_row + (combined_col - 128)); - int values[4]; - #pragma unroll - for (int i = 0; i < 4; i++) { - values[i] = -1; - } - if (combined_col + 4 <= sparse_topk) { - int _vec_load_0[4]; - { - const int4* _ivptr_0 = reinterpret_cast(index_src + 0); - int4 _ivld_0; - _ivld_0 = *_ivptr_0; - _vec_load_0[0 + 0] = _ivld_0.x; - _vec_load_0[0 + 1] = _ivld_0.y; - _vec_load_0[0 + 2] = _ivld_0.z; - _vec_load_0[0 + 3] = _ivld_0.w; - } - #pragma unroll - for (int i_1 = 0; i_1 < 4; i_1++) { - values[i_1] = _vec_load_0[i_1]; - } - } - if (combined_col < sparse_topk && combined_col + 4 > sparse_topk) { - #pragma unroll - for (int i_2 = 0; i_2 < 4; i_2++) { - if (combined_col + i_2 < sparse_topk) { - values[i_2] = index_src[i_2]; - } - } - } - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(smem_indices_addr + (unsigned int)(index_offset_7 * 4)), "r"(values[0]), "r"(values[1]), "r"(values[2]), "r"(values[3]) : "memory"); - unsigned int valid_word = (((values[0] >= 0) ? 1 : 0) | ((values[1] >= 0) ? 1 : 0) << 1 | ((values[2] >= 0) ? 1 : 0) << 2 | ((values[3] >= 0) ? 1 : 0) << 3) << (lane & 7) * 4; - unsigned int _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, valid_word, 1); - valid_word = valid_word | _shfl_xor_0; - unsigned int _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, valid_word, 2); - valid_word = valid_word | _shfl_xor_1; - unsigned int _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, valid_word, 4); - valid_word = valid_word | _shfl_xor_2; - if ((lane & 7) == 0) { - smem_index_valid[valid_slot_1 * 8 + index_pass * 4 + (lane >> 3)] = valid_word; - } - } - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(index_valid_full_addr + (valid_slot_1) * 8); - asm volatile("barrier.sync 2, 160;" ::: "memory"); - asm volatile("barrier.sync 2, 160;" ::: "memory"); - valid_iter_1 = valid_iter_1 + 1; - } - } - } - // ---- Role: empty_warps ---- - if (warp >= 14 && warp <= 15) { - { // empty_warps_main - const int empty_dummy = 0; - #pragma unroll 1 - for (unsigned int work_idx_5 = cluster_id; work_idx_5 < total_work_items; work_idx_5 += num_clusters) { - asm volatile("tcgen05.fence::after_thread_sync;"); - } - } - } - - // Cleanup -} - -} // extern "C" diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_3e5a644d61420c9c7ea1_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_3e5a644d61420c9c7ea1_binding.cu deleted file mode 100644 index 6d2febf4ef8..00000000000 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_3e5a644d61420c9c7ea1_binding.cu +++ /dev/null @@ -1,516 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -// Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_3e5a644d61420c9c7ea1'. -#include -#include -#include -#include - -#include "tvm_ffi_utils.h" - -#include -#include -#include -#include -#include -#include -#include - -struct CakeTensorMap; -extern "C" __global__ void kernel_cake_dsv4_3e5a644d61420c9c7ea1(CakeTensorMap const* tmap_q, CakeTensorMap const* tmap_swa_kv, CakeTensorMap const* tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int num_query_tokens, int has_sinks); - - -namespace cake_host_shim_6afd242c7dfa9fb2 { - -using tvm::ffi::Optional; -using tvm::ffi::TensorView; - -class ScopedCudaDevice { - public: - explicit ScopedCudaDevice(int device_id) { - cudaError_t error = cudaGetDevice(&previous_device_); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed before host-shim launch: cudaError=" - << static_cast(error); - if (previous_device_ != device_id) { - error = cudaSetDevice(device_id); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaSetDevice failed before host-shim launch for cuda:" - << device_id << ": cudaError=" << static_cast(error); - restore_ = true; - } - } - - ScopedCudaDevice(const ScopedCudaDevice&) = delete; - ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; - - ~ScopedCudaDevice() noexcept { - if (restore_) { - (void)cudaSetDevice(previous_device_); - } - } - - private: - int previous_device_ = -1; - bool restore_ = false; -}; - -inline int64_t CakeDeviceMultiprocessorCount(int device_id) { - constexpr int kMaxCachedCudaDevices = 64; - TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) - << "physical-SM-count cache does not cover cuda:" << device_id; - static std::atomic count_by_device[kMaxCachedCudaDevices]{}; - int cached = count_by_device[device_id].load(std::memory_order_acquire); - if (cached > 0) return cached; - - int count = 0; - cudaError_t error = cudaDeviceGetAttribute( - &count, cudaDevAttrMultiProcessorCount, device_id); - TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) - << "querying multiProcessorCount failed for cuda:" << device_id - << ": cudaError=" << static_cast(error) << ", count=" << count; - // Concurrent first launches may repeat the immutable device query, but all - // publication is atomic and every later hot-path lookup is one acquire load. - count_by_device[device_id].store(count, std::memory_order_release); - return count; -} - -inline void CheckCudaTensor(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) - << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; -} - -inline void CheckSameCudaDevice( - const TensorView& t, - const TensorView& reference, - const char* name, - const char* reference_name) { - TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) - << name << " must be on the same CUDA device as " << reference_name - << ": got cuda:" << t.device().device_id - << " versus cuda:" << reference.device().device_id; -} - -inline void CheckCurrentCudaDevice( - const TensorView& reference, - const char* reference_name) { - int current_device = -1; - cudaError_t error = cudaGetDevice(¤t_device); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed while validating " << reference_name - << ": cudaError=" << static_cast(error); - TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) - << "current CUDA device must match " << reference_name - << ": current=cuda:" << current_device - << ", tensor=cuda:" << reference.device().device_id; -} - -inline void CheckContiguous(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; -} - -inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { - DLDataType d = t.dtype(); - TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) - << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits - << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits - << ", lanes=" << (int)d.lanes << ")"; -} - -// A logical axis.outer(trailing) folds every source dim above the trailing -// dimensions. Shape products are independent of physical strides, so verify -// the leading dimensions form one dense row-major chain instead of inventing -// a "folded stride". The descriptor reads its exact adjacent physical step -// separately through stride[-(trailing + 1)]. -inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { - TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) - << name << " cannot fold leading dimensions above " << trailing - << " trailing dims from ndim=" << t.ndim(); - int outer_last = t.ndim() - trailing - 1; - if (outer_last <= 0) { - return; - } - int64_t step = t.stride(outer_last); - TVM_FFI_CHECK(step > 0, ValueError) - << name << " physical strides must be positive"; - int64_t expected = step; - for (int axis = outer_last - 1; axis >= 0; --axis) { - expected *= t.size(axis + 1); - if (t.size(axis) > 1) { - TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) - << name << " leading dims are not physically foldable above " << trailing - << " trailing dims: stride(" << axis << ")=" << t.stride(axis) - << ", expected " << expected; - } - } -} - -// Explicit immutable caller storage: initialize once, outside graph capture, -// exactly as the production TmaDeviceSlot path. The owner must retain the -// allocation and must never expose it as mutable scratch or rewrite its bytes. -template -static inline void PrepareImmutableCallerTmaWorkspace( - const tvm::ffi::Tensor& retained_workspace, - const CUtensorMap (&maps)[N], cudaStream_t stream) { - void* workspace = retained_workspace.data_ptr(); - CUcontext context = nullptr; - CUresult result = cuCtxGetCurrent(&context); - TVM_FFI_CHECK(result == CUDA_SUCCESS && context != nullptr, RuntimeError) - << "immutable TMA workspace requires an active CUDA context"; - CUcontext allocation_context = nullptr; - result = cuPointerGetAttribute( - &allocation_context, CU_POINTER_ATTRIBUTE_CONTEXT, - reinterpret_cast(workspace)); - TVM_FFI_CHECK(result == CUDA_SUCCESS && allocation_context == context, RuntimeError) - << "immutable TMA workspace must belong to the active CUDA context"; - unsigned long long buffer_id = 0; - result = cuPointerGetAttribute( - &buffer_id, CU_POINTER_ATTRIBUTE_BUFFER_ID, - reinterpret_cast(workspace)); - TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) - << "querying immutable TMA allocation identity failed"; - std::string key = std::to_string(reinterpret_cast(context)); - key += ":" + std::to_string(reinterpret_cast(workspace)); - key += ":" + std::to_string(buffer_id); - std::string descriptor_bytes(reinterpret_cast(maps), sizeof(maps)); - struct InitializedWorkspace { - tvm::ffi::Tensor owner; - std::string descriptor_bytes; - }; - static std::mutex mu; - // Keep the actual Tensor allocation alive: a caching allocator may reuse a - // freed suballocation without changing its enclosing CUDA buffer identity. - static auto* initialized = new std::unordered_map(); - std::lock_guard lock(mu); - auto it = initialized->find(key); - if (it != initialized->end()) { - TVM_FFI_CHECK(it->second.descriptor_bytes == descriptor_bytes, ValueError) - << "immutable TMA descriptors changed; provide new private descriptor storage"; - return; - } - CUstreamCaptureStatus capture_status = CU_STREAM_CAPTURE_STATUS_NONE; - result = cuStreamIsCapturing(reinterpret_cast(stream), &capture_status); - TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) - << "querying immutable TMA initialization capture state failed"; - TVM_FFI_CHECK(capture_status == CU_STREAM_CAPTURE_STATUS_NONE, RuntimeError) - << "immutable TMA storage must be initialized before CUDA Graph capture"; - // Synchronous host-to-device initialization completes before publishing the - // slot to any thread/stream, matching TmaDeviceSlot. The consumer retains its - // generated tensor-map acquire; the mutable upload's release fence is unchanged. - result = cuMemcpyHtoD(reinterpret_cast(workspace), maps, sizeof(maps)); - TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) - << "initializing immutable caller TMA descriptors failed"; - initialized->emplace( - std::move(key), InitializedWorkspace{retained_workspace, std::move(descriptor_bytes)}); -} - -static inline void* CallerTmaWorkspaceSlot( - const tvm::ffi::TensorView& workspace, size_t slot) { - return static_cast(workspace.data_ptr()) + slot * sizeof(CUtensorMap); -} - -// 4D TMA descriptor for buffer 'tmap_q' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 3, ValueError) - << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - int64_t d2 = t.size(t.ndim() - 2); - int64_t d3 = t.size(t.ndim() - 3); - TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) - << "TMA source 'tmap_q' trailing dims must be positive"; - TVM_FFI_CHECK(d1 % 64 == 0, ValueError) - << "TMA source 'tmap_q' extent " << d1 - << " must divide exactly by " << 64; - uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(d3)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) - << "TMA box (64, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; - int64_t carrier_stride_0 = d1; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; - int64_t carrier_stride_1 = 64; - TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; - TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; - TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; - int64_t carrier_stride_2 = (d2 * d1); - TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; - TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; - TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; - uint64_t global_strides[3] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - (uint64_t)((carrier_stride_1 * 16) / 8), - (uint64_t)((carrier_stride_2 * 16) / 8), - }; - uint32_t box_dim[4] = {64u, 64u, 1u, 1u}; - uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_swa_kv' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_swa_kv' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_kv'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_compressed_kv' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_compressed_kv(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_compressed_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_compressed_kv' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_compressed_kv' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_compressed_kv"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_compressed_kv' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_kv'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_kv') failed: CUresult=" << (int)r; - return tm; -} - -void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_compressed_kv, TensorView arg_O, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_num_query_tokens, int64_t arg_has_sinks, tvm::ffi::Tensor arg_tma_descriptor_workspace, int64_t grid_x, int64_t grid_y, int64_t grid_z) { - DLDevice dev = arg_tmap_q.device(); - ScopedCudaDevice device_guard(dev.device_id); - CheckCudaTensor(arg_tmap_q, "tmap_q"); - CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); - CheckContiguous(arg_tmap_q, "tmap_q"); - CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); - CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 4, 16, 1); - CheckCudaTensor(arg_tmap_compressed_kv, "tmap_compressed_kv"); - CheckDtype(arg_tmap_compressed_kv, "tmap_compressed_kv", 4, 16, 1); - CheckCudaTensor(arg_O, "O"); - CheckDtype(arg_O, "O", 4, 16, 1); - CheckContiguous(arg_O, "O"); - CheckCudaTensor(arg_swa_indices, "swa_indices"); - CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); - CheckContiguous(arg_swa_indices, "swa_indices"); - CheckCudaTensor(arg_compressed_indices, "compressed_indices"); - CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); - CheckContiguous(arg_compressed_indices, "compressed_indices"); - CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); - CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckCudaTensor(arg_sinks, "sinks"); - CheckDtype(arg_sinks, "sinks", 2, 32, 1); - CheckContiguous(arg_sinks, "sinks"); - CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); - CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); - CheckContiguous(arg_bmm1_scale, "bmm1_scale"); - CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); - CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); - CheckContiguous(arg_bmm2_scale, "bmm2_scale"); - TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) - << "scalar 'num_heads' value " << arg_num_heads - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) - << "scalar 'swa_index_stride' value " << arg_swa_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) - << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) - << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) - << "scalar 'sparse_topk' value " << arg_sparse_topk - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_num_query_tokens >= -2147483648LL && arg_num_query_tokens <= 2147483647LL, ValueError) - << "scalar 'num_query_tokens' value " << arg_num_query_tokens - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) - << "scalar 'has_sinks' value " << arg_has_sinks - << " is outside i32 range [-2147483648, 2147483647]"; - CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); - CheckSameCudaDevice(arg_tmap_compressed_kv, arg_tmap_q, "tmap_compressed_kv", "tmap_q"); - CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); - CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); - CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); - CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); - CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); - CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); - CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); - CheckCudaTensor(arg_tma_descriptor_workspace, "tma_descriptor_workspace"); - CheckSameCudaDevice(arg_tma_descriptor_workspace, arg_tmap_q, "tma_descriptor_workspace", "tmap_q"); - CheckContiguous(arg_tma_descriptor_workspace, "tma_descriptor_workspace"); - CheckDtype(arg_tma_descriptor_workspace, "tma_descriptor_workspace", 1, 8, 1); - TVM_FFI_CHECK(arg_tma_descriptor_workspace.numel() >= 384, ValueError) - << "tma_descriptor_workspace requires at least 384 bytes"; - TVM_FFI_CHECK( - reinterpret_cast(arg_tma_descriptor_workspace.data_ptr()) % 128 == 0, - ValueError) - << "tma_descriptor_workspace must be 128-byte aligned"; - TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) - << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y - << ", " << grid_z << ")"; - - cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); - CUtensorMap h_tmap_q = EncodeTma_tmap_q(arg_tmap_q); - CakeTensorMap const* p_tmap_q = static_cast(CallerTmaWorkspaceSlot(arg_tma_descriptor_workspace, 0u)); - CUtensorMap h_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); - CakeTensorMap const* p_tmap_swa_kv = static_cast(CallerTmaWorkspaceSlot(arg_tma_descriptor_workspace, 1u)); - CUtensorMap h_tmap_compressed_kv = EncodeTma_tmap_compressed_kv(arg_tmap_compressed_kv); - CakeTensorMap const* p_tmap_compressed_kv = static_cast(CallerTmaWorkspaceSlot(arg_tma_descriptor_workspace, 2u)); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); - int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); - int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); - int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); - float* p_sinks = static_cast(arg_sinks.data_ptr()); - float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); - float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); - int32_t v_num_heads = (int32_t)arg_num_heads; - int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; - int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; - int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; - int32_t v_sparse_topk = (int32_t)arg_sparse_topk; - int32_t v_num_query_tokens = (int32_t)arg_num_query_tokens; - int32_t v_has_sinks = (int32_t)arg_has_sinks; - void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_tmap_compressed_kv, &p_O, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_num_query_tokens, &v_has_sinks}; - - static std::mutex smem_mu; - static auto* smem_status_by_device = - new std::unordered_map(); - cudaError_t smem_status = cudaSuccess; - { - std::lock_guard lock(smem_mu); - auto it = smem_status_by_device->find(dev.device_id); - if (it == smem_status_by_device->end()) { - int current_device = -1; - smem_status = cudaGetDevice(¤t_device); - if (smem_status == cudaSuccess && current_device != dev.device_id) { - smem_status = cudaErrorInvalidDevice; - } - if (smem_status == cudaSuccess) { - smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_3e5a644d61420c9c7ea1), - cudaFuncAttributeMaxDynamicSharedMemorySize, - 203904); - } - smem_status_by_device->emplace(dev.device_id, smem_status); - } else { - smem_status = it->second; - } - } - TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_3e5a644d61420c9c7ea1 failed: " - << cudaGetErrorString(smem_status); - dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); - dim3 block(512u, 1u, 1u); - const CUtensorMap workspace_maps[] = {h_tmap_q, h_tmap_swa_kv, h_tmap_compressed_kv}; - PrepareImmutableCallerTmaWorkspace(arg_tma_descriptor_workspace, workspace_maps, stream); - - cudaLaunchConfig_t config{}; - config.gridDim = grid; - config.blockDim = block; - config.dynamicSmemBytes = 203904u; - config.stream = stream; - cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_3e5a644d61420c9c7ea1), kargs); - TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_3e5a644d61420c9c7ea1 failed: " - << cudaGetErrorString(launch_status); -} - -} // namespace cake_host_shim_6afd242c7dfa9fb2 - -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_6afd242c7dfa9fb2::Run); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_3e5a644d61420c9c7ea1_kernel.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_3e5a644d61420c9c7ea1_kernel.cu deleted file mode 100644 index d4ab4dbfb83..00000000000 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_3e5a644d61420c9c7ea1_kernel.cu +++ /dev/null @@ -1,2550 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -typedef signed char int8_t; -typedef unsigned char uint8_t; -typedef unsigned short uint16_t; -typedef unsigned int uint32_t; -#if defined(__CUDACC_RTC__) -typedef unsigned long long uint64_t; -#else -typedef unsigned long uint64_t; -#endif -static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); -typedef signed int int32_t; -typedef short int int16_t; -struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; -struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; -static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); -static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); -template -struct __align__(128) CakeTensorMapPack { CakeTensorMap maps[N]; }; - -#if defined(__CUDACC_RTC__) -typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; -#else -#include -#endif - -static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); -static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); -#include -#include - -#define CAKE_INF CUDART_INF_F -#define TMEM_NCOLS 512 -#define TMEM_TMEM_OFFSET 0 -#define NUM_KV_PIPE_STAGES 4 -#define NUM_PAGE_PIPE_STAGES 6 -#define NUM_WORK_PIPE_STAGES 2 -#define NUM_THROTTLE_PIPE_STAGES 2 -#define SMEM_SMEM_Q_OFF 1024 -#define SMEM_SMEM_Q_STAGE_BYTES 16384 -#define SMEM_SMEM_Q_STRIDE 16384 -#define SMEM_SMEM_KV_OFF 66560 -#define SMEM_SMEM_KV_STAGE_BYTES 32768 -#define SMEM_SMEM_KV_STRIDE 32768 -#define SMEM_SMEM_V_OFF 66560 -#define SMEM_SMEM_V_STAGE_BYTES 32768 -#define SMEM_SMEM_V_STRIDE 32768 -#define SMEM_SMEM_PAGE_OFFSETS_OFF 197632 -#define SMEM_SMEM_PAGE_OFFSETS_STAGE_BYTES 1024 -#define SMEM_SMEM_PAGE_OFFSETS_STRIDE 1024 -#define SMEM_WORK_RESPONSE_OFF 203776 -#define SMEM_WORK_RESPONSE_STAGE_BYTES 16 -#define SMEM_WORK_RESPONSE_STRIDE 16 -#define SMEM_TOTAL 203904 -#define THREADS 512 - -#include - -__device__ __forceinline__ uint32_t elect_sync() { - uint32_t pred = 0; - asm volatile( - "{\n\t" - ".reg .pred %%px;\n\t" - "elect.sync _|%%px, %1;\n\t" - "@%%px mov.s32 %0, 1;\n\t" - "}\n" - : "+r"(pred) - : "r"(0xFFFFFFFF)); - return pred; -} - - -__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { - asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" - :: "r"(mbar_addr), "r"(count) : "memory"); -} - -__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { - asm volatile("mbarrier.init.b64 [%0], %1;" - :: "l"(mbar_addr), "r"(count) : "memory"); -} - - -__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - - -// CTA-local pipelines have short, resident producer/consumer edges. Omitting -// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the -// explicit loop still makes this helper blocking until acquire succeeds. -__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE;\n\t" - "bra.uni LAB_WAIT;\n\t" - "DONE:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Source-faithful relaxed CTA wait used only by a typed protocol that does -// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. -__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, 10000000;\n\t" - "@P1 bra.uni DONE_RELAXED;\n\t" - "bra.uni LAB_WAIT_RELAXED;\n\t" - "DONE_RELAXED:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Exact source ports may request the PTX suspendTimeHint operand explicitly. -// The hint is expressed in nanoseconds and is kept separate from the canonical -// no-hint CTA helper so unrelated schedules retain their existing retry path. -__device__ __forceinline__ void mbarrier_wait_suspend( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_SUSPEND:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_SUSPEND;\n\t" - "bra.uni LAB_WAIT_SUSPEND;\n\t" - "DONE_SUSPEND:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE_CLUSTER;\n\t" - "bra.uni LAB_WAIT_CLUSTER;\n\t" - "DONE_CLUSTER:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - ".reg .u32 WAIT_ADDR;\n\t" - "mov.u32 WAIT_ADDR, %0;\n\t" - "LAB_WAIT_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [WAIT_ADDR], %1, %2;\n\t" - "@P1 bra.uni DONE_HINT;\n\t" - "bra.uni LAB_WAIT_HINT;\n\t" - "DONE_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -// Exact unqualified CTA wait used by source schedules whose PTX intentionally -// omits the acquire qualifier while retaining a typed suspendTimeHint operand. -__device__ __forceinline__ void mbarrier_wait_relaxed_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED_HINT:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra DONE_RELAXED_HINT;\n\t" - "bra LAB_WAIT_RELAXED_HINT;\n\t" - "DONE_RELAXED_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); -} - -__device__ __forceinline__ void mbarrier_wait_cluster_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" - "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" - "DONE_CLUSTER_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_suspend( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait_cluster(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); - } -} - - -__device__ __forceinline__ void tcgen05_mma_f16( - int taddr, uint64_t a_desc, uint64_t b_desc, - uint32_t i_desc, int enable_input_d) { - asm volatile( - "{\n\t" - ".reg .pred p;\n\t" - "setp.ne.b32 p, %4, 0;\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" - "}\n" - :: "r"(taddr), "l"(a_desc), "l"(b_desc), - "r"(i_desc), "r"(enable_input_d) - : "memory"); -} - - -__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { - return (x & 0x3FFFFULL) >> 4ULL; -} - - -union MmaSmemDesc { - uint64_t u64; - uint32_t u32[2]; -}; - -__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { - MmaSmemDesc tmp; - tmp.u64 = smem_desc; - tmp.u32[0] += offset; - smem_desc = tmp.u64; -} - - -__device__ __forceinline__ void elect_commit(int mbar_addr) { - asm volatile( - "{\n\t" - ".reg .pred leader;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" - ".shared::cluster.b64 [%0];\n\t" - "}\n" - :: "r"(mbar_addr)); -} - - -__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { - asm volatile( - "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" - :: "r"(mbar_addr), "r"(bytes) : "memory"); -} - - -__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7," - " %8, %9, %10, %11, %12, %13, %14, %15," - " %16, %17, %18, %19, %20, %21, %22, %23," - " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), - "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), - "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), - "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), - "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), - "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), - "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x16.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7," - " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), - "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), - "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ float approx_exp2(float x) { - float y; - asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float approx_rcp(float x) { - float y; - asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float max_noftz(float a, float b) { - float c; - asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); - return c; -} - - -__device__ __forceinline__ float warp_reduce_max(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); - return val; -} - - -__device__ __forceinline__ float warp_reduce_sum(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val += __shfl_xor_sync(0xFFFFFFFF, val, offset); - return val; -} - - -__device__ __forceinline__ float row_max_reduce(float2 acc) { - return max_noftz(acc.x, acc.y); -} - - -__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { - #pragma unroll - for (int j = 0; j < 16; j++) { - if (j % 2 == 0) - acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); - else - acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); - } -} - - -__device__ __forceinline__ float2 ex2_emulation_f32x2_value(float2 value) { - const float c0 = 1.0f, c1 = 0.695146143436431884765625f; - const float c2 = 0.227564394474029541015625f, c3 = 0.077119089663028717041015625f; - const float magic = 12582912.0f; - float x0 = max_noftz(value.x, -127.0f), x1 = max_noftz(value.y, -127.0f); - float2 xc2 = make_float2(x0, x1), magic2 = make_float2(magic, magic); - float2 xr2; - asm("add.rm.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xr2) - : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&magic2)); - float2 c3_2 = make_float2(c3, c3), c2_2 = make_float2(c2, c2); - float2 c1_2 = make_float2(c1, c1), c0_2 = make_float2(c0, c0); - float2 xrb2, xfrac2; - asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xrb2) - : "l"(*(unsigned long long*)&xr2), "l"(*(unsigned long long*)&magic2)); - asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xfrac2) - : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&xrb2)); - float2 poly2; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&c3_2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c2_2)); - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c1_2)); - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c0_2)); - int x0r_i, x1r_i, p0_i, p1_i; - asm("mov.b64 {%0, %1}, %2;" : "=r"(x0r_i), "=r"(x1r_i) : "l"(*(unsigned long long*)&xr2)); - asm("mov.b64 {%0, %1}, %2;" : "=r"(p0_i), "=r"(p1_i) : "l"(*(unsigned long long*)&poly2)); - float r0, r1; - asm("mov.b32 %0, %1;" : "=f"(r0) : "r"((x0r_i << 23) + p0_i)); - asm("mov.b32 %0, %1;" : "=f"(r1) : "r"((x1r_i << 23) + p1_i)); - return make_float2(r0, r1); -} - -__device__ __forceinline__ void ex2_emulation_f32x2(float* x0_ptr, float* x1_ptr) { - float2 result = ex2_emulation_f32x2_value(make_float2(*x0_ptr, *x1_ptr)); - *x0_ptr = result.x; *x1_ptr = result.y; -} - -__device__ __forceinline__ void softmax_frag_exp2_cast( - float* sv, uint32_t* pv, int use_emu) -{ - #pragma unroll - for (int j = 0; j < 16; j++) { - if (use_emu && j >= 12) - ex2_emulation_f32x2(&sv[j*2], &sv[j*2+1]); - else { - sv[j*2] = approx_exp2(sv[j*2]); - sv[j*2+1] = approx_exp2(sv[j*2+1]); - } - } - #pragma unroll - for (int j = 0; j < 16; j++) { - __nv_bfloat162 bf = __float22bfloat162_rn({sv[j*2], sv[j*2+1]}); - pv[j] = reinterpret_cast(bf); - } -} - - - -__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { - const float2* sv2 = reinterpret_cast(sv); - #pragma unroll - for (int j = 0; j < 16; j++) { - asm("add.f32x2 %0, %1, %2;" - : "+l"(reinterpret_cast(*acc)) - : "l"(reinterpret_cast(*acc)), - "l"(reinterpret_cast(sv2[j]))); - } -} - - -__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { - asm("mul.rn.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { - asm("mul.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { - asm("add.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { - asm("add.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { - asm("sub.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { - asm("sub.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("add.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { - float2 r; - asm("sub.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("sub.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ void fma_scale_x32( - float* sv, const float2* scale2, const float2* neg_max2) -{ - float2* sv_2 = reinterpret_cast(sv); - #pragma unroll - for (int j = 0; j < 16; j++) - fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); -} - -__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("mul.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - - -// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) - -__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("add.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("add.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("add.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("add.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("add.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("add.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("add.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - - -__device__ __forceinline__ void fence_async_shared() { - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); -} - - -__device__ __forceinline__ uint64_t make_smem_desc(int addr) { - const int SBO = 1024; - return desc_encode(addr) - | (desc_encode(SBO) << 32ULL) - | (1ULL << 46ULL) - | (2ULL << 61ULL); -} - - -__device__ __forceinline__ void tma_2d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3}], [%4];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_4d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.4d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5}], [%6];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr) { - // Canonical .shared::cta form for non-multicast gather4, matching - // trtllm-gen / cuda_ptx and the PTX ISA qualifier order - // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, - // .shared::cluster is reserved for the multicast variant (ctaMask). - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { - asm volatile( - "tcgen05.commit.cta_group::1.mbarrier::arrive::one" - ".shared::cluster.b64 [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { - uint32_t result; - asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" - : "=r"(result) : "r"(val)); - return result; -} - -extern "C" { - -__global__ __launch_bounds__(512, 1) void -kernel_cake_dsv4_3e5a644d61420c9c7ea1(CakeTensorMap const* tmap_q, CakeTensorMap const* tmap_swa_kv, CakeTensorMap const* tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int num_query_tokens, int has_sinks) -{ - const int tid = threadIdx.x; - const int warp = make_warp_uniform(tid / 32); - const int lane = tid % 32; - - extern __shared__ __align__(1024) char smem_raw[]; - int smem; - smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); - - const int mbar_base = smem; - #define q_full_addr (mbar_base + 0) - #define q_empty_addr (mbar_base + 8) - #define kv_full_addr (mbar_base + 16) - #define kv_empty_addr (mbar_base + 48) - #define s_full_addr (mbar_base + 80) - #define p_full_addr (mbar_base + 96) - #define stats_addr (mbar_base + 112) - #define stats_empty_addr (mbar_base + 128) - #define o_empty_addr (mbar_base + 144) - #define o_full_addr (mbar_base + 152) - #define page_full_addr (mbar_base + 160) - #define page_empty_addr (mbar_base + 208) - #define work_full_addr (mbar_base + 256) - #define work_empty_addr (mbar_base + 272) - #define throttle_full_addr (mbar_base + 288) - #define throttle_empty_addr (mbar_base + 304) - - const int bid = blockIdx.x; - const int num_bids = gridDim.x; - - const int cta_rank = 0; - if (tid == 0) { - asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_q)) : "memory"); - asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_swa_kv)) : "memory"); - asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_compressed_kv)) : "memory"); - } - __syncthreads(); - - - // Kernel setup ops - __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); - const int smem_q_addr = smem + 1024; - __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); - const int smem_kv_addr = smem + 66560; - __nv_bfloat16* smem_v = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); - const int smem_v_addr = smem + 66560; - int* smem_page_offsets = reinterpret_cast(smem_raw + 197632); - const int smem_page_offsets_addr = smem + 197632; - unsigned int* work_response = reinterpret_cast(smem_raw + 203776); - const int work_response_addr = smem + 203776; - - // Mbarrier init (16 pipeline groups, 0 ordered-sequence groups, 40 barriers) - // Mbarriers at smem_raw[0..320) - - if (warp == 0) { - uint32_t leader = elect_sync(); - if (leader) { - // q_full: 1 barriers, init_count=1 - mbarrier_init(smem + 0, 1); - // q_empty: 1 barriers, init_count=1 - mbarrier_init(smem + 8, 1); - // --- pipeline 'kv_pipe' --- - // kv_full: 4 barriers, init_count=1 - mbarrier_init(smem + 16, 1); - mbarrier_init(smem + 24, 1); - mbarrier_init(smem + 32, 1); - mbarrier_init(smem + 40, 1); - // kv_empty: 4 barriers, init_count=1 - mbarrier_init(smem + 48, 1); - mbarrier_init(smem + 56, 1); - mbarrier_init(smem + 64, 1); - mbarrier_init(smem + 72, 1); - // s_full: 2 barriers, init_count=1 - mbarrier_init(smem + 80, 1); - mbarrier_init(smem + 88, 1); - // p_full: 2 barriers, init_count=128 - mbarrier_init(smem + 96, 128); - mbarrier_init(smem + 104, 128); - // stats: 2 barriers, init_count=128 - mbarrier_init(smem + 112, 128); - mbarrier_init(smem + 120, 128); - // stats_empty: 2 barriers, init_count=128 - mbarrier_init(smem + 128, 128); - mbarrier_init(smem + 136, 128); - // o_empty: 1 barriers, init_count=128 - mbarrier_init(smem + 144, 128); - // o_full: 1 barriers, init_count=1 - mbarrier_init(smem + 152, 1); - // --- pipeline 'page_pipe' --- - // page_full: 6 barriers, init_count=32 - mbarrier_init(smem + 160, 32); - mbarrier_init(smem + 168, 32); - mbarrier_init(smem + 176, 32); - mbarrier_init(smem + 184, 32); - mbarrier_init(smem + 192, 32); - mbarrier_init(smem + 200, 32); - // page_empty: 6 barriers, init_count=128 - mbarrier_init(smem + 208, 128); - mbarrier_init(smem + 216, 128); - mbarrier_init(smem + 224, 128); - mbarrier_init(smem + 232, 128); - mbarrier_init(smem + 240, 128); - mbarrier_init(smem + 248, 128); - // --- pipeline 'work_pipe' --- - // work_full: 2 barriers, init_count=1 - mbarrier_init(smem + 256, 1); - mbarrier_init(smem + 264, 1); - // work_empty: 2 barriers, init_count=512 - mbarrier_init(smem + 272, 512); - mbarrier_init(smem + 280, 512); - // --- pipeline 'throttle_pipe' --- - // throttle_full: 2 barriers, init_count=128 - mbarrier_init(smem + 288, 128); - mbarrier_init(smem + 296, 128); - // throttle_empty: 2 barriers, init_count=32 - mbarrier_init(smem + 304, 32); - mbarrier_init(smem + 312, 32); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - - __syncwarp(); - - // TMEM alloc (512 columns, 512 used) - volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 320); - if (warp == 0) { - int _tmem_hold = smem + 320; - asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); - __syncwarp(); - asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); - } - - __syncthreads(); - asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); - - const int taddr = tmem_addr_storage[0]; - - // Kernel post-init ops - const int tmem_tmem = taddr; - - // ---- Ordered hardware-WG register redistribution ---- - // Inc phase consumes the registers released above. - if (warp >= 8 && warp <= 11) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 136;"); - } - - // ---- Role: softmax ---- - if (warp <= 3) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 152;"); - { // softmax_main - asm volatile("griddepcontrol.wait;" ::: "memory"); - float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; - unsigned int work_stage = 0; - unsigned int query_idx = blockIdx.x; - unsigned int global_tile = 0; - unsigned int stats_cursor = 0; - const int warp_in_compute = warp; - const int tmem_row_origin = warp_in_compute * 32; - const int logical_row_origin = warp_in_compute * 16; - const int my_row = logical_row_origin + lane % 16; - const int col_half = lane / 16; - unsigned int _phase_work_full = 0; - #pragma unroll 1 - for (unsigned int work_iter = 0; work_iter < num_query_tokens; work_iter++) { - if (query_idx >= (unsigned int)num_query_tokens) { - break; - } - int _max_4 = ((sparse_topk_lens[query_idx] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx] + sparse_topk_lens_offset) : (0)); - int _min_2 = ((_max_4) < (sparse_topk) ? (_max_4) : (sparse_topk)); - int active_topk = _min_2; - int _max_5 = (((active_topk + 128 - 1) / 128) > (1) ? ((active_topk + 128 - 1) / 128) : (1)); - int num_loop_steps = _max_5; - float row_max = -CAKE_INF; - float row_sum = 0.0f; - if (has_sinks != 0 && my_row < num_heads) { - row_max = sinks[my_row] * 1.4426950408889634f / softmax_scale_log2; - row_sum = 1.0f; - } - if (num_loop_steps > 0) { - unsigned int first_stats_stage = stats_cursor & 1; - unsigned int first_stats_empty_phase = stats_cursor >> 1 & 1 ^ 1; - mbarrier_wait(stats_empty_addr + (first_stats_stage) * 8, first_stats_empty_phase); - } - #pragma unroll 1 - for (int tile = 0; tile < num_loop_steps; tile++) { - unsigned int sync_stage = global_tile & 1; - unsigned int sync_phase = global_tile >> 1 & 1; - mbarrier_wait(s_full_addr + (sync_stage) * 8, sync_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - int score_col = ((sync_stage != 0) ? 128 : 0); - int score_addr = taddr + (unsigned int)score_col + (unsigned int)(tmem_row_origin << 16); - float _tmem_load_0[64]; - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_0[0])), "=r"(*reinterpret_cast(&_tmem_load_0[1])), "=r"(*reinterpret_cast(&_tmem_load_0[2])), "=r"(*reinterpret_cast(&_tmem_load_0[3])), "=r"(*reinterpret_cast(&_tmem_load_0[4])), "=r"(*reinterpret_cast(&_tmem_load_0[5])), "=r"(*reinterpret_cast(&_tmem_load_0[6])), "=r"(*reinterpret_cast(&_tmem_load_0[7])), "=r"(*reinterpret_cast(&_tmem_load_0[8])), "=r"(*reinterpret_cast(&_tmem_load_0[9])), "=r"(*reinterpret_cast(&_tmem_load_0[10])), "=r"(*reinterpret_cast(&_tmem_load_0[11])), "=r"(*reinterpret_cast(&_tmem_load_0[12])), "=r"(*reinterpret_cast(&_tmem_load_0[13])), "=r"(*reinterpret_cast(&_tmem_load_0[14])), "=r"(*reinterpret_cast(&_tmem_load_0[15])), "=r"(*reinterpret_cast(&_tmem_load_0[16])), "=r"(*reinterpret_cast(&_tmem_load_0[17])), "=r"(*reinterpret_cast(&_tmem_load_0[18])), "=r"(*reinterpret_cast(&_tmem_load_0[19])), "=r"(*reinterpret_cast(&_tmem_load_0[20])), "=r"(*reinterpret_cast(&_tmem_load_0[21])), "=r"(*reinterpret_cast(&_tmem_load_0[22])), "=r"(*reinterpret_cast(&_tmem_load_0[23])), "=r"(*reinterpret_cast(&_tmem_load_0[24])), "=r"(*reinterpret_cast(&_tmem_load_0[25])), "=r"(*reinterpret_cast(&_tmem_load_0[26])), "=r"(*reinterpret_cast(&_tmem_load_0[27])), "=r"(*reinterpret_cast(&_tmem_load_0[28])), "=r"(*reinterpret_cast(&_tmem_load_0[29])), "=r"(*reinterpret_cast(&_tmem_load_0[30])), "=r"(*reinterpret_cast(&_tmem_load_0[31])) - : "r"(score_addr)); - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_0[32])), "=r"(*reinterpret_cast(&_tmem_load_0[33])), "=r"(*reinterpret_cast(&_tmem_load_0[34])), "=r"(*reinterpret_cast(&_tmem_load_0[35])), "=r"(*reinterpret_cast(&_tmem_load_0[36])), "=r"(*reinterpret_cast(&_tmem_load_0[37])), "=r"(*reinterpret_cast(&_tmem_load_0[38])), "=r"(*reinterpret_cast(&_tmem_load_0[39])), "=r"(*reinterpret_cast(&_tmem_load_0[40])), "=r"(*reinterpret_cast(&_tmem_load_0[41])), "=r"(*reinterpret_cast(&_tmem_load_0[42])), "=r"(*reinterpret_cast(&_tmem_load_0[43])), "=r"(*reinterpret_cast(&_tmem_load_0[44])), "=r"(*reinterpret_cast(&_tmem_load_0[45])), "=r"(*reinterpret_cast(&_tmem_load_0[46])), "=r"(*reinterpret_cast(&_tmem_load_0[47])), "=r"(*reinterpret_cast(&_tmem_load_0[48])), "=r"(*reinterpret_cast(&_tmem_load_0[49])), "=r"(*reinterpret_cast(&_tmem_load_0[50])), "=r"(*reinterpret_cast(&_tmem_load_0[51])), "=r"(*reinterpret_cast(&_tmem_load_0[52])), "=r"(*reinterpret_cast(&_tmem_load_0[53])), "=r"(*reinterpret_cast(&_tmem_load_0[54])), "=r"(*reinterpret_cast(&_tmem_load_0[55])), "=r"(*reinterpret_cast(&_tmem_load_0[56])), "=r"(*reinterpret_cast(&_tmem_load_0[57])), "=r"(*reinterpret_cast(&_tmem_load_0[58])), "=r"(*reinterpret_cast(&_tmem_load_0[59])), "=r"(*reinterpret_cast(&_tmem_load_0[60])), "=r"(*reinterpret_cast(&_tmem_load_0[61])), "=r"(*reinterpret_cast(&_tmem_load_0[62])), "=r"(*reinterpret_cast(&_tmem_load_0[63])) - : "r"(score_addr + 32)); - int valid_cols = active_topk - tile * 128 - col_half * 64; - if (valid_cols < 0) { - valid_cols = 0; - } - if (valid_cols > 64) { - valid_cols = 64; - } - uint32_t _slice_lo_mask_0; - { - int _lim_0 = valid_cols; - if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } - else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); - } - } - if (!(_slice_lo_mask_0 & (1u << 0))) _tmem_load_0[0] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 1))) _tmem_load_0[1] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 2))) _tmem_load_0[2] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 3))) _tmem_load_0[3] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 4))) _tmem_load_0[4] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 5))) _tmem_load_0[5] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 6))) _tmem_load_0[6] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 7))) _tmem_load_0[7] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 8))) _tmem_load_0[8] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 9))) _tmem_load_0[9] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 10))) _tmem_load_0[10] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 11))) _tmem_load_0[11] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 12))) _tmem_load_0[12] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 13))) _tmem_load_0[13] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 14))) _tmem_load_0[14] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 15))) _tmem_load_0[15] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 16))) _tmem_load_0[16] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 17))) _tmem_load_0[17] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 18))) _tmem_load_0[18] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 19))) _tmem_load_0[19] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 20))) _tmem_load_0[20] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 21))) _tmem_load_0[21] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 22))) _tmem_load_0[22] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 23))) _tmem_load_0[23] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 24))) _tmem_load_0[24] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 25))) _tmem_load_0[25] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 26))) _tmem_load_0[26] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 27))) _tmem_load_0[27] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 28))) _tmem_load_0[28] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 29))) _tmem_load_0[29] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 30))) _tmem_load_0[30] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 31))) _tmem_load_0[31] = -CAKE_INF; - uint32_t _slice_lo_mask_1; - { - int _lim_1 = valid_cols - 32; - if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } - else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); - } - } - if (!(_slice_lo_mask_1 & (1u << 0))) _tmem_load_0[32] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 1))) _tmem_load_0[33] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 2))) _tmem_load_0[34] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 3))) _tmem_load_0[35] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 4))) _tmem_load_0[36] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 5))) _tmem_load_0[37] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 6))) _tmem_load_0[38] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 7))) _tmem_load_0[39] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 8))) _tmem_load_0[40] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 9))) _tmem_load_0[41] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 10))) _tmem_load_0[42] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 11))) _tmem_load_0[43] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 12))) _tmem_load_0[44] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 13))) _tmem_load_0[45] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 14))) _tmem_load_0[46] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 15))) _tmem_load_0[47] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 16))) _tmem_load_0[48] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 17))) _tmem_load_0[49] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 18))) _tmem_load_0[50] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 19))) _tmem_load_0[51] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 20))) _tmem_load_0[52] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 21))) _tmem_load_0[53] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 22))) _tmem_load_0[54] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 23))) _tmem_load_0[55] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 24))) _tmem_load_0[56] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 25))) _tmem_load_0[57] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 26))) _tmem_load_0[58] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 27))) _tmem_load_0[59] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 28))) _tmem_load_0[60] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 29))) _tmem_load_0[61] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 30))) _tmem_load_0[62] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 31))) _tmem_load_0[63] = -CAKE_INF; - float2 _reg_reduce_max2_2 = {-CAKE_INF, -CAKE_INF}; - row_max_x32_accum(&_tmem_load_0[0], _reg_reduce_max2_2); - row_max_x32_accum(&_tmem_load_0[32], _reg_reduce_max2_2); - float _tmem_load_0_max = row_max_reduce(_reg_reduce_max2_2); - float tile_max = _tmem_load_0_max; - float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, tile_max, 16); - float _max_6 = max_noftz(tile_max, _shfl_xor_0); - tile_max = _max_6; - float old_max = row_max; - float _max_7 = max_noftz(row_max, tile_max); - float new_max = _max_7; - float _fma_0 = __fmaf_rn(old_max, softmax_scale_log2, (-new_max) * softmax_scale_log2); - float delta = _fma_0; - float _exp2_0 = approx_exp2(delta); - float acc_scale = ((old_max > -CAKE_INF) ? _exp2_0 : 1.0f); - unsigned int stats_stage = stats_cursor & 1; - int stats_col = ((stats_stage != 0) ? 128 : 0); - float max_stats[2]; - max_stats[0] = old_max; - max_stats[1] = new_max; - asm volatile( - "tcgen05.st.sync.aligned.32x32b.x2.b32" - " [%0], {%1, %2};" - :: "r"(taddr + (unsigned int)stats_col), "f"(max_stats[0]), "f"(max_stats[1])); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - mbarrier_arrive(stats_addr + (stats_stage) * 8); - stats_cursor = stats_cursor + 1; - unsigned int next_stats_stage = stats_cursor & 1; - unsigned int next_stats_empty_phase = stats_cursor >> 1 & 1 ^ 1; - mbarrier_wait(stats_empty_addr + (next_stats_stage) * 8, next_stats_empty_phase); - bool empty_row = new_max == -CAKE_INF; - float safe_max = ((empty_row) ? 0.0f : new_max); - float max_scaled = safe_max * softmax_scale_log2; - const float2 _fma_b2_3 = {softmax_scale_log2, softmax_scale_log2}; - const float2 _fma_c2_4 = {-max_scaled, -max_scaled}; - #pragma unroll - for (int _lf = 0; _lf < 32; _lf++) - fma_f32x2_inplace(&reinterpret_cast(_tmem_load_0)[_lf], _fma_b2_3, _fma_c2_4); - #pragma unroll - for (int _le = 0; _le < 64; _le++) { - _tmem_load_0[_le] = approx_exp2(_tmem_load_0[_le]); - } - float2 _reg_reduce_sum2_5 = make_float2(0.0f, 0.0f); - softmax_block_sum(&_tmem_load_0[0], &_reg_reduce_sum2_5); - softmax_block_sum(&_tmem_load_0[32], &_reg_reduce_sum2_5); - float _tmem_load_0_sum = _reg_reduce_sum2_5.x + _reg_reduce_sum2_5.y; - float block_sum = _tmem_load_0_sum; - float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, block_sum, 16); - block_sum = block_sum + _shfl_xor_1; - float _fma_1 = __fmaf_rn(row_sum, acc_scale, block_sum); - row_sum = _fma_1; - row_max = new_max; - if (tile == num_loop_steps - 1) { - float final_stats[2]; - float _max_8 = max_noftz(row_sum, 1.1754943508222875e-38f); - final_stats[0] = _max_8; - final_stats[1] = row_max; - int final_stats_col = ((next_stats_stage != 0) ? 128 : 0); - asm volatile( - "tcgen05.st.sync.aligned.32x32b.x2.b32" - " [%0], {%1, %2};" - :: "r"(taddr + (unsigned int)final_stats_col), "f"(final_stats[0]), "f"(final_stats[1])); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - mbarrier_arrive(stats_addr + (next_stats_stage) * 8); - stats_cursor = stats_cursor + 1; - unsigned int tail_stats_stage = stats_cursor & 1; - unsigned int tail_stats_empty_phase = stats_cursor >> 1 & 1 ^ 1; - mbarrier_wait(stats_empty_addr + (tail_stats_stage) * 8, tail_stats_empty_phase); - mbarrier_arrive(stats_addr + (tail_stats_stage) * 8); - stats_cursor = stats_cursor + 1; - } - unsigned int packed_p[32]; - #pragma unroll - for (int _lp = 0; _lp < 32; _lp++) { - __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_0[_lp*2 + 0], _tmem_load_0[_lp*2+1 + 0])); - packed_p[_lp] = *(uint32_t*)&_bf2; - } - int p_addr = taddr + (unsigned int)score_col + 32 + (unsigned int)(tmem_row_origin << 16); - asm volatile( - "tcgen05.st.sync.aligned.16x32bx2.x16.b32" - " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" - :: "r"(p_addr), "r"(*reinterpret_cast(&packed_p[0])), "r"(*reinterpret_cast(&packed_p[1])), "r"(*reinterpret_cast(&packed_p[2])), "r"(*reinterpret_cast(&packed_p[3])), "r"(*reinterpret_cast(&packed_p[4])), "r"(*reinterpret_cast(&packed_p[5])), "r"(*reinterpret_cast(&packed_p[6])), "r"(*reinterpret_cast(&packed_p[7])), "r"(*reinterpret_cast(&packed_p[8])), "r"(*reinterpret_cast(&packed_p[9])), "r"(*reinterpret_cast(&packed_p[10])), "r"(*reinterpret_cast(&packed_p[11])), "r"(*reinterpret_cast(&packed_p[12])), "r"(*reinterpret_cast(&packed_p[13])), "r"(*reinterpret_cast(&packed_p[14])), "r"(*reinterpret_cast(&packed_p[15]))); - asm volatile( - "tcgen05.st.sync.aligned.16x32bx2.x16.b32" - " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" - :: "r"(p_addr + 16), "r"(*reinterpret_cast(&(packed_p + 16)[0])), "r"(*reinterpret_cast(&(packed_p + 16)[1])), "r"(*reinterpret_cast(&(packed_p + 16)[2])), "r"(*reinterpret_cast(&(packed_p + 16)[3])), "r"(*reinterpret_cast(&(packed_p + 16)[4])), "r"(*reinterpret_cast(&(packed_p + 16)[5])), "r"(*reinterpret_cast(&(packed_p + 16)[6])), "r"(*reinterpret_cast(&(packed_p + 16)[7])), "r"(*reinterpret_cast(&(packed_p + 16)[8])), "r"(*reinterpret_cast(&(packed_p + 16)[9])), "r"(*reinterpret_cast(&(packed_p + 16)[10])), "r"(*reinterpret_cast(&(packed_p + 16)[11])), "r"(*reinterpret_cast(&(packed_p + 16)[12])), "r"(*reinterpret_cast(&(packed_p + 16)[13])), "r"(*reinterpret_cast(&(packed_p + 16)[14])), "r"(*reinterpret_cast(&(packed_p + 16)[15]))); - int p_addr_high = p_addr + 1048576; - asm volatile( - "tcgen05.st.sync.aligned.16x32bx2.x16.b32" - " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" - :: "r"(p_addr_high), "r"(*reinterpret_cast(&packed_p[0])), "r"(*reinterpret_cast(&packed_p[1])), "r"(*reinterpret_cast(&packed_p[2])), "r"(*reinterpret_cast(&packed_p[3])), "r"(*reinterpret_cast(&packed_p[4])), "r"(*reinterpret_cast(&packed_p[5])), "r"(*reinterpret_cast(&packed_p[6])), "r"(*reinterpret_cast(&packed_p[7])), "r"(*reinterpret_cast(&packed_p[8])), "r"(*reinterpret_cast(&packed_p[9])), "r"(*reinterpret_cast(&packed_p[10])), "r"(*reinterpret_cast(&packed_p[11])), "r"(*reinterpret_cast(&packed_p[12])), "r"(*reinterpret_cast(&packed_p[13])), "r"(*reinterpret_cast(&packed_p[14])), "r"(*reinterpret_cast(&packed_p[15]))); - asm volatile( - "tcgen05.st.sync.aligned.16x32bx2.x16.b32" - " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" - :: "r"(p_addr_high + 16), "r"(*reinterpret_cast(&(packed_p + 16)[0])), "r"(*reinterpret_cast(&(packed_p + 16)[1])), "r"(*reinterpret_cast(&(packed_p + 16)[2])), "r"(*reinterpret_cast(&(packed_p + 16)[3])), "r"(*reinterpret_cast(&(packed_p + 16)[4])), "r"(*reinterpret_cast(&(packed_p + 16)[5])), "r"(*reinterpret_cast(&(packed_p + 16)[6])), "r"(*reinterpret_cast(&(packed_p + 16)[7])), "r"(*reinterpret_cast(&(packed_p + 16)[8])), "r"(*reinterpret_cast(&(packed_p + 16)[9])), "r"(*reinterpret_cast(&(packed_p + 16)[10])), "r"(*reinterpret_cast(&(packed_p + 16)[11])), "r"(*reinterpret_cast(&(packed_p + 16)[12])), "r"(*reinterpret_cast(&(packed_p + 16)[13])), "r"(*reinterpret_cast(&(packed_p + 16)[14])), "r"(*reinterpret_cast(&(packed_p + 16)[15]))); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(p_full_addr + (sync_stage) * 8); - global_tile = global_tile + 1; - } - mbarrier_wait(work_full_addr + (work_stage) * 8, _phase_work_full); - uint32_t _clc_valid_2 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_2) - : "r"(work_response_addr + work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_2 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_2) - : "r"(work_response_addr + work_stage * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(work_empty_addr + (work_stage) * 8); - work_stage += 1; - if (work_stage == 2) { work_stage = 0; _phase_work_full ^= 1; } - unsigned int valid = _clc_valid_2; - query_idx = _clc_ctaid_2; - if (valid == 0) { - break; - } - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - } - } - // ---- Role: correction ---- - if (warp >= 4 && warp <= 7) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 144;"); - { // correction_main - asm volatile("griddepcontrol.wait;" ::: "memory"); - float softmax_scale_log2_1 = bmm1_scale[0] * 1.4426950408889634f; - float output_scale = bmm2_scale[0]; - unsigned int work_stage_1 = 0; - unsigned int query_idx_1 = blockIdx.x; - unsigned int global_tile_1 = 0; - unsigned int stats_cursor_1 = 0; - const int warp_in_role = warp - 4; - const int tmem_row_origin_1 = warp_in_role * 32; - const int logical_row_origin_1 = warp_in_role * 16; - const int my_row_1 = logical_row_origin_1 + lane % 16; - const int col_half_1 = lane / 16; - unsigned int _phase_work_full_1 = 0; - #pragma unroll 1 - for (unsigned int work_iter_1 = 0; work_iter_1 < num_query_tokens; work_iter_1++) { - if (query_idx_1 >= (unsigned int)num_query_tokens) { - break; - } - int _max_9 = ((sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) : (0)); - int _min_3 = ((_max_9) < (sparse_topk) ? (_max_9) : (sparse_topk)); - int active_topk_1 = _min_3; - int _max_10 = (((active_topk_1 + 128 - 1) / 128) > (1) ? ((active_topk_1 + 128 - 1) / 128) : (1)); - int num_loop_steps_1 = _max_10; - unsigned int first_stats_stage_1 = stats_cursor_1 & 1; - unsigned int first_stats_phase = stats_cursor_1 >> 1 & 1; - mbarrier_wait(stats_addr + (first_stats_stage_1) * 8, first_stats_phase); - mbarrier_arrive(stats_empty_addr + (first_stats_stage_1) * 8); - stats_cursor_1 = stats_cursor_1 + 1; - global_tile_1 = global_tile_1 + 1; - #pragma unroll 1 - for (int tile_1 = 0; tile_1 < num_loop_steps_1 - 1; tile_1++) { - unsigned int stats_stage_1 = stats_cursor_1 & 1; - unsigned int stats_phase = stats_cursor_1 >> 1 & 1; - mbarrier_wait(stats_addr + (stats_stage_1) * 8, stats_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - int stats_col_1 = ((stats_stage_1 != 0) ? 128 : 0); - float _tmem_load_1[2]; - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x2.b32" - " {%0, %1}, [%2];" - : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]) - : "r"(taddr + (unsigned int)stats_col_1)); - mbarrier_arrive(stats_empty_addr + (stats_stage_1) * 8); - stats_cursor_1 = stats_cursor_1 + 1; - float max_diff = _tmem_load_1[0] - _tmem_load_1[1]; - float _exp2_1 = approx_exp2(softmax_scale_log2_1 * max_diff); - float acc_scale_1 = ((max_diff != 0.0f) ? _exp2_1 : 1.0f); - mbarrier_wait(o_full_addr, global_tile_1 - 1 & 1); - asm volatile("tcgen05.fence::after_thread_sync;"); - int _vote_0 = __any_sync(0xFFFFFFFF, acc_scale_1 < 1.0f); - int any_rescale = _vote_0; - if (any_rescale != 0) { - #pragma unroll - for (int v_stage = 0; v_stage < 4; v_stage++) { - int output_addr = taddr + 256 + (unsigned int)((v_stage & 1) * 128) + (unsigned int)((v_stage >> 1) * 16 << 16) + (unsigned int)(tmem_row_origin_1 << 16); - float _tmem_load_2[64]; - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_2[0])), "=r"(*reinterpret_cast(&_tmem_load_2[1])), "=r"(*reinterpret_cast(&_tmem_load_2[2])), "=r"(*reinterpret_cast(&_tmem_load_2[3])), "=r"(*reinterpret_cast(&_tmem_load_2[4])), "=r"(*reinterpret_cast(&_tmem_load_2[5])), "=r"(*reinterpret_cast(&_tmem_load_2[6])), "=r"(*reinterpret_cast(&_tmem_load_2[7])), "=r"(*reinterpret_cast(&_tmem_load_2[8])), "=r"(*reinterpret_cast(&_tmem_load_2[9])), "=r"(*reinterpret_cast(&_tmem_load_2[10])), "=r"(*reinterpret_cast(&_tmem_load_2[11])), "=r"(*reinterpret_cast(&_tmem_load_2[12])), "=r"(*reinterpret_cast(&_tmem_load_2[13])), "=r"(*reinterpret_cast(&_tmem_load_2[14])), "=r"(*reinterpret_cast(&_tmem_load_2[15])), "=r"(*reinterpret_cast(&_tmem_load_2[16])), "=r"(*reinterpret_cast(&_tmem_load_2[17])), "=r"(*reinterpret_cast(&_tmem_load_2[18])), "=r"(*reinterpret_cast(&_tmem_load_2[19])), "=r"(*reinterpret_cast(&_tmem_load_2[20])), "=r"(*reinterpret_cast(&_tmem_load_2[21])), "=r"(*reinterpret_cast(&_tmem_load_2[22])), "=r"(*reinterpret_cast(&_tmem_load_2[23])), "=r"(*reinterpret_cast(&_tmem_load_2[24])), "=r"(*reinterpret_cast(&_tmem_load_2[25])), "=r"(*reinterpret_cast(&_tmem_load_2[26])), "=r"(*reinterpret_cast(&_tmem_load_2[27])), "=r"(*reinterpret_cast(&_tmem_load_2[28])), "=r"(*reinterpret_cast(&_tmem_load_2[29])), "=r"(*reinterpret_cast(&_tmem_load_2[30])), "=r"(*reinterpret_cast(&_tmem_load_2[31])) - : "r"(output_addr)); - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_2[32])), "=r"(*reinterpret_cast(&_tmem_load_2[33])), "=r"(*reinterpret_cast(&_tmem_load_2[34])), "=r"(*reinterpret_cast(&_tmem_load_2[35])), "=r"(*reinterpret_cast(&_tmem_load_2[36])), "=r"(*reinterpret_cast(&_tmem_load_2[37])), "=r"(*reinterpret_cast(&_tmem_load_2[38])), "=r"(*reinterpret_cast(&_tmem_load_2[39])), "=r"(*reinterpret_cast(&_tmem_load_2[40])), "=r"(*reinterpret_cast(&_tmem_load_2[41])), "=r"(*reinterpret_cast(&_tmem_load_2[42])), "=r"(*reinterpret_cast(&_tmem_load_2[43])), "=r"(*reinterpret_cast(&_tmem_load_2[44])), "=r"(*reinterpret_cast(&_tmem_load_2[45])), "=r"(*reinterpret_cast(&_tmem_load_2[46])), "=r"(*reinterpret_cast(&_tmem_load_2[47])), "=r"(*reinterpret_cast(&_tmem_load_2[48])), "=r"(*reinterpret_cast(&_tmem_load_2[49])), "=r"(*reinterpret_cast(&_tmem_load_2[50])), "=r"(*reinterpret_cast(&_tmem_load_2[51])), "=r"(*reinterpret_cast(&_tmem_load_2[52])), "=r"(*reinterpret_cast(&_tmem_load_2[53])), "=r"(*reinterpret_cast(&_tmem_load_2[54])), "=r"(*reinterpret_cast(&_tmem_load_2[55])), "=r"(*reinterpret_cast(&_tmem_load_2[56])), "=r"(*reinterpret_cast(&_tmem_load_2[57])), "=r"(*reinterpret_cast(&_tmem_load_2[58])), "=r"(*reinterpret_cast(&_tmem_load_2[59])), "=r"(*reinterpret_cast(&_tmem_load_2[60])), "=r"(*reinterpret_cast(&_tmem_load_2[61])), "=r"(*reinterpret_cast(&_tmem_load_2[62])), "=r"(*reinterpret_cast(&_tmem_load_2[63])) - : "r"(output_addr + 32)); - const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; - #pragma unroll - for (int _ls = 0; _ls < 32; _ls++) - mul_f32x2_inplace(&reinterpret_cast(_tmem_load_2)[_ls], _scale2_0); - asm volatile( - "tcgen05.st.sync.aligned.16x32bx2.x64.b32" - " [%0], 64, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64};" - :: "r"(output_addr), "r"(*reinterpret_cast(&_tmem_load_2[0])), "r"(*reinterpret_cast(&_tmem_load_2[1])), "r"(*reinterpret_cast(&_tmem_load_2[2])), "r"(*reinterpret_cast(&_tmem_load_2[3])), "r"(*reinterpret_cast(&_tmem_load_2[4])), "r"(*reinterpret_cast(&_tmem_load_2[5])), "r"(*reinterpret_cast(&_tmem_load_2[6])), "r"(*reinterpret_cast(&_tmem_load_2[7])), "r"(*reinterpret_cast(&_tmem_load_2[8])), "r"(*reinterpret_cast(&_tmem_load_2[9])), "r"(*reinterpret_cast(&_tmem_load_2[10])), "r"(*reinterpret_cast(&_tmem_load_2[11])), "r"(*reinterpret_cast(&_tmem_load_2[12])), "r"(*reinterpret_cast(&_tmem_load_2[13])), "r"(*reinterpret_cast(&_tmem_load_2[14])), "r"(*reinterpret_cast(&_tmem_load_2[15])), "r"(*reinterpret_cast(&_tmem_load_2[16])), "r"(*reinterpret_cast(&_tmem_load_2[17])), "r"(*reinterpret_cast(&_tmem_load_2[18])), "r"(*reinterpret_cast(&_tmem_load_2[19])), "r"(*reinterpret_cast(&_tmem_load_2[20])), "r"(*reinterpret_cast(&_tmem_load_2[21])), "r"(*reinterpret_cast(&_tmem_load_2[22])), "r"(*reinterpret_cast(&_tmem_load_2[23])), "r"(*reinterpret_cast(&_tmem_load_2[24])), "r"(*reinterpret_cast(&_tmem_load_2[25])), "r"(*reinterpret_cast(&_tmem_load_2[26])), "r"(*reinterpret_cast(&_tmem_load_2[27])), "r"(*reinterpret_cast(&_tmem_load_2[28])), "r"(*reinterpret_cast(&_tmem_load_2[29])), "r"(*reinterpret_cast(&_tmem_load_2[30])), "r"(*reinterpret_cast(&_tmem_load_2[31])), "r"(*reinterpret_cast(&_tmem_load_2[32])), "r"(*reinterpret_cast(&_tmem_load_2[33])), "r"(*reinterpret_cast(&_tmem_load_2[34])), "r"(*reinterpret_cast(&_tmem_load_2[35])), "r"(*reinterpret_cast(&_tmem_load_2[36])), "r"(*reinterpret_cast(&_tmem_load_2[37])), "r"(*reinterpret_cast(&_tmem_load_2[38])), "r"(*reinterpret_cast(&_tmem_load_2[39])), "r"(*reinterpret_cast(&_tmem_load_2[40])), "r"(*reinterpret_cast(&_tmem_load_2[41])), "r"(*reinterpret_cast(&_tmem_load_2[42])), "r"(*reinterpret_cast(&_tmem_load_2[43])), "r"(*reinterpret_cast(&_tmem_load_2[44])), "r"(*reinterpret_cast(&_tmem_load_2[45])), "r"(*reinterpret_cast(&_tmem_load_2[46])), "r"(*reinterpret_cast(&_tmem_load_2[47])), "r"(*reinterpret_cast(&_tmem_load_2[48])), "r"(*reinterpret_cast(&_tmem_load_2[49])), "r"(*reinterpret_cast(&_tmem_load_2[50])), "r"(*reinterpret_cast(&_tmem_load_2[51])), "r"(*reinterpret_cast(&_tmem_load_2[52])), "r"(*reinterpret_cast(&_tmem_load_2[53])), "r"(*reinterpret_cast(&_tmem_load_2[54])), "r"(*reinterpret_cast(&_tmem_load_2[55])), "r"(*reinterpret_cast(&_tmem_load_2[56])), "r"(*reinterpret_cast(&_tmem_load_2[57])), "r"(*reinterpret_cast(&_tmem_load_2[58])), "r"(*reinterpret_cast(&_tmem_load_2[59])), "r"(*reinterpret_cast(&_tmem_load_2[60])), "r"(*reinterpret_cast(&_tmem_load_2[61])), "r"(*reinterpret_cast(&_tmem_load_2[62])), "r"(*reinterpret_cast(&_tmem_load_2[63]))); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - } - } - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(o_empty_addr); - global_tile_1 = global_tile_1 + 1; - } - unsigned int final_stats_stage = stats_cursor_1 & 1; - unsigned int final_stats_phase = stats_cursor_1 >> 1 & 1; - mbarrier_wait(stats_addr + (final_stats_stage) * 8, final_stats_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - int final_stats_col_1 = ((final_stats_stage != 0) ? 128 : 0); - float _tmem_load_3[2]; - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x2.b32" - " {%0, %1}, [%2];" - : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]) - : "r"(taddr + (unsigned int)final_stats_col_1)); - mbarrier_arrive(stats_empty_addr + (final_stats_stage) * 8); - stats_cursor_1 = stats_cursor_1 + 1; - unsigned int tail_stats_stage_1 = stats_cursor_1 & 1; - unsigned int tail_stats_phase = stats_cursor_1 >> 1 & 1; - mbarrier_wait(stats_addr + (tail_stats_stage_1) * 8, tail_stats_phase); - mbarrier_arrive(stats_empty_addr + (tail_stats_stage_1) * 8); - stats_cursor_1 = stats_cursor_1 + 1; - mbarrier_wait(o_full_addr, global_tile_1 - 1 & 1); - asm volatile("tcgen05.fence::after_thread_sync;"); - float final_sum = _tmem_load_3[0]; - float _rcp_0 = approx_rcp(final_sum); - float inv_sum = _rcp_0; - int output_base = (query_idx_1 * (unsigned int)num_heads + (unsigned int)my_row_1) * 512; - #pragma unroll - for (int v_stage_1 = 0; v_stage_1 < 4; v_stage_1++) { - int output_addr_1 = taddr + 256 + (unsigned int)((v_stage_1 & 1) * 128) + (unsigned int)((v_stage_1 >> 1) * 16 << 16) + (unsigned int)(tmem_row_origin_1 << 16); - float _tmem_load_4[64]; - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_4[0])), "=r"(*reinterpret_cast(&_tmem_load_4[1])), "=r"(*reinterpret_cast(&_tmem_load_4[2])), "=r"(*reinterpret_cast(&_tmem_load_4[3])), "=r"(*reinterpret_cast(&_tmem_load_4[4])), "=r"(*reinterpret_cast(&_tmem_load_4[5])), "=r"(*reinterpret_cast(&_tmem_load_4[6])), "=r"(*reinterpret_cast(&_tmem_load_4[7])), "=r"(*reinterpret_cast(&_tmem_load_4[8])), "=r"(*reinterpret_cast(&_tmem_load_4[9])), "=r"(*reinterpret_cast(&_tmem_load_4[10])), "=r"(*reinterpret_cast(&_tmem_load_4[11])), "=r"(*reinterpret_cast(&_tmem_load_4[12])), "=r"(*reinterpret_cast(&_tmem_load_4[13])), "=r"(*reinterpret_cast(&_tmem_load_4[14])), "=r"(*reinterpret_cast(&_tmem_load_4[15])), "=r"(*reinterpret_cast(&_tmem_load_4[16])), "=r"(*reinterpret_cast(&_tmem_load_4[17])), "=r"(*reinterpret_cast(&_tmem_load_4[18])), "=r"(*reinterpret_cast(&_tmem_load_4[19])), "=r"(*reinterpret_cast(&_tmem_load_4[20])), "=r"(*reinterpret_cast(&_tmem_load_4[21])), "=r"(*reinterpret_cast(&_tmem_load_4[22])), "=r"(*reinterpret_cast(&_tmem_load_4[23])), "=r"(*reinterpret_cast(&_tmem_load_4[24])), "=r"(*reinterpret_cast(&_tmem_load_4[25])), "=r"(*reinterpret_cast(&_tmem_load_4[26])), "=r"(*reinterpret_cast(&_tmem_load_4[27])), "=r"(*reinterpret_cast(&_tmem_load_4[28])), "=r"(*reinterpret_cast(&_tmem_load_4[29])), "=r"(*reinterpret_cast(&_tmem_load_4[30])), "=r"(*reinterpret_cast(&_tmem_load_4[31])) - : "r"(output_addr_1)); - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_4[32])), "=r"(*reinterpret_cast(&_tmem_load_4[33])), "=r"(*reinterpret_cast(&_tmem_load_4[34])), "=r"(*reinterpret_cast(&_tmem_load_4[35])), "=r"(*reinterpret_cast(&_tmem_load_4[36])), "=r"(*reinterpret_cast(&_tmem_load_4[37])), "=r"(*reinterpret_cast(&_tmem_load_4[38])), "=r"(*reinterpret_cast(&_tmem_load_4[39])), "=r"(*reinterpret_cast(&_tmem_load_4[40])), "=r"(*reinterpret_cast(&_tmem_load_4[41])), "=r"(*reinterpret_cast(&_tmem_load_4[42])), "=r"(*reinterpret_cast(&_tmem_load_4[43])), "=r"(*reinterpret_cast(&_tmem_load_4[44])), "=r"(*reinterpret_cast(&_tmem_load_4[45])), "=r"(*reinterpret_cast(&_tmem_load_4[46])), "=r"(*reinterpret_cast(&_tmem_load_4[47])), "=r"(*reinterpret_cast(&_tmem_load_4[48])), "=r"(*reinterpret_cast(&_tmem_load_4[49])), "=r"(*reinterpret_cast(&_tmem_load_4[50])), "=r"(*reinterpret_cast(&_tmem_load_4[51])), "=r"(*reinterpret_cast(&_tmem_load_4[52])), "=r"(*reinterpret_cast(&_tmem_load_4[53])), "=r"(*reinterpret_cast(&_tmem_load_4[54])), "=r"(*reinterpret_cast(&_tmem_load_4[55])), "=r"(*reinterpret_cast(&_tmem_load_4[56])), "=r"(*reinterpret_cast(&_tmem_load_4[57])), "=r"(*reinterpret_cast(&_tmem_load_4[58])), "=r"(*reinterpret_cast(&_tmem_load_4[59])), "=r"(*reinterpret_cast(&_tmem_load_4[60])), "=r"(*reinterpret_cast(&_tmem_load_4[61])), "=r"(*reinterpret_cast(&_tmem_load_4[62])), "=r"(*reinterpret_cast(&_tmem_load_4[63])) - : "r"(output_addr_1 + 32)); - int col_base = v_stage_1 * 128 + col_half_1 * 64; - if (my_row_1 < num_heads) { - #pragma unroll - for (int offset = 0; offset < 64; offset += 8) { - { - const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; - #if __CUDA_ARCH__ >= 1000 - #pragma unroll - for (int _ps = 0; _ps < 4; _ps++) - mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_4[offset])[_ps], _prescale2_1); - #else - #pragma unroll - for (int _ps = 0; _ps < 8; _ps++) - _tmem_load_4[offset + _ps] *= inv_sum * output_scale; - #endif - __nv_bfloat162 _pk[4]; - _pk[0] = __floats2bfloat162_rn(_tmem_load_4[offset + 0], _tmem_load_4[offset + 1]); - _pk[1] = __floats2bfloat162_rn(_tmem_load_4[offset + 2], _tmem_load_4[offset + 3]); - _pk[2] = __floats2bfloat162_rn(_tmem_load_4[offset + 4], _tmem_load_4[offset + 5]); - _pk[3] = __floats2bfloat162_rn(_tmem_load_4[offset + 6], _tmem_load_4[offset + 7]); - *reinterpret_cast(&((__nv_bfloat16*)(O + (output_base + col_base + offset)))[0]) = *reinterpret_cast(&_pk[0]); - } - } - } - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(o_empty_addr); - mbarrier_wait(work_full_addr + (work_stage_1) * 8, _phase_work_full_1); - uint32_t _clc_valid_3 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_3) - : "r"(work_response_addr + work_stage_1 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_3 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_3) - : "r"(work_response_addr + work_stage_1 * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(work_empty_addr + (work_stage_1) * 8); - work_stage_1 += 1; - if (work_stage_1 == 2) { work_stage_1 = 0; _phase_work_full_1 ^= 1; } - unsigned int valid_1 = _clc_valid_3; - query_idx_1 = _clc_ctaid_3; - if (valid_1 == 0) { - break; - } - } - asm volatile("barrier.sync 7, 128;" ::: "memory"); - if (warp == 4) { - int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); - asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); - } - } - } - // ---- Role: mma_warp ---- - if (warp == 8) { - { // mma_warp_main - unsigned int work_stage_2 = 0; - unsigned int query_idx_2 = blockIdx.x; - unsigned int global_tile_2 = 0; - unsigned int _phase_work_full_2 = 0; - #pragma unroll 1 - for (unsigned int work_iter_2 = 0; work_iter_2 < num_query_tokens; work_iter_2++) { - if (query_idx_2 >= (unsigned int)num_query_tokens) { - break; - } - int _max_2 = ((sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) : (0)); - int _min_1 = ((_max_2) < (sparse_topk) ? (_max_2) : (sparse_topk)); - int active_topk_2 = _min_1; - int _max_3 = (((active_topk_2 + 128 - 1) / 128) > (1) ? ((active_topk_2 + 128 - 1) / 128) : (1)); - int num_loop_steps_2 = _max_3; - unsigned int query_phase = work_iter_2 & 1; - mbarrier_wait(q_full_addr, query_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - if (num_loop_steps_2 > 0) { - unsigned int score_stage = global_tile_2 & 1; - int score_col_1 = ((score_stage != 0) ? 128 : 0); - #pragma unroll - for (int k_stage = 0; k_stage < 4; k_stage++) { - mbarrier_wait(kv_full_addr + (k_stage) * 8, 0); - asm volatile("tcgen05.fence::after_thread_sync;"); - int _mma_a_lo_0 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (k_stage) * 1024); - int _mma_b_lo_0 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (k_stage) * 2048); - { - uint64_t _mma_ss_a_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_0); - uint64_t _mma_ss_b_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_0); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, ((k_stage == 0) ? 0 : 1)); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 506U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 1018U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); - } - } - if (k_stage == 3) { - elect_commit(s_full_addr + (score_stage) * 8); - } - elect_commit(kv_empty_addr + (k_stage) * 8); - } - #pragma unroll 1 - for (int tile_2 = 0; tile_2 < num_loop_steps_2 - 1; tile_2++) { - unsigned int next_global_tile = global_tile_2 + 1; - unsigned int next_score_stage = next_global_tile & 1; - int next_score_col = ((next_score_stage != 0) ? 128 : 0); - #pragma unroll - for (int k_stage_1 = 0; k_stage_1 < 4; k_stage_1++) { - mbarrier_wait(kv_full_addr + (k_stage_1) * 8, 1); - asm volatile("tcgen05.fence::after_thread_sync;"); - int _mma_a_lo_1 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (k_stage_1) * 1024); - int _mma_b_lo_1 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (k_stage_1) * 2048); - { - uint64_t _mma_ss_a_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_1); - uint64_t _mma_ss_b_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_1); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (next_score_col)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, ((k_stage_1 == 0) ? 0 : 1)); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (next_score_col)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (next_score_col)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (next_score_col)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 506U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 1018U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (next_score_col)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (next_score_col)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (next_score_col)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (next_score_col)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); - } - } - if (k_stage_1 == 3) { - elect_commit(s_full_addr + (next_score_stage) * 8); - } - elect_commit(kv_empty_addr + (k_stage_1) * 8); - } - unsigned int sync_stage_1 = global_tile_2 & 1; - unsigned int sync_phase_1 = global_tile_2 >> 1 & 1; - int current_score_col = ((sync_stage_1 != 0) ? 128 : 0); - mbarrier_wait(p_full_addr + (sync_stage_1) * 8, sync_phase_1); - mbarrier_wait(o_empty_addr, 1 - (global_tile_2 & 1)); - asm volatile("tcgen05.fence::after_thread_sync;"); - #pragma unroll - for (int v_stage_2 = 0; v_stage_2 < 4; v_stage_2++) { - mbarrier_wait(kv_full_addr + (v_stage_2) * 8, 0); - asm volatile("tcgen05.fence::after_thread_sync;"); - int output_addr_2 = 256 + (v_stage_2 & 1) * 128 + ((v_stage_2 >> 1) * 16 << 16); - int _mma_b_lo_2 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (v_stage_2) * 2048); - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 dhi, blo, ta, id;\n\t" - ".reg .b64 db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "" - "mov.b32 dhi, 0x40004040;\n\t" - "mov.b32 id, 69272720;\n\t" - "mov.b32 ta, %2;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p0;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "}\n" - :: "r"((tmem_tmem + (output_addr_2))), "r"(_mma_b_lo_2), "r"(tmem_tmem + (current_score_col + 32 + ((v_stage_2 >> 1) * 16 << 16))), "r"(((tile_2 == 0) ? 0 : 1))); - elect_commit(kv_empty_addr + (v_stage_2) * 8); - } - elect_commit(o_full_addr); - global_tile_2 = next_global_tile; - } - elect_commit(q_empty_addr); - int last_tile = num_loop_steps_2 - 1; - unsigned int last_sync_stage = global_tile_2 & 1; - unsigned int last_sync_phase = global_tile_2 >> 1 & 1; - int last_score_col = ((last_sync_stage != 0) ? 128 : 0); - mbarrier_wait(p_full_addr + (last_sync_stage) * 8, last_sync_phase); - mbarrier_wait(o_empty_addr, 1 - (global_tile_2 & 1)); - asm volatile("tcgen05.fence::after_thread_sync;"); - #pragma unroll - for (int v_stage_3 = 0; v_stage_3 < 4; v_stage_3++) { - mbarrier_wait(kv_full_addr + (v_stage_3) * 8, 1); - asm volatile("tcgen05.fence::after_thread_sync;"); - int output_addr_3 = 256 + (v_stage_3 & 1) * 128 + ((v_stage_3 >> 1) * 16 << 16); - int _mma_b_lo_3 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (v_stage_3) * 2048); - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 dhi, blo, ta, id;\n\t" - ".reg .b64 db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "" - "mov.b32 dhi, 0x40004040;\n\t" - "mov.b32 id, 69272720;\n\t" - "mov.b32 ta, %2;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p0;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "}\n" - :: "r"((tmem_tmem + (output_addr_3))), "r"(_mma_b_lo_3), "r"(tmem_tmem + (last_score_col + 32 + ((v_stage_3 >> 1) * 16 << 16))), "r"(((last_tile == 0) ? 0 : 1))); - elect_commit(kv_empty_addr + (v_stage_3) * 8); - } - elect_commit(o_full_addr); - global_tile_2 = global_tile_2 + 1; - } else { - elect_commit(q_empty_addr); - } - mbarrier_wait(work_full_addr + (work_stage_2) * 8, _phase_work_full_2); - uint32_t _clc_valid_1 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_1) - : "r"(work_response_addr + work_stage_2 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_1 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_1) - : "r"(work_response_addr + work_stage_2 * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(work_empty_addr + (work_stage_2) * 8); - work_stage_2 += 1; - if (work_stage_2 == 2) { work_stage_2 = 0; _phase_work_full_2 ^= 1; } - unsigned int valid_2 = _clc_valid_1; - query_idx_2 = _clc_ctaid_1; - if (valid_2 == 0) { - break; - } - } - } - } - // ---- Role: page_warp ---- - if (warp == 9) { - { // page_warp_main - unsigned int page_stage = 0; - unsigned int work_stage_3 = 0; - unsigned int query_idx_3 = blockIdx.x; - unsigned int _phase_page_empty = 1; - unsigned int _phase_work_full_3 = 0; - #pragma unroll 1 - for (unsigned int work_iter_3 = 0; work_iter_3 < num_query_tokens; work_iter_3++) { - if (query_idx_3 >= (unsigned int)num_query_tokens) { - break; - } - int* swa_row = swa_indices + (query_idx_3 * (unsigned int)swa_index_stride); - int* compressed_row = compressed_indices + (query_idx_3 * (unsigned int)compressed_index_stride); - int _max_11 = ((sparse_topk_lens[query_idx_3] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_3] + sparse_topk_lens_offset) : (0)); - int _min_4 = ((_max_11) < (sparse_topk) ? (_max_11) : (sparse_topk)); - int active_topk_3 = _min_4; - int _max_12 = (((active_topk_3 + 128 - 1) / 128) > (1) ? ((active_topk_3 + 128 - 1) / 128) : (1)); - int num_loop_steps_3 = _max_12; - int num_page_pairs = (num_loop_steps_3 + 1) / 2; - int _max_13 = ((active_topk_3 - 1 & -4) > (0) ? (active_topk_3 - 1 & -4) : (0)); - int last_aligned = _max_13; - #pragma unroll 1 - for (int pair = 0; pair < num_page_pairs; pair++) { - int pair_base = pair * 256; - #pragma unroll - for (int duplicate = 0; duplicate < 2; duplicate++) { - mbarrier_wait(page_empty_addr + (page_stage) * 8, _phase_page_empty); - int lane_base = lane * 4; - int page0 = pair_base + lane_base; - int page1 = page0 + 128; - int safe0 = ((page0 <= last_aligned) ? page0 : last_aligned); - int safe1 = ((page1 <= last_aligned) ? page1 : last_aligned); - int* src0 = ((safe0 < 128) ? swa_row + safe0 : compressed_row + (safe0 - 128)); - int* src1 = ((safe1 < 128) ? swa_row + safe1 : compressed_row + (safe1 - 128)); - asm volatile("cp.async.ca.shared::cta.global [%0], [%1], 16;" - :: "r"(smem_page_offsets_addr + page_stage * 1024 + (unsigned int)(lane_base * 4)), "l"(src0)); - asm volatile("cp.async.ca.shared::cta.global [%0], [%1], 16;" - :: "r"(smem_page_offsets_addr + page_stage * 1024 + (unsigned int)((128 + lane_base) * 4)), "l"(src1)); - asm volatile( - "{\n\t" - "cp.async.mbarrier.arrive.noinc.shared::cta.b64 [%0];\n\t" - "}" - :: "r"(page_full_addr + (page_stage) * 8) : "memory"); - page_stage += 1; - if (page_stage == 6) { page_stage = 0; _phase_page_empty ^= 1; } - } - } - mbarrier_wait(work_full_addr + (work_stage_3) * 8, _phase_work_full_3); - uint32_t _clc_valid_4 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_4) - : "r"(work_response_addr + work_stage_3 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_4 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_4) - : "r"(work_response_addr + work_stage_3 * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(work_empty_addr + (work_stage_3) * 8); - work_stage_3 += 1; - if (work_stage_3 == 2) { work_stage_3 = 0; _phase_work_full_3 ^= 1; } - unsigned int valid_3 = _clc_valid_4; - query_idx_3 = _clc_ctaid_4; - if (valid_3 == 0) { - break; - } - } - } - } - // ---- Role: scheduler_warp ---- - if (warp == 10) { - { // scheduler_warp_main - unsigned int work_stage_4 = 0; - unsigned int throttle_stage = 0; - unsigned int query_idx_4 = blockIdx.x; - unsigned int _phase_throttle_full = 0; - unsigned int _phase_work_empty = 1; - unsigned int _phase_work_full_4 = 0; - #pragma unroll 1 - for (unsigned int work_iter_4 = 0; work_iter_4 < num_query_tokens; work_iter_4++) { - if (query_idx_4 >= (unsigned int)num_query_tokens) { - break; - } - mbarrier_wait(throttle_full_addr + (throttle_stage) * 8, _phase_throttle_full); - mbarrier_arrive(throttle_empty_addr + (throttle_stage) * 8); - throttle_stage += 1; - if (throttle_stage == 2) { throttle_stage = 0; _phase_throttle_full ^= 1; } - if (elect_sync()) { - mbarrier_wait(work_empty_addr + (work_stage_4) * 8, _phase_work_empty); - mbarrier_arrive_expect_tx(work_full_addr + (work_stage_4) * 8, 16); - asm volatile( - "fence.proxy.async.shared::cta;\n\t" - "clusterlaunchcontrol.try_cancel.async.shared::cta" - ".mbarrier::complete_tx::bytes.b128" - " [%0], [%1];" - :: "r"(work_response_addr + work_stage_4 * 16 + 0 * 16), "r"(work_full_addr + work_stage_4 * 8) - : "memory"); - } - mbarrier_wait(work_full_addr + (work_stage_4) * 8, _phase_work_full_4); - uint32_t _clc_valid_5 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_5) - : "r"(work_response_addr + work_stage_4 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_5 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_5) - : "r"(work_response_addr + work_stage_4 * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(work_empty_addr + (work_stage_4) * 8); - work_stage_4 += 1; - if (work_stage_4 == 2) { work_stage_4 = 0; _phase_work_empty ^= 1; _phase_work_full_4 ^= 1; } - unsigned int valid_4 = _clc_valid_5; - query_idx_4 = _clc_ctaid_5; - if (valid_4 == 0) { - break; - } - } - } - } - // ---- Role: padding_warp ---- - if (warp == 11) { - { // padding_warp_main - unsigned int work_stage_5 = 0; - unsigned int query_idx_5 = blockIdx.x; - unsigned int _phase_work_full_5 = 0; - #pragma unroll 1 - for (unsigned int work_iter_5 = 0; work_iter_5 < num_query_tokens; work_iter_5++) { - if (query_idx_5 >= (unsigned int)num_query_tokens) { - break; - } - mbarrier_wait(work_full_addr + (work_stage_5) * 8, _phase_work_full_5); - uint32_t _clc_valid_6 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_6) - : "r"(work_response_addr + work_stage_5 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_6 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_6) - : "r"(work_response_addr + work_stage_5 * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(work_empty_addr + (work_stage_5) * 8); - work_stage_5 += 1; - if (work_stage_5 == 2) { work_stage_5 = 0; _phase_work_full_5 ^= 1; } - unsigned int valid_5 = _clc_valid_6; - query_idx_5 = _clc_ctaid_6; - if (valid_5 == 0) { - break; - } - } - } - } - // ---- Role: load_warp ---- - if (warp >= 12 && warp <= 15) { - asm volatile("setmaxnreg.dec.sync.aligned.u32 72;"); - { // load_warp_main - asm volatile("griddepcontrol.wait;" ::: "memory"); - const int load_warp_rank = warp - 12; - unsigned int load_page_stage = 0; - unsigned int page_k_stage = 0; - unsigned int page_v_stage = 0; - int page_k_addr = 0; - int page_v_addr = 0; - unsigned int work_stage_6 = 0; - unsigned int throttle_stage_1 = 0; - unsigned int query_idx_6 = blockIdx.x; - unsigned int _phase_throttle_empty = 1; - unsigned int _phase_page_full = 0; - unsigned int _phase_work_full_6 = 0; - #pragma unroll 1 - for (unsigned int work_iter_6 = 0; work_iter_6 < num_query_tokens; work_iter_6++) { - if (query_idx_6 >= (unsigned int)num_query_tokens) { - break; - } - int _max_0 = ((sparse_topk_lens[query_idx_6] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_6] + sparse_topk_lens_offset) : (0)); - int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); - int active_topk_4 = _min_0; - int _max_1 = (((active_topk_4 + 128 - 1) / 128) > (1) ? ((active_topk_4 + 128 - 1) / 128) : (1)); - int num_loop_steps_4 = _max_1; - mbarrier_wait(throttle_empty_addr + (throttle_stage_1) * 8, _phase_throttle_empty); - mbarrier_arrive(throttle_full_addr + (throttle_stage_1) * 8); - throttle_stage_1 += 1; - if (throttle_stage_1 == 2) { throttle_stage_1 = 0; _phase_throttle_empty ^= 1; } - unsigned int query_phase_1 = work_iter_6 & 1; - mbarrier_wait(q_empty_addr, 1 - query_phase_1); - if (warp == 12) { - if (elect_sync()) { - mbarrier_arrive_expect_tx(q_full_addr, 65536); - #pragma unroll - for (int q_subload = 0; q_subload < 8; q_subload++) { - tma_4d_gmem2smem(smem_q_addr + (unsigned int)(q_subload * 8192), tmap_q, 0, 0, q_subload, query_idx_6, q_full_addr); - } - } - } - if (num_loop_steps_4 > 0) { - mbarrier_wait(page_full_addr + (load_page_stage) * 8, _phase_page_full); - page_k_stage = load_page_stage; - page_k_addr = smem_page_offsets_addr + page_k_stage * 1024; - load_page_stage += 1; - if (load_page_stage == 6) { load_page_stage = 0; _phase_page_full ^= 1; } - #pragma unroll - for (int kv_stage = 0; kv_stage < 4; kv_stage++) { - mbarrier_wait(kv_empty_addr + (kv_stage) * 8, 1); - if (warp == 12) { - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage) * 8, 32768); - } - } - #pragma unroll - for (int group = 0; group < 8; group++) { - int row_offset = load_warp_rank * 4 + group * 16; - int raw_rows[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows[0])), "=r"(*reinterpret_cast(&raw_rows[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows[(0) + 3])) - : "r"(page_k_addr + row_offset * 4)); - int row0 = raw_rows[0]; - int row1 = raw_rows[1]; - int row2 = raw_rows[2]; - int row3 = raw_rows[3]; - int dst = smem_kv_addr + (unsigned int)(kv_stage * 32768) + (unsigned int)(row_offset * 64 * 2); - if (elect_sync()) { - { - tma_gather4_gmem2smem(dst, tmap_swa_kv, kv_stage * 128, row0, row1, row2, row3, kv_full_addr + (kv_stage) * 8); - tma_gather4_gmem2smem(dst + 16384, tmap_swa_kv, kv_stage * 128 + 64, row0, row1, row2, row3, kv_full_addr + (kv_stage) * 8); - } - } - } - } - if (num_loop_steps_4 > 1) { - #pragma unroll - for (int kv_stage_1 = 0; kv_stage_1 < 4; kv_stage_1++) { - mbarrier_wait(kv_empty_addr + (kv_stage_1) * 8, 0); - if (warp == 12) { - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_1) * 8, 32768); - } - } - #pragma unroll - for (int group_1 = 0; group_1 < 8; group_1++) { - int row_offset_1 = load_warp_rank * 4 + group_1 * 16; - int raw_rows_1[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_1[0])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 3])) - : "r"(page_k_addr + 512 + row_offset_1 * 4)); - int row0_1 = raw_rows_1[0]; - int row1_1 = raw_rows_1[1]; - int row2_1 = raw_rows_1[2]; - int row3_1 = raw_rows_1[3]; - int dst_1 = smem_kv_addr + (unsigned int)(kv_stage_1 * 32768) + (unsigned int)(row_offset_1 * 64 * 2); - if (elect_sync()) { - { - tma_gather4_gmem2smem(dst_1, tmap_compressed_kv, kv_stage_1 * 128, row0_1, row1_1, row2_1, row3_1, kv_full_addr + (kv_stage_1) * 8); - tma_gather4_gmem2smem(dst_1 + 16384, tmap_compressed_kv, kv_stage_1 * 128 + 64, row0_1, row1_1, row2_1, row3_1, kv_full_addr + (kv_stage_1) * 8); - } - } - } - } - if (num_loop_steps_4 > 2) { - mbarrier_arrive(page_empty_addr + (page_k_stage) * 8); - } - mbarrier_wait(page_full_addr + (load_page_stage) * 8, _phase_page_full); - page_v_stage = load_page_stage; - page_v_addr = smem_page_offsets_addr + page_v_stage * 1024; - load_page_stage += 1; - if (load_page_stage == 6) { load_page_stage = 0; _phase_page_full ^= 1; } - #pragma unroll - for (int kv_stage_2 = 0; kv_stage_2 < 4; kv_stage_2++) { - mbarrier_wait(kv_empty_addr + (kv_stage_2) * 8, 1); - if (warp == 12) { - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_2) * 8, 32768); - } - } - #pragma unroll - for (int group_2 = 0; group_2 < 8; group_2++) { - int row_offset_2 = load_warp_rank * 4 + group_2 * 16; - int raw_rows_2[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_2[0])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 3])) - : "r"(page_v_addr + row_offset_2 * 4)); - int row0_2 = raw_rows_2[0]; - int row1_2 = raw_rows_2[1]; - int row2_2 = raw_rows_2[2]; - int row3_2 = raw_rows_2[3]; - int dst_2 = smem_v_addr + (unsigned int)(kv_stage_2 * 32768) + (unsigned int)(row_offset_2 * 64 * 2); - if (elect_sync()) { - { - tma_gather4_gmem2smem(dst_2, tmap_swa_kv, kv_stage_2 * 128, row0_2, row1_2, row2_2, row3_2, kv_full_addr + (kv_stage_2) * 8); - tma_gather4_gmem2smem(dst_2 + 16384, tmap_swa_kv, kv_stage_2 * 128 + 64, row0_2, row1_2, row2_2, row3_2, kv_full_addr + (kv_stage_2) * 8); - } - } - } - } - } - #pragma unroll 1 - for (int tile_3 = 1; tile_3 < num_loop_steps_4 - 1; tile_3++) { - int next_tile = tile_3 + 1; - if ((next_tile & 1) == 0) { - mbarrier_wait(page_full_addr + (load_page_stage) * 8, _phase_page_full); - page_k_stage = load_page_stage; - page_k_addr = smem_page_offsets_addr + page_k_stage * 1024; - load_page_stage += 1; - if (load_page_stage == 6) { load_page_stage = 0; _phase_page_full ^= 1; } - } - #pragma unroll - for (int kv_stage_3 = 0; kv_stage_3 < 4; kv_stage_3++) { - mbarrier_wait(kv_empty_addr + (kv_stage_3) * 8, 0); - if (warp == 12) { - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_3) * 8, 32768); - } - } - #pragma unroll - for (int group_3 = 0; group_3 < 8; group_3++) { - int row_offset_3 = load_warp_rank * 4 + group_3 * 16; - int raw_rows_3[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_3[0])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 3])) - : "r"(page_k_addr + (next_tile & 1) * 128 * 4 + row_offset_3 * 4)); - int row0_3 = raw_rows_3[0]; - int row1_3 = raw_rows_3[1]; - int row2_3 = raw_rows_3[2]; - int row3_3 = raw_rows_3[3]; - int dst_3 = smem_kv_addr + (unsigned int)(kv_stage_3 * 32768) + (unsigned int)(row_offset_3 * 64 * 2); - if (elect_sync()) { - { - tma_gather4_gmem2smem(dst_3, tmap_compressed_kv, kv_stage_3 * 128, row0_3, row1_3, row2_3, row3_3, kv_full_addr + (kv_stage_3) * 8); - tma_gather4_gmem2smem(dst_3 + 16384, tmap_compressed_kv, kv_stage_3 * 128 + 64, row0_3, row1_3, row2_3, row3_3, kv_full_addr + (kv_stage_3) * 8); - } - } - } - } - if ((next_tile & 1) != 0 && tile_3 < num_loop_steps_4 - 2) { - mbarrier_arrive(page_empty_addr + (page_k_stage) * 8); - } - if ((tile_3 & 1) == 0) { - mbarrier_wait(page_full_addr + (load_page_stage) * 8, _phase_page_full); - page_v_stage = load_page_stage; - page_v_addr = smem_page_offsets_addr + page_v_stage * 1024; - load_page_stage += 1; - if (load_page_stage == 6) { load_page_stage = 0; _phase_page_full ^= 1; } - } - #pragma unroll - for (int kv_stage_4 = 0; kv_stage_4 < 4; kv_stage_4++) { - mbarrier_wait(kv_empty_addr + (kv_stage_4) * 8, 1); - if (warp == 12) { - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_4) * 8, 32768); - } - } - #pragma unroll - for (int group_4 = 0; group_4 < 8; group_4++) { - int row_offset_4 = load_warp_rank * 4 + group_4 * 16; - int raw_rows_4[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_4[0])), "=r"(*reinterpret_cast(&raw_rows_4[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_4[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_4[(0) + 3])) - : "r"(page_v_addr + (tile_3 & 1) * 128 * 4 + row_offset_4 * 4)); - int row0_4 = raw_rows_4[0]; - int row1_4 = raw_rows_4[1]; - int row2_4 = raw_rows_4[2]; - int row3_4 = raw_rows_4[3]; - int dst_4 = smem_v_addr + (unsigned int)(kv_stage_4 * 32768) + (unsigned int)(row_offset_4 * 64 * 2); - if (elect_sync()) { - { - tma_gather4_gmem2smem(dst_4, tmap_compressed_kv, kv_stage_4 * 128, row0_4, row1_4, row2_4, row3_4, kv_full_addr + (kv_stage_4) * 8); - tma_gather4_gmem2smem(dst_4 + 16384, tmap_compressed_kv, kv_stage_4 * 128 + 64, row0_4, row1_4, row2_4, row3_4, kv_full_addr + (kv_stage_4) * 8); - } - } - } - } - if ((tile_3 & 1) != 0) { - mbarrier_arrive(page_empty_addr + (page_v_stage) * 8); - } - } - int last_tile_1 = num_loop_steps_4 - 1; - mbarrier_arrive(page_empty_addr + (page_k_stage) * 8); - if ((last_tile_1 & 1) == 0) { - mbarrier_wait(page_full_addr + (load_page_stage) * 8, _phase_page_full); - page_v_stage = load_page_stage; - page_v_addr = smem_page_offsets_addr + page_v_stage * 1024; - load_page_stage += 1; - if (load_page_stage == 6) { load_page_stage = 0; _phase_page_full ^= 1; } - } - if (num_loop_steps_4 == 1) { - #pragma unroll - for (int kv_stage_5 = 0; kv_stage_5 < 4; kv_stage_5++) { - mbarrier_wait(kv_empty_addr + (kv_stage_5) * 8, 0); - if (warp == 12) { - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_5) * 8, 32768); - } - } - #pragma unroll - for (int group_5 = 0; group_5 < 8; group_5++) { - int row_offset_5 = load_warp_rank * 4 + group_5 * 16; - int raw_rows_5[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_5[0])), "=r"(*reinterpret_cast(&raw_rows_5[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_5[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_5[(0) + 3])) - : "r"(page_v_addr + row_offset_5 * 4)); - int row0_5 = raw_rows_5[0]; - int row1_5 = raw_rows_5[1]; - int row2_5 = raw_rows_5[2]; - int row3_5 = raw_rows_5[3]; - int dst_5 = smem_v_addr + (unsigned int)(kv_stage_5 * 32768) + (unsigned int)(row_offset_5 * 64 * 2); - if (elect_sync()) { - { - tma_gather4_gmem2smem(dst_5, tmap_swa_kv, kv_stage_5 * 128, row0_5, row1_5, row2_5, row3_5, kv_full_addr + (kv_stage_5) * 8); - tma_gather4_gmem2smem(dst_5 + 16384, tmap_swa_kv, kv_stage_5 * 128 + 64, row0_5, row1_5, row2_5, row3_5, kv_full_addr + (kv_stage_5) * 8); - } - } - } - } - } else { - #pragma unroll - for (int kv_stage_6 = 0; kv_stage_6 < 4; kv_stage_6++) { - mbarrier_wait(kv_empty_addr + (kv_stage_6) * 8, 0); - if (warp == 12) { - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_6) * 8, 32768); - } - } - #pragma unroll - for (int group_6 = 0; group_6 < 8; group_6++) { - int row_offset_6 = load_warp_rank * 4 + group_6 * 16; - int raw_rows_6[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_6[0])), "=r"(*reinterpret_cast(&raw_rows_6[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_6[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_6[(0) + 3])) - : "r"(page_v_addr + (last_tile_1 & 1) * 128 * 4 + row_offset_6 * 4)); - int row0_6 = raw_rows_6[0]; - int row1_6 = raw_rows_6[1]; - int row2_6 = raw_rows_6[2]; - int row3_6 = raw_rows_6[3]; - int dst_6 = smem_v_addr + (unsigned int)(kv_stage_6 * 32768) + (unsigned int)(row_offset_6 * 64 * 2); - if (elect_sync()) { - { - tma_gather4_gmem2smem(dst_6, tmap_compressed_kv, kv_stage_6 * 128, row0_6, row1_6, row2_6, row3_6, kv_full_addr + (kv_stage_6) * 8); - tma_gather4_gmem2smem(dst_6 + 16384, tmap_compressed_kv, kv_stage_6 * 128 + 64, row0_6, row1_6, row2_6, row3_6, kv_full_addr + (kv_stage_6) * 8); - } - } - } - } - } - mbarrier_arrive(page_empty_addr + (page_v_stage) * 8); - } - mbarrier_wait(work_full_addr + (work_stage_6) * 8, _phase_work_full_6); - uint32_t _clc_valid_0 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_0) - : "r"(work_response_addr + work_stage_6 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_0 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_0) - : "r"(work_response_addr + work_stage_6 * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(work_empty_addr + (work_stage_6) * 8); - work_stage_6 += 1; - if (work_stage_6 == 2) { work_stage_6 = 0; _phase_work_full_6 ^= 1; } - unsigned int valid_6 = _clc_valid_0; - query_idx_6 = _clc_ctaid_0; - if (valid_6 == 0) { - break; - } - } - } - } - - // Cleanup -} - -} // extern "C" diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_4c41fb8347f137d4d46f_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_4c41fb8347f137d4d46f_binding.cu deleted file mode 100644 index 9ea7ff58e8d..00000000000 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_4c41fb8347f137d4d46f_binding.cu +++ /dev/null @@ -1,555 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -// Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_4c41fb8347f137d4d46f'. -#include -#include -#include -#include - -#include "tvm_ffi_utils.h" - -#include -#include -#include -#include -#include -#include -#include - -extern "C" __global__ void kernel_cake_dsv4_4c41fb8347f137d4d46f(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); - - -namespace cake_host_shim_22e55bd245aced16 { - -using tvm::ffi::Optional; -using tvm::ffi::TensorView; - -class ScopedCudaDevice { - public: - explicit ScopedCudaDevice(int device_id) { - cudaError_t error = cudaGetDevice(&previous_device_); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed before host-shim launch: cudaError=" - << static_cast(error); - if (previous_device_ != device_id) { - error = cudaSetDevice(device_id); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaSetDevice failed before host-shim launch for cuda:" - << device_id << ": cudaError=" << static_cast(error); - restore_ = true; - } - } - - ScopedCudaDevice(const ScopedCudaDevice&) = delete; - ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; - - ~ScopedCudaDevice() noexcept { - if (restore_) { - (void)cudaSetDevice(previous_device_); - } - } - - private: - int previous_device_ = -1; - bool restore_ = false; -}; - -inline int64_t CakeDeviceMultiprocessorCount(int device_id) { - constexpr int kMaxCachedCudaDevices = 64; - TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) - << "physical-SM-count cache does not cover cuda:" << device_id; - static std::atomic count_by_device[kMaxCachedCudaDevices]{}; - int cached = count_by_device[device_id].load(std::memory_order_acquire); - if (cached > 0) return cached; - - int count = 0; - cudaError_t error = cudaDeviceGetAttribute( - &count, cudaDevAttrMultiProcessorCount, device_id); - TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) - << "querying multiProcessorCount failed for cuda:" << device_id - << ": cudaError=" << static_cast(error) << ", count=" << count; - // Concurrent first launches may repeat the immutable device query, but all - // publication is atomic and every later hot-path lookup is one acquire load. - count_by_device[device_id].store(count, std::memory_order_release); - return count; -} - -inline void CheckCudaTensor(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) - << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; -} - -inline void CheckSameCudaDevice( - const TensorView& t, - const TensorView& reference, - const char* name, - const char* reference_name) { - TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) - << name << " must be on the same CUDA device as " << reference_name - << ": got cuda:" << t.device().device_id - << " versus cuda:" << reference.device().device_id; -} - -inline void CheckCurrentCudaDevice( - const TensorView& reference, - const char* reference_name) { - int current_device = -1; - cudaError_t error = cudaGetDevice(¤t_device); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed while validating " << reference_name - << ": cudaError=" << static_cast(error); - TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) - << "current CUDA device must match " << reference_name - << ": current=cuda:" << current_device - << ", tensor=cuda:" << reference.device().device_id; -} - -inline void CheckContiguous(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; -} - -inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { - DLDataType d = t.dtype(); - TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) - << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits - << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits - << ", lanes=" << (int)d.lanes << ")"; -} - -// A logical axis.outer(trailing) folds every source dim above the trailing -// dimensions. Shape products are independent of physical strides, so verify -// the leading dimensions form one dense row-major chain instead of inventing -// a "folded stride". The descriptor reads its exact adjacent physical step -// separately through stride[-(trailing + 1)]. -inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { - TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) - << name << " cannot fold leading dimensions above " << trailing - << " trailing dims from ndim=" << t.ndim(); - int outer_last = t.ndim() - trailing - 1; - if (outer_last <= 0) { - return; - } - int64_t step = t.stride(outer_last); - TVM_FFI_CHECK(step > 0, ValueError) - << name << " physical strides must be positive"; - int64_t expected = step; - for (int axis = outer_last - 1; axis >= 0; --axis) { - expected *= t.size(axis + 1); - if (t.size(axis) > 1) { - TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) - << name << " leading dims are not physically foldable above " << trailing - << " trailing dims: stride(" << axis << ")=" << t.stride(axis) - << ", expected " << expected; - } - } -} - -// 4D TMA descriptor for buffer 'tmap_q' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 3, ValueError) - << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - int64_t d2 = t.size(t.ndim() - 2); - int64_t d3 = t.size(t.ndim() - 3); - TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) - << "TMA source 'tmap_q' trailing dims must be positive"; - TVM_FFI_CHECK(d1 % 64 == 0, ValueError) - << "TMA source 'tmap_q' extent " << d1 - << " must divide exactly by " << 64; - uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(d3)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) - << "TMA box (64, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; - int64_t carrier_stride_0 = d1; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; - int64_t carrier_stride_1 = 64; - TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; - TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; - TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; - int64_t carrier_stride_2 = (d2 * d1); - TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; - TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; - TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; - uint64_t global_strides[3] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - (uint64_t)((carrier_stride_1 * 16) / 8), - (uint64_t)((carrier_stride_2 * 16) / 8), - }; - uint32_t box_dim[4] = {64u, 64u, 1u, 1u}; - uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_swa_k' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_swa_k(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_swa_k' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_swa_k' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_swa_k' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_swa_k"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_swa_k' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_swa_k' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_k'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_swa_k') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_compressed_k' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_compressed_k(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_compressed_k' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_compressed_k' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_compressed_k' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_compressed_k"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_compressed_k' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_compressed_k' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_k'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_k') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_swa_v' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_swa_v(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_swa_v' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_swa_v' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_swa_v' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_swa_v"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_swa_v' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_swa_v' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_v'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_swa_v') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_compressed_v' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_compressed_v(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_compressed_v' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_compressed_v' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_compressed_v' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_compressed_v"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_compressed_v' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_compressed_v' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_v'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_v') failed: CUresult=" << (int)r; - return tm; -} - -void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_compressed_k, TensorView arg_tmap_swa_v, TensorView arg_tmap_compressed_v, TensorView arg_O, TensorView arg_partial_lse, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_num_query_tokens, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, int64_t grid_x, int64_t grid_y, int64_t grid_z) { - DLDevice dev = arg_tmap_q.device(); - ScopedCudaDevice device_guard(dev.device_id); - CheckCudaTensor(arg_tmap_q, "tmap_q"); - CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); - CheckContiguous(arg_tmap_q, "tmap_q"); - CheckCudaTensor(arg_tmap_swa_k, "tmap_swa_k"); - CheckDtype(arg_tmap_swa_k, "tmap_swa_k", 4, 16, 1); - CheckCudaTensor(arg_tmap_compressed_k, "tmap_compressed_k"); - CheckDtype(arg_tmap_compressed_k, "tmap_compressed_k", 4, 16, 1); - CheckCudaTensor(arg_tmap_swa_v, "tmap_swa_v"); - CheckDtype(arg_tmap_swa_v, "tmap_swa_v", 4, 16, 1); - CheckCudaTensor(arg_tmap_compressed_v, "tmap_compressed_v"); - CheckDtype(arg_tmap_compressed_v, "tmap_compressed_v", 4, 16, 1); - CheckCudaTensor(arg_O, "O"); - CheckDtype(arg_O, "O", 4, 16, 1); - CheckContiguous(arg_O, "O"); - CheckCudaTensor(arg_partial_lse, "partial_lse"); - CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); - CheckContiguous(arg_partial_lse, "partial_lse"); - CheckCudaTensor(arg_swa_indices, "swa_indices"); - CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); - CheckContiguous(arg_swa_indices, "swa_indices"); - CheckCudaTensor(arg_compressed_indices, "compressed_indices"); - CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); - CheckContiguous(arg_compressed_indices, "compressed_indices"); - CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); - CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckCudaTensor(arg_sinks, "sinks"); - CheckDtype(arg_sinks, "sinks", 2, 32, 1); - CheckContiguous(arg_sinks, "sinks"); - CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); - CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); - CheckContiguous(arg_bmm1_scale, "bmm1_scale"); - CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); - CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); - CheckContiguous(arg_bmm2_scale, "bmm2_scale"); - TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) - << "scalar 'num_heads' value " << arg_num_heads - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_num_query_tokens >= -2147483648LL && arg_num_query_tokens <= 2147483647LL, ValueError) - << "scalar 'num_query_tokens' value " << arg_num_query_tokens - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) - << "scalar 'swa_index_stride' value " << arg_swa_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) - << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) - << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) - << "scalar 'sparse_topk' value " << arg_sparse_topk - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) - << "scalar 'has_sinks' value " << arg_has_sinks - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) - << "scalar 'total_work_items' value " << arg_total_work_items - << " is outside i32 range [-2147483648, 2147483647]"; - CheckSameCudaDevice(arg_tmap_swa_k, arg_tmap_q, "tmap_swa_k", "tmap_q"); - CheckSameCudaDevice(arg_tmap_compressed_k, arg_tmap_q, "tmap_compressed_k", "tmap_q"); - CheckSameCudaDevice(arg_tmap_swa_v, arg_tmap_q, "tmap_swa_v", "tmap_q"); - CheckSameCudaDevice(arg_tmap_compressed_v, arg_tmap_q, "tmap_compressed_v", "tmap_q"); - CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); - CheckSameCudaDevice(arg_partial_lse, arg_tmap_q, "partial_lse", "tmap_q"); - CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); - CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); - CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); - CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); - CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); - CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); - TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) - << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y - << ", " << grid_z << ")"; - TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) - << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z - << ") must be divisible by cluster dims (2, 1, 1)"; - - cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); - CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); - CUtensorMap p_tmap_swa_k = EncodeTma_tmap_swa_k(arg_tmap_swa_k); - CUtensorMap p_tmap_compressed_k = EncodeTma_tmap_compressed_k(arg_tmap_compressed_k); - CUtensorMap p_tmap_swa_v = EncodeTma_tmap_swa_v(arg_tmap_swa_v); - CUtensorMap p_tmap_compressed_v = EncodeTma_tmap_compressed_v(arg_tmap_compressed_v); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); - float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); - int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); - int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); - int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); - float* p_sinks = static_cast(arg_sinks.data_ptr()); - float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); - float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); - int32_t v_num_heads = (int32_t)arg_num_heads; - int32_t v_num_query_tokens = (int32_t)arg_num_query_tokens; - int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; - int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; - int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; - int32_t v_sparse_topk = (int32_t)arg_sparse_topk; - int32_t v_has_sinks = (int32_t)arg_has_sinks; - int32_t v_total_work_items = (int32_t)arg_total_work_items; - void* kargs[] = {&p_tmap_q, &p_tmap_swa_k, &p_tmap_compressed_k, &p_tmap_swa_v, &p_tmap_compressed_v, &p_O, &p_partial_lse, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_num_query_tokens, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_has_sinks, &v_total_work_items}; - - static std::mutex smem_mu; - static auto* smem_status_by_device = - new std::unordered_map(); - cudaError_t smem_status = cudaSuccess; - { - std::lock_guard lock(smem_mu); - auto it = smem_status_by_device->find(dev.device_id); - if (it == smem_status_by_device->end()) { - int current_device = -1; - smem_status = cudaGetDevice(¤t_device); - if (smem_status == cudaSuccess && current_device != dev.device_id) { - smem_status = cudaErrorInvalidDevice; - } - if (smem_status == cudaSuccess) { - smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_4c41fb8347f137d4d46f), - cudaFuncAttributeMaxDynamicSharedMemorySize, - 220032); - } - smem_status_by_device->emplace(dev.device_id, smem_status); - } else { - smem_status = it->second; - } - } - TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_4c41fb8347f137d4d46f failed: " - << cudaGetErrorString(smem_status); - dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); - dim3 block(512u, 1u, 1u); - - // Extended direct-source launch for a linked CUDA kernel symbol. - cudaLaunchAttribute attrs[2]{}; - int n = 0; - attrs[n].id = cudaLaunchAttributeClusterDimension; - attrs[n].val.clusterDim.x = 2u; - attrs[n].val.clusterDim.y = 1u; - attrs[n].val.clusterDim.z = 1u; - ++n; - attrs[n].id = cudaLaunchAttributeClusterSchedulingPolicyPreference; - attrs[n].val.clusterSchedulingPolicyPreference = cudaClusterSchedulingPolicySpread; - ++n; - cudaLaunchConfig_t config{}; - config.gridDim = grid; - config.blockDim = block; - config.dynamicSmemBytes = 220032u; - config.stream = stream; - config.attrs = attrs; - config.numAttrs = n; - cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_4c41fb8347f137d4d46f), kargs); - TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_4c41fb8347f137d4d46f failed: " - << cudaGetErrorString(launch_status); -} - -} // namespace cake_host_shim_22e55bd245aced16 - -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_22e55bd245aced16::Run); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_4c41fb8347f137d4d46f_kernel.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_4c41fb8347f137d4d46f_kernel.cu deleted file mode 100644 index d55adfe2ced..00000000000 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_4c41fb8347f137d4d46f_kernel.cu +++ /dev/null @@ -1,2352 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -typedef signed char int8_t; -typedef unsigned char uint8_t; -typedef unsigned short uint16_t; -typedef unsigned int uint32_t; -#if defined(__CUDACC_RTC__) -typedef unsigned long long uint64_t; -#else -typedef unsigned long uint64_t; -#endif -static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); -typedef signed int int32_t; -typedef short int int16_t; -struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; -struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; -static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); -static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); - -#if defined(__CUDACC_RTC__) -typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; -#else -#include -#endif - -static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); -static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); -#include -#include - -__device__ __forceinline__ int make_warp_uniform(int x) { - int result; - asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" - : "=r"(result) : "r"(x)); - return result; -} - -#define CAKE_INF CUDART_INF_F -#define TMEM_NCOLS 512 -#define TMEM_TMEM_SCRATCH_OFFSET 0 -#define NUM_K_PIPE_STAGES 3 -#define NUM_V_PIPE_STAGES 2 -#define SMEM_SMEM_Q_OFF 1024 -#define SMEM_SMEM_Q_STAGE_BYTES 8192 -#define SMEM_SMEM_Q_STRIDE 8192 -#define SMEM_SMEM_K_OFF 66560 -#define SMEM_SMEM_K_STAGE_BYTES 16384 -#define SMEM_SMEM_K_STRIDE 16384 -#define SMEM_SMEM_V_LO_OFF 115712 -#define SMEM_SMEM_V_LO_STAGE_BYTES 16384 -#define SMEM_SMEM_V_LO_STRIDE 32768 -#define SMEM_SMEM_V_HI_OFF 132096 -#define SMEM_SMEM_V_HI_STAGE_BYTES 16384 -#define SMEM_SMEM_V_HI_STRIDE 32768 -#define SMEM_SMEM_P_OFF 181248 -#define SMEM_SMEM_P_STAGE_BYTES 8192 -#define SMEM_SMEM_P_STRIDE 8192 -#define SMEM_SMEM_INDICES_OFF 214016 -#define SMEM_SMEM_INDICES_STAGE_BYTES 2048 -#define SMEM_SMEM_INDICES_STRIDE 2048 -#define SMEM_SMEM_SOFTMAX_EXCHANGE_OFF 218624 -#define SMEM_SMEM_SOFTMAX_EXCHANGE_STAGE_BYTES 512 -#define SMEM_SMEM_SOFTMAX_EXCHANGE_STRIDE 512 -#define SMEM_SMEM_EPILOGUE_EXCHANGE_OFF 219136 -#define SMEM_SMEM_EPILOGUE_EXCHANGE_STAGE_BYTES 512 -#define SMEM_SMEM_EPILOGUE_EXCHANGE_STRIDE 512 -#define SMEM_SMEM_INDEX_VALID_OFF 219648 -#define SMEM_SMEM_INDEX_VALID_STAGE_BYTES 288 -#define SMEM_SMEM_INDEX_VALID_STRIDE 288 -#define SMEM_TOTAL 220032 -#define THREADS 512 - -#include - -__device__ __forceinline__ uint32_t elect_sync() { - uint32_t pred = 0; - asm volatile( - "{\n\t" - ".reg .pred %%px;\n\t" - "elect.sync _|%%px, %1;\n\t" - "@%%px mov.s32 %0, 1;\n\t" - "}\n" - : "+r"(pred) - : "r"(0xFFFFFFFF)); - return pred; -} - - -__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { - asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" - :: "r"(mbar_addr), "r"(count) : "memory"); -} - -__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { - asm volatile("mbarrier.init.b64 [%0], %1;" - :: "l"(mbar_addr), "r"(count) : "memory"); -} - - -__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - - -// CTA-local pipelines have short, resident producer/consumer edges. Omitting -// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the -// explicit loop still makes this helper blocking until acquire succeeds. -__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE;\n\t" - "bra.uni LAB_WAIT;\n\t" - "DONE:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Source-faithful relaxed CTA wait used only by a typed protocol that does -// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. -__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, 10000000;\n\t" - "@P1 bra.uni DONE_RELAXED;\n\t" - "bra.uni LAB_WAIT_RELAXED;\n\t" - "DONE_RELAXED:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Exact source ports may request the PTX suspendTimeHint operand explicitly. -// The hint is expressed in nanoseconds and is kept separate from the canonical -// no-hint CTA helper so unrelated schedules retain their existing retry path. -__device__ __forceinline__ void mbarrier_wait_suspend( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_SUSPEND:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_SUSPEND;\n\t" - "bra.uni LAB_WAIT_SUSPEND;\n\t" - "DONE_SUSPEND:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE_CLUSTER;\n\t" - "bra.uni LAB_WAIT_CLUSTER;\n\t" - "DONE_CLUSTER:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - ".reg .u32 WAIT_ADDR;\n\t" - "mov.u32 WAIT_ADDR, %0;\n\t" - "LAB_WAIT_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [WAIT_ADDR], %1, %2;\n\t" - "@P1 bra.uni DONE_HINT;\n\t" - "bra.uni LAB_WAIT_HINT;\n\t" - "DONE_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -// Exact unqualified CTA wait used by source schedules whose PTX intentionally -// omits the acquire qualifier while retaining a typed suspendTimeHint operand. -__device__ __forceinline__ void mbarrier_wait_relaxed_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED_HINT:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra DONE_RELAXED_HINT;\n\t" - "bra LAB_WAIT_RELAXED_HINT;\n\t" - "DONE_RELAXED_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); -} - -__device__ __forceinline__ void mbarrier_wait_cluster_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" - "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" - "DONE_CLUSTER_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_suspend( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait_cluster(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); - } -} - - -__device__ __forceinline__ void tcgen05_mma_f16_cta2( - int taddr, uint64_t a_desc, uint64_t b_desc, - uint32_t i_desc, int enable_input_d) { - asm volatile( - "{\n\t" - ".reg .pred p;\n\t" - ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - "setp.ne.b32 p, %4, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" - "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "tcgen05.mma.cta_group::2.kind::f16 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" - "}\n" - :: "r"(taddr), "l"(a_desc), "l"(b_desc), - "r"(i_desc), "r"(enable_input_d) - : "memory"); -} - - -__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { - return (x & 0x3FFFFULL) >> 4ULL; -} - - -union MmaSmemDesc { - uint64_t u64; - uint32_t u32[2]; -}; - -__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { - MmaSmemDesc tmp; - tmp.u64 = smem_desc; - tmp.u32[0] += offset; - smem_desc = tmp.u64; -} - - -__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { - asm volatile( - "{\n\t" - ".reg .pred leader;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" - ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" - "}\n" - :: "r"(mbar_addr), "h"(cta_mask) : "memory"); -} - - -__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { - asm volatile( - "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" - :: "r"(mbar_addr), "r"(bytes) : "memory"); -} - - -__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7," - " %8, %9, %10, %11, %12, %13, %14, %15," - " %16, %17, %18, %19, %20, %21, %22, %23," - " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), - "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), - "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), - "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), - "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), - "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), - "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x16.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7," - " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), - "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), - "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ void tmem_ld_x4(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x4.b32" - " {%0, %1, %2, %3}, [%4];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { - asm volatile( - "tcgen05.st.sync.aligned.32x32b.x16.b32" - " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," - " %9, %10, %11, %12, %13, %14, %15, %16};" - :: "r"(tmem_addr), - "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), - "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), - "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), - "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); -} - - -__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { - asm volatile( - "tcgen05.st.sync.aligned.32x32b.x4.b32" - " [%0], {%1, %2, %3, %4};" - :: "r"(tmem_addr), - "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); -} - - -__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { - uint32_t addr; - asm("{\n\t" - ".reg .u64 u64addr;\n\t" - "cvta.to.shared.u64 u64addr, %1;\n\t" - "cvt.u32.u64 %0, u64addr;\n\t" - "}\n" : "=r"(addr) : "l"(ptr)); - return addr; -} - - -__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { - uint32_t remote; - asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" - : "=r"(remote) : "r"(local_addr), "r"(rank)); - return remote; -} - - -__device__ __forceinline__ float approx_exp2(float x) { - float y; - asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float approx_rcp(float x) { - float y; - asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float max_noftz(float a, float b) { - float c; - asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); - return c; -} - - -__device__ __forceinline__ float warp_reduce_max(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); - return val; -} - - -__device__ __forceinline__ float warp_reduce_sum(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val += __shfl_xor_sync(0xFFFFFFFF, val, offset); - return val; -} - - -__device__ __forceinline__ float row_max_reduce(float2 acc) { - return max_noftz(acc.x, acc.y); -} - - -__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { - #pragma unroll - for (int j = 0; j < 16; j++) { - if (j % 2 == 0) - acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); - else - acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); - } -} - - -__device__ __forceinline__ float2 ex2_emulation_f32x2_value(float2 value) { - const float c0 = 1.0f, c1 = 0.695146143436431884765625f; - const float c2 = 0.227564394474029541015625f, c3 = 0.077119089663028717041015625f; - const float magic = 12582912.0f; - float x0 = max_noftz(value.x, -127.0f), x1 = max_noftz(value.y, -127.0f); - float2 xc2 = make_float2(x0, x1), magic2 = make_float2(magic, magic); - float2 xr2; - asm("add.rm.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xr2) - : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&magic2)); - float2 c3_2 = make_float2(c3, c3), c2_2 = make_float2(c2, c2); - float2 c1_2 = make_float2(c1, c1), c0_2 = make_float2(c0, c0); - float2 xrb2, xfrac2; - asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xrb2) - : "l"(*(unsigned long long*)&xr2), "l"(*(unsigned long long*)&magic2)); - asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xfrac2) - : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&xrb2)); - float2 poly2; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&c3_2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c2_2)); - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c1_2)); - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c0_2)); - int x0r_i, x1r_i, p0_i, p1_i; - asm("mov.b64 {%0, %1}, %2;" : "=r"(x0r_i), "=r"(x1r_i) : "l"(*(unsigned long long*)&xr2)); - asm("mov.b64 {%0, %1}, %2;" : "=r"(p0_i), "=r"(p1_i) : "l"(*(unsigned long long*)&poly2)); - float r0, r1; - asm("mov.b32 %0, %1;" : "=f"(r0) : "r"((x0r_i << 23) + p0_i)); - asm("mov.b32 %0, %1;" : "=f"(r1) : "r"((x1r_i << 23) + p1_i)); - return make_float2(r0, r1); -} - -__device__ __forceinline__ void ex2_emulation_f32x2(float* x0_ptr, float* x1_ptr) { - float2 result = ex2_emulation_f32x2_value(make_float2(*x0_ptr, *x1_ptr)); - *x0_ptr = result.x; *x1_ptr = result.y; -} - -__device__ __forceinline__ void softmax_frag_exp2_cast( - float* sv, uint32_t* pv, int use_emu) -{ - #pragma unroll - for (int j = 0; j < 16; j++) { - if (use_emu && j >= 12) - ex2_emulation_f32x2(&sv[j*2], &sv[j*2+1]); - else { - sv[j*2] = approx_exp2(sv[j*2]); - sv[j*2+1] = approx_exp2(sv[j*2+1]); - } - } - #pragma unroll - for (int j = 0; j < 16; j++) { - __nv_bfloat162 bf = __float22bfloat162_rn({sv[j*2], sv[j*2+1]}); - pv[j] = reinterpret_cast(bf); - } -} - - - -__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { - const float2* sv2 = reinterpret_cast(sv); - #pragma unroll - for (int j = 0; j < 16; j++) { - asm("add.f32x2 %0, %1, %2;" - : "+l"(reinterpret_cast(*acc)) - : "l"(reinterpret_cast(*acc)), - "l"(reinterpret_cast(sv2[j]))); - } -} - - -__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { - asm("mul.rn.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { - asm("mul.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { - asm("add.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { - asm("add.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { - asm("sub.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { - asm("sub.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("add.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { - float2 r; - asm("sub.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("sub.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ void fma_scale_x32( - float* sv, const float2* scale2, const float2* neg_max2) -{ - float2* sv_2 = reinterpret_cast(sv); - #pragma unroll - for (int j = 0; j < 16; j++) - fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); -} - -__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("mul.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - - -// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) - -__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("add.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("add.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("add.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("add.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("add.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("add.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("add.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - - -__device__ __forceinline__ void fence_async_shared() { - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); -} - - -__device__ __forceinline__ uint64_t make_smem_desc(int addr) { - const int SBO = 1024; - return desc_encode(addr) - | (desc_encode(SBO) << 32ULL) - | (1ULL << 46ULL) - | (2ULL << 61ULL); -} - - -__device__ __forceinline__ void tma_2d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3}], [%4];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_4d_gmem2smem_cta2( - int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.4d.shared::cluster.global" - ".mbarrier::complete_tx::bytes.cta_group::2" - " [%0], [%1, {%2, %3, %4, %5}], [%6];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr) { - // Canonical .shared::cta form for non-multicast gather4, matching - // trtllm-gen / cuda_ptx and the PTX ISA qualifier order - // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, - // .shared::cluster is reserved for the multicast variant (ctaMask). - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem_cta2( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr) { - // Canonical .shared::cta form; see tma_gather4_gmem2smem above. - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" - ".mbarrier::complete_tx::bytes.cta_group::2" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem_mc( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr, unsigned short cta_mask) { - // Multicast variant: the PTX grammar ties the .shared::cluster - // destination to .multicast::cluster + ctaMask (cf. cuda_ptx / - // SM100_TMA_LOAD_MULTICAST_2D_GATHER4 in CUTLASS). - asm volatile( - "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" - ".mbarrier::complete_tx::bytes.multicast::cluster" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr), "h"(cta_mask) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem_mc_cta2( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr, unsigned short cta_mask) { - // Multicast + cta_group::2 variant; see tma_gather4_gmem2smem_mc. - asm volatile( - "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" - ".mbarrier::complete_tx::bytes.multicast::cluster.cta_group::2" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr), "h"(cta_mask) : "memory"); -} - - -__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { - asm volatile( - "{\n\t" - ".reg .b16 lo, hi;\n\t" - "mov.b32 {lo, hi}, %1;\n\t" - "tcgen05.commit.cta_group::2.mbarrier::arrive::one" - ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" - "}\n" - :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); -} - -extern "C" { - -__global__ __launch_bounds__(512, 1) __cluster_dims__(2,1,1) void -kernel_cake_dsv4_4c41fb8347f137d4d46f(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items) -{ - const int tid = threadIdx.x; - const int warp = make_warp_uniform(tid / 32); - const int lane = tid % 32; - - extern __shared__ __align__(1024) char smem_raw[]; - int smem; - smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); - - const int mbar_base = smem; - #define q_full_addr (mbar_base + 0) - #define q_empty_addr (mbar_base + 8) - #define k_full_addr (mbar_base + 16) - #define k_empty_addr (mbar_base + 40) - #define v_full_addr (mbar_base + 64) - #define v_empty_addr (mbar_base + 80) - #define s_full_addr (mbar_base + 96) - #define s_empty_addr (mbar_base + 112) - #define p_full_addr (mbar_base + 128) - #define p_empty_addr (mbar_base + 144) - #define o_empty_addr (mbar_base + 160) - #define stats_addr (mbar_base + 168) - #define stats_empty_addr (mbar_base + 184) - #define o_full_addr (mbar_base + 200) - #define o_first_slice_seeded_addr (mbar_base + 208) - #define tmem_dealloc_peer_addr (mbar_base + 216) - #define index_valid_full_addr (mbar_base + 224) - #define index_valid_empty_addr (mbar_base + 240) - - const int bid = blockIdx.x; - const int num_bids = gridDim.x; - const unsigned int clusters_x = gridDim.x / 2; - const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; - const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; - - int cta_rank; - asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); - - // Kernel setup ops - __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); - const int smem_q_addr = smem + 1024; - __nv_bfloat16* smem_k = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); - const int smem_k_addr = smem + 66560; - __nv_bfloat16* smem_v_lo = reinterpret_cast<__nv_bfloat16*>(smem_raw + 115712); - const int smem_v_lo_addr = smem + 115712; - __nv_bfloat16* smem_v_hi = reinterpret_cast<__nv_bfloat16*>(smem_raw + 132096); - const int smem_v_hi_addr = smem + 132096; - __nv_bfloat16* smem_p = reinterpret_cast<__nv_bfloat16*>(smem_raw + 181248); - const int smem_p_addr = smem + 181248; - int* smem_indices = reinterpret_cast(smem_raw + 214016); - const int smem_indices_addr = smem + 214016; - float* smem_softmax_exchange = reinterpret_cast(smem_raw + 218624); - const int smem_softmax_exchange_addr = smem + 218624; - float* smem_epilogue_exchange = reinterpret_cast(smem_raw + 219136); - const int smem_epilogue_exchange_addr = smem + 219136; - unsigned int* smem_index_valid = reinterpret_cast(smem_raw + 219648); - const int smem_index_valid_addr = smem + 219648; - asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); - asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); - - // Mbarrier init (18 pipeline groups, 0 ordered-sequence groups, 32 barriers) - // Mbarriers at smem_raw[0..256) - - if (warp == 0) { - uint32_t leader = elect_sync(); - if (leader) { - // q_full: 1 barriers, init_count=1 - mbarrier_init(smem + 0, 1); - // q_empty: 1 barriers, init_count=1 - mbarrier_init(smem + 8, 1); - // --- pipeline 'k_pipe' --- - // k_full: 3 barriers, init_count=1 - mbarrier_init(smem + 16, 1); - mbarrier_init(smem + 24, 1); - mbarrier_init(smem + 32, 1); - // k_empty: 3 barriers, init_count=1 - mbarrier_init(smem + 40, 1); - mbarrier_init(smem + 48, 1); - mbarrier_init(smem + 56, 1); - // --- pipeline 'v_pipe' --- - // v_full: 2 barriers, init_count=1 - mbarrier_init(smem + 64, 1); - mbarrier_init(smem + 72, 1); - // v_empty: 2 barriers, init_count=1 - mbarrier_init(smem + 80, 1); - mbarrier_init(smem + 88, 1); - // s_full: 2 barriers, init_count=1 - mbarrier_init(smem + 96, 1); - mbarrier_init(smem + 104, 1); - // s_empty: 2 barriers, init_count=256 - mbarrier_init(smem + 112, 256); - mbarrier_init(smem + 120, 256); - // p_full: 2 barriers, init_count=256 - mbarrier_init(smem + 128, 256); - mbarrier_init(smem + 136, 256); - // p_empty: 2 barriers, init_count=1 - mbarrier_init(smem + 144, 1); - mbarrier_init(smem + 152, 1); - // o_empty: 1 barriers, init_count=256 - mbarrier_init(smem + 160, 256); - // stats: 2 barriers, init_count=128 - mbarrier_init(smem + 168, 128); - mbarrier_init(smem + 176, 128); - // stats_empty: 2 barriers, init_count=128 - mbarrier_init(smem + 184, 128); - mbarrier_init(smem + 192, 128); - // o_full: 1 barriers, init_count=1 - mbarrier_init(smem + 200, 1); - // o_first_slice_seeded: 1 barriers, init_count=256 - mbarrier_init(smem + 208, 256); - // tmem_dealloc_peer: 1 barriers, init_count=32 - mbarrier_init(smem + 216, 32); - // index_valid_full: 2 barriers, init_count=32 - mbarrier_init(smem + 224, 32); - mbarrier_init(smem + 232, 32); - // index_valid_empty: 2 barriers, init_count=128 - mbarrier_init(smem + 240, 128); - mbarrier_init(smem + 248, 128); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - - asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); - asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); - - // TMEM alloc (512 columns, 512 used) - volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 256); - if (warp == 0) { - int _tmem_hold = smem + 256; - asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); - __syncwarp(); - asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); - } - - __syncthreads(); - asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); - asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); - - const int taddr = tmem_addr_storage[0]; - - // Kernel post-init ops - const int tmem_tmem_scratch = taddr; - - // ---- Ordered hardware-WG register redistribution ---- - // Dec phase frees registers before any WG attempts inc. - if (warp >= 8 && warp <= 15) { - asm volatile("setmaxnreg.dec.sync.aligned.u32 56;"); - } - - // ---- Role: softmax_wg ---- - if (warp <= 3) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); - { // softmax_wg_main - const int softmax_dummy = 0; - float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; - const int head_row_base = warp % 2 * 32; - const int score_row_base = warp % 4 * 32; - const int n_half = warp % 4 / 2; - const int my_row = head_row_base + lane; - const int exchange_idx = n_half * 64 + my_row; - int tile_cursor = 0; - int valid_iter = 0; - float seed_zero[4]; - #pragma unroll - for (int seed_i = 0; seed_i < 4; seed_i++) { - seed_zero[seed_i] = 0.0f; - } - #pragma unroll - for (int seed_stage = 0; seed_stage < 2; seed_stage++) { - #pragma unroll - for (int seed_col = 128 + seed_stage * 128; seed_col < 128 + (seed_stage + 1) * 128; seed_col += 4) { - int seed_addr = taddr + (unsigned int)seed_col + (unsigned int)(score_row_base << 16); - tmem_st_x4_f32(seed_addr, seed_zero); - } - } - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(o_first_slice_seeded_addr); - int peer_rank = cta_rank ^ 1; - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(o_first_slice_seeded_addr), "r"(peer_rank) : "memory"); - #pragma unroll 1 - for (unsigned int work_idx = cluster_id; work_idx < total_work_items; work_idx += num_clusters) { - int split_idx = work_idx % (unsigned int)(((1) ? 4 : 1)); - int query_idx = work_idx / (unsigned int)(((1) ? 4 : 1)); - int _max_0 = ((sparse_topk_lens[query_idx] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx] + sparse_topk_lens_offset) : (0)); - int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); - int active_topk = _min_0; - int num_kv_tiles = ((sparse_topk + 128 - 1) / 128 + ((1) ? 4 : 1) - 1) / ((1) ? 4 : 1); - int first_kv_tile = split_idx * num_kv_tiles; - int head_idx = cta_rank * 64 + my_row; - float row_max_val = -CAKE_INF; - float row_sum_val = 0.0f; - if (has_sinks != 0 && head_idx < num_heads && split_idx == 0) { - row_max_val = sinks[head_idx] * 1.4426950408889634f / softmax_scale_log2; - if (n_half == 0) { - row_sum_val = 1.0f; - } - } - int valid_slot = valid_iter & 1; - int valid_full_phase = valid_iter >> 1 & 1; - mbarrier_wait(index_valid_full_addr + (valid_slot) * 8, valid_full_phase); - int valid_word_base = valid_slot * 36 + n_half * 2; - #pragma unroll 1 - for (int tile = 0; tile < num_kv_tiles; tile++) { - int pipeline_tile = tile_cursor + tile; - int score_phase = pipeline_tile & 1; - int s_full_phase = pipeline_tile >> 1 & 1; - mbarrier_wait(s_full_addr + (score_phase) * 8, s_full_phase); - int score_offset = ((score_phase != 0) ? 64 : 0); - int score_base = taddr + (unsigned int)score_offset + (unsigned int)(score_row_base << 16); - float _tmem_load_0[64]; - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]) - : "r"(score_base)); - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]) - : "r"(score_base + 32)); - int valid_cols = active_topk - (first_kv_tile + tile) * 128 - n_half * 64; - if (valid_cols < 0) { - valid_cols = 0; - } - if (valid_cols > 64) { - valid_cols = 64; - } - unsigned int keep_lo = smem_index_valid[valid_word_base + tile * 4]; - unsigned int keep_hi = smem_index_valid[valid_word_base + tile * 4 + 1]; - uint32_t _thresh_mask_0; - { - int _lim_0 = valid_cols; - if (_lim_0 <= 0) { _thresh_mask_0 = 0u; } - else if (_lim_0 >= 32) { _thresh_mask_0 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_thresh_mask_0) : "r"(_lim_0)); - } - } - uint32_t _thresh_mask_1; - { - int _lim_1 = valid_cols - 32; - if (_lim_1 <= 0) { _thresh_mask_1 = 0u; } - else if (_lim_1 >= 32) { _thresh_mask_1 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_thresh_mask_1) : "r"(_lim_1)); - } - } - if (!(_thresh_mask_0 & keep_lo & (1u << 0))) _tmem_load_0[0] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 1))) _tmem_load_0[1] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 2))) _tmem_load_0[2] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 3))) _tmem_load_0[3] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 4))) _tmem_load_0[4] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 5))) _tmem_load_0[5] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 6))) _tmem_load_0[6] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 7))) _tmem_load_0[7] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 8))) _tmem_load_0[8] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 9))) _tmem_load_0[9] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 10))) _tmem_load_0[10] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 11))) _tmem_load_0[11] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 12))) _tmem_load_0[12] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 13))) _tmem_load_0[13] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 14))) _tmem_load_0[14] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 15))) _tmem_load_0[15] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 16))) _tmem_load_0[16] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 17))) _tmem_load_0[17] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 18))) _tmem_load_0[18] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 19))) _tmem_load_0[19] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 20))) _tmem_load_0[20] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 21))) _tmem_load_0[21] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 22))) _tmem_load_0[22] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 23))) _tmem_load_0[23] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 24))) _tmem_load_0[24] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 25))) _tmem_load_0[25] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 26))) _tmem_load_0[26] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 27))) _tmem_load_0[27] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 28))) _tmem_load_0[28] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 29))) _tmem_load_0[29] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 30))) _tmem_load_0[30] = -CAKE_INF; - if (!(_thresh_mask_0 & keep_lo & (1u << 31))) _tmem_load_0[31] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 0))) _tmem_load_0[32] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 1))) _tmem_load_0[33] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 2))) _tmem_load_0[34] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 3))) _tmem_load_0[35] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 4))) _tmem_load_0[36] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 5))) _tmem_load_0[37] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 6))) _tmem_load_0[38] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 7))) _tmem_load_0[39] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 8))) _tmem_load_0[40] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 9))) _tmem_load_0[41] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 10))) _tmem_load_0[42] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 11))) _tmem_load_0[43] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 12))) _tmem_load_0[44] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 13))) _tmem_load_0[45] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 14))) _tmem_load_0[46] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 15))) _tmem_load_0[47] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 16))) _tmem_load_0[48] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 17))) _tmem_load_0[49] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 18))) _tmem_load_0[50] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 19))) _tmem_load_0[51] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 20))) _tmem_load_0[52] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 21))) _tmem_load_0[53] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 22))) _tmem_load_0[54] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 23))) _tmem_load_0[55] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 24))) _tmem_load_0[56] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 25))) _tmem_load_0[57] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 26))) _tmem_load_0[58] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 27))) _tmem_load_0[59] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 28))) _tmem_load_0[60] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 29))) _tmem_load_0[61] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 30))) _tmem_load_0[62] = -CAKE_INF; - if (!(_thresh_mask_1 & keep_hi & (1u << 31))) _tmem_load_0[63] = -CAKE_INF; - float2 _reg_reduce_max2_2 = {-CAKE_INF, -CAKE_INF}; - row_max_x32_accum(&_tmem_load_0[0], _reg_reduce_max2_2); - row_max_x32_accum(&_tmem_load_0[32], _reg_reduce_max2_2); - float _tmem_load_0_max = row_max_reduce(_reg_reduce_max2_2); - float _max_1 = max_noftz(_tmem_load_0_max, row_max_val); - float new_max = _max_1; - smem_softmax_exchange[exchange_idx] = new_max; - asm volatile("barrier.sync 3, 128;" ::: "memory"); - float _max_2 = max_noftz(new_max, smem_softmax_exchange[exchange_idx ^ 64]); - new_max = _max_2; - asm volatile("barrier.sync 3, 128;" ::: "memory"); - float no_correction = (((new_max - row_max_val) * softmax_scale_log2 <= 0.0f) ? 1.0f : 0.0f); - float delta = softmax_scale_log2 * (row_max_val - new_max); - float _exp2_0 = approx_exp2(delta); - float exp_delta = _exp2_0; - float acc_scale = ((row_max_val > -CAKE_INF) ? exp_delta : 1.0f); - row_max_val = new_max; - float safe_max = ((new_max == -CAKE_INF) ? 0.0f : new_max); - float max_scaled = safe_max * softmax_scale_log2; - const float2 _fma_b2_3 = {softmax_scale_log2, softmax_scale_log2}; - const float2 _fma_c2_4 = {-max_scaled, -max_scaled}; - #pragma unroll - for (int _lf = 0; _lf < 32; _lf++) - fma_f32x2_inplace(&reinterpret_cast(_tmem_load_0)[_lf], _fma_b2_3, _fma_c2_4); - #pragma unroll - for (int _le = 0; _le < 64; _le++) { - _tmem_load_0[_le] = approx_exp2(_tmem_load_0[_le]); - } - float2 _reg_reduce_sum2_5 = make_float2(0.0f, 0.0f); - softmax_block_sum(&_tmem_load_0[0], &_reg_reduce_sum2_5); - softmax_block_sum(&_tmem_load_0[32], &_reg_reduce_sum2_5); - float _tmem_load_0_sum = _reg_reduce_sum2_5.x + _reg_reduce_sum2_5.y; - row_sum_val = row_sum_val * acc_scale + _tmem_load_0_sum; - int stats_empty_phase = pipeline_tile >> 1 & 1 ^ 1; - mbarrier_wait(stats_empty_addr + (score_phase) * 8, stats_empty_phase); - float meta[4]; - meta[0] = row_sum_val; - meta[1] = row_max_val; - meta[2] = acc_scale; - meta[3] = no_correction; - int meta_addr = taddr + 384 + (unsigned int)(score_phase * 8) + (unsigned int)(n_half * 4) + (unsigned int)(head_row_base << 16); - tmem_st_x4_f32(meta_addr, meta); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - mbarrier_arrive(stats_addr + (score_phase) * 8); - int p_stage = pipeline_tile % 2; - int p_empty_phase = pipeline_tile / 2 & 1 ^ 1; - mbarrier_wait(p_empty_addr + (p_stage) * 8, p_empty_phase); - uint32_t _tmem_load_0_bf16[32]; - #pragma unroll - for (int _lp = 0; _lp < 32; _lp++) { - __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_0[_lp*2 + 0], _tmem_load_0[_lp*2+1 + 0])); - _tmem_load_0_bf16[_lp] = *(uint32_t*)&_bf2; - } - int smem_p_stage = p_stage * 2 + n_half; - int p_base_smem = smem_p_addr + (unsigned int)(smem_p_stage * 8192); - #pragma unroll - for (int vec = 0; vec < 8; vec++) { - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_base_smem + (my_row * 128 + vec * 16 ^ (my_row * 128 + vec * 16 >> 7 & 7) << 4))), "r"(_tmem_load_0_bf16[vec * 4]), "r"(_tmem_load_0_bf16[vec * 4 + 1]), "r"(_tmem_load_0_bf16[vec * 4 + 2]), "r"(_tmem_load_0_bf16[vec * 4 + 3]) : "memory"); - } - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(p_full_addr + (unsigned int)(p_stage * 8)), "r"(0) : "memory"); - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(s_empty_addr + (unsigned int)(score_phase * 8)), "r"(0) : "memory"); - } - mbarrier_arrive(index_valid_empty_addr + (valid_slot) * 8); - valid_iter = valid_iter + 1; - tile_cursor = tile_cursor + num_kv_tiles; - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - } - } - // ---- Role: correction_wg ---- - if (warp >= 4 && warp <= 7) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); - { // correction_wg_main - const int correction_dummy = 0; - float softmax_scale_log2_1 = bmm1_scale[0] * 1.4426950408889634f; - float output_scale = bmm2_scale[0]; - const int head_row_base_1 = warp % 2 * 32; - const int output_row_base = warp % 4 * 32; - const int n_half_1 = warp % 4 / 2; - const int my_row_1 = head_row_base_1 + lane; - const int corr_row = head_row_base_1 << 16; - int tile_cursor_1 = 0; - unsigned int _phase_o_full_0 = 0; - #pragma unroll 1 - for (unsigned int work_idx_1 = cluster_id; work_idx_1 < total_work_items; work_idx_1 += num_clusters) { - int split_idx_1 = work_idx_1 % (unsigned int)(((1) ? 4 : 1)); - int query_idx_1 = work_idx_1 / (unsigned int)(((1) ? 4 : 1)); - int num_kv_tiles_1 = ((sparse_topk + 128 - 1) / 128 + ((1) ? 4 : 1) - 1) / ((1) ? 4 : 1); - float final_sum = 0.0f; - float final_max = -CAKE_INF; - #pragma unroll 1 - for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { - int pipeline_tile_1 = tile_cursor_1 + tile_1; - int score_phase_1 = pipeline_tile_1 & 1; - int stats_phase = pipeline_tile_1 >> 1 & 1; - mbarrier_wait(stats_addr + (score_phase_1) * 8, stats_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - int meta_addr_1 = taddr + 384 + (unsigned int)(score_phase_1 * 8) + (unsigned int)(n_half_1 * 4) + (unsigned int)corr_row; - float _tmem_load_1[4]; - tmem_ld_x4(&_tmem_load_1[0], meta_addr_1); - final_sum = _tmem_load_1[0]; - final_max = _tmem_load_1[1]; - float acc_scale_1 = _tmem_load_1[2]; - float no_correction_1 = _tmem_load_1[3]; - if (tile_1 > 0) { - mbarrier_wait(o_full_addr, _phase_o_full_0); - _phase_o_full_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - int _vote_0 = __all_sync(0xFFFFFFFF, no_correction_1 == 1.0f); - int skip_correction = _vote_0; - if (skip_correction == 0) { - #pragma unroll - for (int local_slice = 0; local_slice < 2; local_slice++) { - #pragma unroll - for (int acc_stage = 0; acc_stage < 2; acc_stage++) { - int o_base = taddr + 128 + (unsigned int)(acc_stage * 128) + (unsigned int)(local_slice * 64) + (unsigned int)(output_row_base << 16); - #pragma unroll - for (int c = 0; c < 64; c += 16) { - float _tmem_load_2[16]; - tmem_ld_x16(&_tmem_load_2[0], o_base + c); - const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; - #pragma unroll - for (int _ls = 0; _ls < 8; _ls++) - mul_f32x2_inplace(&reinterpret_cast(_tmem_load_2)[_ls], _scale2_0); - tmem_st_x16_f32(o_base + c, _tmem_load_2); - } - } - } - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - } - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(o_empty_addr), "r"(0) : "memory"); - } - mbarrier_arrive(stats_empty_addr + (score_phase_1) * 8); - } - mbarrier_wait(o_full_addr, _phase_o_full_0); - _phase_o_full_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - const int exchange_idx_1 = n_half_1 * 64 + my_row_1; - smem_epilogue_exchange[exchange_idx_1] = final_sum; - asm volatile("barrier.sync 4, 128;" ::: "memory"); - final_sum = final_sum + smem_epilogue_exchange[exchange_idx_1 ^ 64]; - asm volatile("barrier.sync 4, 128;" ::: "memory"); - float _rcp_0 = approx_rcp(final_sum); - float inv_sum = ((final_sum > 0.0f) ? _rcp_0 : 0.0f); - int head_idx_1 = cta_rank * 64 + my_row_1; - int output_offset = ((1) ? ((query_idx_1 * num_heads + head_idx_1) * ((1) ? 4 : 1) + split_idx_1) * 512 : (query_idx_1 * num_heads + head_idx_1) * 512); - if (n_half_1 == 0 && head_idx_1 < num_heads) { - int stat_offset = (query_idx_1 * num_heads + head_idx_1) * ((1) ? 4 : 1) + split_idx_1; - float _log2_0; - asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(final_sum)); - partial_lse[stat_offset] = ((final_sum > 0.0f) ? final_max * softmax_scale_log2_1 + _log2_0 : -CAKE_INF); - } - #pragma unroll - for (int acc_stage_1 = 0; acc_stage_1 < 2; acc_stage_1++) { - #pragma unroll - for (int local_slice_1 = 0; local_slice_1 < 2; local_slice_1++) { - int logical_slice = acc_stage_1 * 2 * 2 + n_half_1 * 2 + local_slice_1; - int o_base_1 = taddr + 128 + (unsigned int)(acc_stage_1 * 128) + (unsigned int)(local_slice_1 * 64) + (unsigned int)(output_row_base << 16); - #pragma unroll - for (int c_1 = 0; c_1 < 64; c_1 += 32) { - float _tmem_load_3[32]; - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]) - : "r"(o_base_1 + c_1)); - int gmem_base = output_offset + logical_slice * 64 + c_1; - #pragma unroll - for (int j = 0; j < 32; j += 16) { - { - const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; - #if __CUDA_ARCH__ >= 1000 - #pragma unroll - for (int _ps = 0; _ps < 8; _ps++) - mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_3[j])[_ps], _prescale2_1); - #else - #pragma unroll - for (int _ps = 0; _ps < 16; _ps++) - _tmem_load_3[j + _ps] *= inv_sum * output_scale; - #endif - { - __nv_bfloat162 _pk0 = __floats2bfloat162_rn(_tmem_load_3[j + 0], _tmem_load_3[j + 1]); - unsigned _pk_u0 = *reinterpret_cast(&_pk0); - __nv_bfloat162 _pk1 = __floats2bfloat162_rn(_tmem_load_3[j + 2], _tmem_load_3[j + 3]); - unsigned _pk_u1 = *reinterpret_cast(&_pk1); - __nv_bfloat162 _pk2 = __floats2bfloat162_rn(_tmem_load_3[j + 4], _tmem_load_3[j + 5]); - unsigned _pk_u2 = *reinterpret_cast(&_pk2); - __nv_bfloat162 _pk3 = __floats2bfloat162_rn(_tmem_load_3[j + 6], _tmem_load_3[j + 7]); - unsigned _pk_u3 = *reinterpret_cast(&_pk3); - __nv_bfloat162 _pk4 = __floats2bfloat162_rn(_tmem_load_3[j + 8], _tmem_load_3[j + 9]); - unsigned _pk_u4 = *reinterpret_cast(&_pk4); - __nv_bfloat162 _pk5 = __floats2bfloat162_rn(_tmem_load_3[j + 10], _tmem_load_3[j + 11]); - unsigned _pk_u5 = *reinterpret_cast(&_pk5); - __nv_bfloat162 _pk6 = __floats2bfloat162_rn(_tmem_load_3[j + 12], _tmem_load_3[j + 13]); - unsigned _pk_u6 = *reinterpret_cast(&_pk6); - __nv_bfloat162 _pk7 = __floats2bfloat162_rn(_tmem_load_3[j + 14], _tmem_load_3[j + 15]); - unsigned _pk_u7 = *reinterpret_cast(&_pk7); - asm volatile( - "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" - :: "l"((void*)(&((__nv_bfloat16*)(O + (gmem_base + j)))[0])), "r"(_pk_u0), "r"(_pk_u1), "r"(_pk_u2), "r"(_pk_u3), "r"(_pk_u4), "r"(_pk_u5), "r"(_pk_u6), "r"(_pk_u7) : "memory"); - } - } - } - } - } - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(o_empty_addr), "r"(0) : "memory"); - tile_cursor_1 = tile_cursor_1 + num_kv_tiles_1; - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - } - } - // ---- Role: mma_warp ---- - if (warp == 8) { - { // mma_warp_main - const int mma_dummy = 0; - unsigned int k_stage = 0; - unsigned int v_stage = 0; - int tile_cursor_2 = 0; - unsigned int _phase_o_first_slice_seeded_0 = 0; - mbarrier_wait(o_first_slice_seeded_addr, _phase_o_first_slice_seeded_0); - _phase_o_first_slice_seeded_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - unsigned int _phase_q_full_0 = 0; - unsigned int _phase_k_full = 0; - unsigned int _phase_o_empty_0 = 1; - unsigned int _phase_v_full = 0; - #pragma unroll 1 - for (unsigned int work_idx_2 = cluster_id; work_idx_2 < total_work_items; work_idx_2 += num_clusters) { - int num_kv_tiles_2 = ((sparse_topk + 128 - 1) / 128 + ((1) ? 4 : 1) - 1) / ((1) ? 4 : 1); - if (cta_rank == 0) { - mbarrier_wait(q_full_addr, _phase_q_full_0); - _phase_q_full_0 ^= 1; - int first_pv = 1; - #pragma unroll 1 - for (int tile_2 = 0; tile_2 < num_kv_tiles_2; tile_2++) { - int pipeline_tile_2 = tile_cursor_2 + tile_2; - int score_phase_2 = pipeline_tile_2 & 1; - int s_empty_phase = pipeline_tile_2 >> 1 & 1 ^ 1; - mbarrier_wait(s_empty_addr + (score_phase_2) * 8, s_empty_phase); - int score_col = ((score_phase_2 != 0) ? 64 : 0); - #pragma unroll - for (int qk_stage = 0; qk_stage < 4; qk_stage++) { - mbarrier_wait(k_full_addr + (k_stage) * 8, _phase_k_full); - int _mma_a_lo_0 = (((smem_q_addr) >> 4) & 0x3FFF) + (qk_stage * 2) * 512; - int _mma_b_lo_0 = (((smem_k_addr) >> 4) & 0x3FFF) + (k_stage) * 1024; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 136316048;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (score_col))), "r"(((qk_stage == 0) ? 0 : 1))); - int _mma_a_lo_1 = (((smem_q_addr) >> 4) & 0x3FFF) + (qk_stage * 2 + 1) * 512; - int _mma_b_lo_1 = (((smem_k_addr + 8192) >> 4) & 0x3FFF) + (k_stage) * 1024; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 136316048;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); - elect_commit_cg2_multicast(k_empty_addr + (k_stage) * 8, (uint16_t)(3)); - k_stage += 1; - if (k_stage == 3) { k_stage = 0; _phase_k_full ^= 1; } - } - elect_commit_cg2_multicast(s_full_addr + (score_phase_2) * 8, (uint16_t)(3)); - if (tile_2 > 0) { - int prev_pipeline_tile = pipeline_tile_2 - 1; - int p_stage_1 = prev_pipeline_tile % 2; - int p_full_phase = prev_pipeline_tile / 2 & 1; - mbarrier_wait(o_empty_addr, _phase_o_empty_0); - _phase_o_empty_0 ^= 1; - mbarrier_wait(p_full_addr + (p_stage_1) * 8, p_full_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - #pragma unroll - for (int ps = 0; ps < 2; ps++) { - int smem_p_stage_1 = p_stage_1 * 2 + ps; - mbarrier_wait(v_full_addr + (v_stage) * 8, _phase_v_full); - #pragma unroll - for (int acc_stage_2 = 0; acc_stage_2 < 2; acc_stage_2++) { - int output_col = 128 + acc_stage_2 * 128; - if (acc_stage_2 == 0) { - int _mma_a_lo_2 = (((smem_p_addr) >> 4) & 0x3FFF) + (smem_p_stage_1) * 512; - int _mma_b_lo_2 = ((((smem_v_lo_addr) >> 4) & 0x3FFF) | 0x2000000) + (v_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 138478736;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (output_col))), "r"(((((ps == 0) ? first_pv : 0)) ? 0 : 1))); - } else { - int _mma_a_lo_3 = (((smem_p_addr) >> 4) & 0x3FFF) + (smem_p_stage_1) * 512; - int _mma_b_lo_3 = ((((smem_v_hi_addr) >> 4) & 0x3FFF) | 0x2000000) + (v_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 138478736;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"((tmem_tmem_scratch + (output_col))), "r"(((((ps == 0) ? first_pv : 0)) ? 0 : 1))); - } - } - elect_commit_cg2_multicast(v_empty_addr + (v_stage) * 8, (uint16_t)(3)); - v_stage += 1; - if (v_stage == 2) { v_stage = 0; _phase_v_full ^= 1; } - } - first_pv = 0; - elect_commit_cg2_multicast(p_empty_addr + (p_stage_1) * 8, (uint16_t)(3)); - elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); - } - } - int last_pipeline_tile = tile_cursor_2 + num_kv_tiles_2 - 1; - int last_p_stage = last_pipeline_tile % 2; - int last_p_phase = last_pipeline_tile / 2 & 1; - mbarrier_wait(o_empty_addr, _phase_o_empty_0); - _phase_o_empty_0 ^= 1; - mbarrier_wait(p_full_addr + (last_p_stage) * 8, last_p_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - #pragma unroll - for (int ps_1 = 0; ps_1 < 2; ps_1++) { - int smem_p_stage_2 = last_p_stage * 2 + ps_1; - mbarrier_wait(v_full_addr + (v_stage) * 8, _phase_v_full); - #pragma unroll - for (int acc_stage_3 = 0; acc_stage_3 < 2; acc_stage_3++) { - int output_col_1 = 128 + acc_stage_3 * 128; - if (acc_stage_3 == 0) { - int _mma_a_lo_4 = (((smem_p_addr) >> 4) & 0x3FFF) + (smem_p_stage_2) * 512; - int _mma_b_lo_4 = ((((smem_v_lo_addr) >> 4) & 0x3FFF) | 0x2000000) + (v_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 138478736;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"((tmem_tmem_scratch + (output_col_1))), "r"(((((ps_1 == 0) ? first_pv : 0)) ? 0 : 1))); - } else { - int _mma_a_lo_5 = (((smem_p_addr) >> 4) & 0x3FFF) + (smem_p_stage_2) * 512; - int _mma_b_lo_5 = ((((smem_v_hi_addr) >> 4) & 0x3FFF) | 0x2000000) + (v_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 138478736;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_5), "r"(_mma_b_lo_5), "r"((tmem_tmem_scratch + (output_col_1))), "r"(((((ps_1 == 0) ? first_pv : 0)) ? 0 : 1))); - } - } - elect_commit_cg2_multicast(v_empty_addr + (v_stage) * 8, (uint16_t)(3)); - v_stage += 1; - if (v_stage == 2) { v_stage = 0; _phase_v_full ^= 1; } - } - elect_commit_cg2_multicast(p_empty_addr + (last_p_stage) * 8, (uint16_t)(3)); - elect_commit_cg2_multicast(o_full_addr, (uint16_t)(3)); - elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); - } - tile_cursor_2 = tile_cursor_2 + num_kv_tiles_2; - } - if (cta_rank == 0) { - #pragma unroll - for (int tail_offset = 0; tail_offset < 2; tail_offset++) { - int tail_tile = tile_cursor_2 + tail_offset; - int tail_stage = tail_tile & 1; - int tail_phase = tail_tile >> 1 & 1 ^ 1; - mbarrier_wait(s_empty_addr + (tail_stage) * 8, tail_phase); - } - mbarrier_wait(o_empty_addr, _phase_o_empty_0); - _phase_o_empty_0 ^= 1; - } - asm volatile("tcgen05.fence::before_thread_sync;"); - int peer_rank_1 = cta_rank ^ 1; - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(tmem_dealloc_peer_addr), "r"(peer_rank_1) : "memory"); - mbarrier_wait(tmem_dealloc_peer_addr, 0); - int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); - asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); - } - } - // ---- Role: load_wg ---- - if (warp >= 9 && warp <= 12) { - { // load_wg_main - const int load_dummy = 0; - const int load_warp_rank = warp - 9; - unsigned int k_stage_1 = 0; - unsigned int v_stage_1 = 0; - unsigned int _phase_q_empty_0 = 1; - unsigned int _phase_k_empty = 1; - unsigned int _phase_v_empty = 1; - #pragma unroll 1 - for (unsigned int work_idx_3 = cluster_id; work_idx_3 < total_work_items; work_idx_3 += num_clusters) { - int split_idx_2 = work_idx_3 % (unsigned int)(((1) ? 4 : 1)); - int query_idx_2 = work_idx_3 / (unsigned int)(((1) ? 4 : 1)); - int num_kv_tiles_3 = ((sparse_topk + 128 - 1) / 128 + ((1) ? 4 : 1) - 1) / ((1) ? 4 : 1); - int first_kv_tile_1 = split_idx_2 * num_kv_tiles_3; - mbarrier_wait(q_empty_addr, _phase_q_empty_0); - _phase_q_empty_0 ^= 1; - if (load_warp_rank == 0) { - if (cta_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(131072)) : "memory"); - } - } - if (elect_sync()) { - #pragma unroll - for (int q_stage = 0; q_stage < 8; q_stage++) { - tma_4d_gmem2smem_cta2(smem_q_addr + (unsigned int)(q_stage * 8192), (&tmap_q), 0, cta_rank * 64, q_stage, query_idx_2, ((q_full_addr) & 0xFEFFFFFF)); - } - } - } - asm volatile("barrier.sync 2, 160;" ::: "memory"); - #pragma unroll - for (int qk_stage_1 = 0; qk_stage_1 < 4; qk_stage_1++) { - mbarrier_wait(k_empty_addr + (k_stage_1) * 8, _phase_k_empty); - if (load_warp_rank == 0) { - if (cta_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); - } - } - } - int group = load_warp_rank * 4 + lane; - if (lane < 4) { - int index_offset = cta_rank * 64 + group * 4; - int raw_rows[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows[0])), "=r"(*reinterpret_cast(&raw_rows[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows[(0) + 3])) - : "r"(smem_indices_addr + (unsigned int)(index_offset * 4))); - if (first_kv_tile_1 == 0) { - tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group * 512), (&tmap_swa_k), qk_stage_1 * 128, ((raw_rows[0] >= 0) ? raw_rows[0] : ((1) ? group * 4 : 0)), ((raw_rows[1] >= 0) ? raw_rows[1] : ((1) ? group * 4 + 1 : 0)), ((raw_rows[2] >= 0) ? raw_rows[2] : ((1) ? group * 4 + 2 : 0)), ((raw_rows[3] >= 0) ? raw_rows[3] : ((1) ? group * 4 + 3 : 0)), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group * 512) + 8192, (&tmap_swa_k), qk_stage_1 * 128 + 64, ((raw_rows[0] >= 0) ? raw_rows[0] : ((1) ? group * 4 : 0)), ((raw_rows[1] >= 0) ? raw_rows[1] : ((1) ? group * 4 + 1 : 0)), ((raw_rows[2] >= 0) ? raw_rows[2] : ((1) ? group * 4 + 2 : 0)), ((raw_rows[3] >= 0) ? raw_rows[3] : ((1) ? group * 4 + 3 : 0)), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - } else { - tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group * 512), (&tmap_compressed_k), qk_stage_1 * 128, ((raw_rows[0] >= 0) ? raw_rows[0] : ((1) ? group * 4 : 0)), ((raw_rows[1] >= 0) ? raw_rows[1] : ((1) ? group * 4 + 1 : 0)), ((raw_rows[2] >= 0) ? raw_rows[2] : ((1) ? group * 4 + 2 : 0)), ((raw_rows[3] >= 0) ? raw_rows[3] : ((1) ? group * 4 + 3 : 0)), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group * 512) + 8192, (&tmap_compressed_k), qk_stage_1 * 128 + 64, ((raw_rows[0] >= 0) ? raw_rows[0] : ((1) ? group * 4 : 0)), ((raw_rows[1] >= 0) ? raw_rows[1] : ((1) ? group * 4 + 1 : 0)), ((raw_rows[2] >= 0) ? raw_rows[2] : ((1) ? group * 4 + 2 : 0)), ((raw_rows[3] >= 0) ? raw_rows[3] : ((1) ? group * 4 + 3 : 0)), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - } - } - k_stage_1 += 1; - if (k_stage_1 == 3) { k_stage_1 = 0; _phase_k_empty ^= 1; } - } - #pragma unroll 1 - for (int tile_3 = 1; tile_3 < num_kv_tiles_3; tile_3++) { - #pragma unroll - for (int qk_stage_2 = 0; qk_stage_2 < 4; qk_stage_2++) { - mbarrier_wait(k_empty_addr + (k_stage_1) * 8, _phase_k_empty); - if (load_warp_rank == 0) { - if (cta_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); - } - } - } - int group_1 = load_warp_rank * 4 + lane; - if (lane < 4) { - int index_offset_1 = tile_3 * 128 + cta_rank * 64 + group_1 * 4; - int raw_rows_1[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_1[0])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 3])) - : "r"(smem_indices_addr + (unsigned int)(index_offset_1 * 4))); - tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group_1 * 512), (&tmap_compressed_k), qk_stage_2 * 128, ((raw_rows_1[0] >= 0) ? raw_rows_1[0] : ((1) ? group_1 * 4 : 0)), ((raw_rows_1[1] >= 0) ? raw_rows_1[1] : ((1) ? group_1 * 4 + 1 : 0)), ((raw_rows_1[2] >= 0) ? raw_rows_1[2] : ((1) ? group_1 * 4 + 2 : 0)), ((raw_rows_1[3] >= 0) ? raw_rows_1[3] : ((1) ? group_1 * 4 + 3 : 0)), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - tma_gather4_gmem2smem_mc_cta2(smem_k_addr + k_stage_1 * 16384 + (unsigned int)(group_1 * 512) + 8192, (&tmap_compressed_k), qk_stage_2 * 128 + 64, ((raw_rows_1[0] >= 0) ? raw_rows_1[0] : ((1) ? group_1 * 4 : 0)), ((raw_rows_1[1] >= 0) ? raw_rows_1[1] : ((1) ? group_1 * 4 + 1 : 0)), ((raw_rows_1[2] >= 0) ? raw_rows_1[2] : ((1) ? group_1 * 4 + 2 : 0)), ((raw_rows_1[3] >= 0) ? raw_rows_1[3] : ((1) ? group_1 * 4 + 3 : 0)), ((k_full_addr + (k_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - } - k_stage_1 += 1; - if (k_stage_1 == 3) { k_stage_1 = 0; _phase_k_empty ^= 1; } - } - int prev_tile = tile_3 - 1; - #pragma unroll - for (int ps_2 = 0; ps_2 < 2; ps_2++) { - mbarrier_wait(v_empty_addr + (v_stage_1) * 8, _phase_v_empty); - if (load_warp_rank == 0) { - if (cta_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(65536)) : "memory"); - } - } - } - int group_2 = load_warp_rank * 4 + lane; - if (lane < 4) { - int index_offset_2 = prev_tile * 128 + ps_2 * 64 + group_2 * 4; - int raw_rows_2[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_2[0])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 3])) - : "r"(smem_indices_addr + (unsigned int)(index_offset_2 * 4))); - #pragma unroll - for (int acc_stage_4 = 0; acc_stage_4 < 2; acc_stage_4++) { - int dst_v = ((acc_stage_4 == 0) ? smem_v_lo_addr + v_stage_1 * 32768 : smem_v_hi_addr + v_stage_1 * 32768); - if (first_kv_tile_1 + prev_tile == 0) { - tma_gather4_gmem2smem_mc_cta2(dst_v + group_2 * 512, (&tmap_swa_v), acc_stage_4 * 128 * 2 + cta_rank * 128, ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : ((1) ? group_2 * 4 : 0)), ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : ((1) ? group_2 * 4 + 1 : 0)), ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : ((1) ? group_2 * 4 + 2 : 0)), ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : ((1) ? group_2 * 4 + 3 : 0)), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - tma_gather4_gmem2smem_mc_cta2(dst_v + group_2 * 512 + 8192, (&tmap_swa_v), acc_stage_4 * 128 * 2 + cta_rank * 128 + 64, ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : ((1) ? group_2 * 4 : 0)), ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : ((1) ? group_2 * 4 + 1 : 0)), ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : ((1) ? group_2 * 4 + 2 : 0)), ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : ((1) ? group_2 * 4 + 3 : 0)), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - } else { - tma_gather4_gmem2smem_mc_cta2(dst_v + group_2 * 512, (&tmap_compressed_v), acc_stage_4 * 128 * 2 + cta_rank * 128, ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : ((1) ? group_2 * 4 : 0)), ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : ((1) ? group_2 * 4 + 1 : 0)), ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : ((1) ? group_2 * 4 + 2 : 0)), ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : ((1) ? group_2 * 4 + 3 : 0)), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - tma_gather4_gmem2smem_mc_cta2(dst_v + group_2 * 512 + 8192, (&tmap_compressed_v), acc_stage_4 * 128 * 2 + cta_rank * 128 + 64, ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : ((1) ? group_2 * 4 : 0)), ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : ((1) ? group_2 * 4 + 1 : 0)), ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : ((1) ? group_2 * 4 + 2 : 0)), ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : ((1) ? group_2 * 4 + 3 : 0)), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - } - } - } - v_stage_1 += 1; - if (v_stage_1 == 2) { v_stage_1 = 0; _phase_v_empty ^= 1; } - } - } - int last_tile = num_kv_tiles_3 - 1; - #pragma unroll - for (int ps_3 = 0; ps_3 < 2; ps_3++) { - mbarrier_wait(v_empty_addr + (v_stage_1) * 8, _phase_v_empty); - if (load_warp_rank == 0) { - if (cta_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(65536)) : "memory"); - } - } - } - int group_3 = load_warp_rank * 4 + lane; - if (lane < 4) { - int index_offset_3 = last_tile * 128 + ps_3 * 64 + group_3 * 4; - int raw_rows_3[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_3[0])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 3])) - : "r"(smem_indices_addr + (unsigned int)(index_offset_3 * 4))); - #pragma unroll - for (int acc_stage_5 = 0; acc_stage_5 < 2; acc_stage_5++) { - int dst_v_1 = ((acc_stage_5 == 0) ? smem_v_lo_addr + v_stage_1 * 32768 : smem_v_hi_addr + v_stage_1 * 32768); - if (first_kv_tile_1 + last_tile == 0) { - tma_gather4_gmem2smem_mc_cta2(dst_v_1 + group_3 * 512, (&tmap_swa_v), acc_stage_5 * 128 * 2 + cta_rank * 128, ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : ((1) ? group_3 * 4 : 0)), ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : ((1) ? group_3 * 4 + 1 : 0)), ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : ((1) ? group_3 * 4 + 2 : 0)), ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : ((1) ? group_3 * 4 + 3 : 0)), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - tma_gather4_gmem2smem_mc_cta2(dst_v_1 + group_3 * 512 + 8192, (&tmap_swa_v), acc_stage_5 * 128 * 2 + cta_rank * 128 + 64, ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : ((1) ? group_3 * 4 : 0)), ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : ((1) ? group_3 * 4 + 1 : 0)), ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : ((1) ? group_3 * 4 + 2 : 0)), ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : ((1) ? group_3 * 4 + 3 : 0)), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - } else { - tma_gather4_gmem2smem_mc_cta2(dst_v_1 + group_3 * 512, (&tmap_compressed_v), acc_stage_5 * 128 * 2 + cta_rank * 128, ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : ((1) ? group_3 * 4 : 0)), ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : ((1) ? group_3 * 4 + 1 : 0)), ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : ((1) ? group_3 * 4 + 2 : 0)), ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : ((1) ? group_3 * 4 + 3 : 0)), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - tma_gather4_gmem2smem_mc_cta2(dst_v_1 + group_3 * 512 + 8192, (&tmap_compressed_v), acc_stage_5 * 128 * 2 + cta_rank * 128 + 64, ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : ((1) ? group_3 * 4 : 0)), ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : ((1) ? group_3 * 4 + 1 : 0)), ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : ((1) ? group_3 * 4 + 2 : 0)), ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : ((1) ? group_3 * 4 + 3 : 0)), ((v_full_addr + (v_stage_1) * 8) & 0xFEFFFFFF), 1 << cta_rank); - } - } - } - v_stage_1 += 1; - if (v_stage_1 == 2) { v_stage_1 = 0; _phase_v_empty ^= 1; } - } - asm volatile("barrier.sync 2, 160;" ::: "memory"); - } - mbarrier_wait(q_empty_addr, _phase_q_empty_0); - _phase_q_empty_0 ^= 1; - } - } - // ---- Role: index_warp ---- - if (warp == 13) { - { // index_warp_main - const int index_dummy = 0; - int valid_iter_1 = 0; - #pragma unroll 1 - for (unsigned int work_idx_4 = cluster_id; work_idx_4 < total_work_items; work_idx_4 += num_clusters) { - int split_idx_3 = work_idx_4 % (unsigned int)(((1) ? 4 : 1)); - int query_idx_3 = work_idx_4 / (unsigned int)(((1) ? 4 : 1)); - int owner_col_base = split_idx_3 * (4 / ((1) ? 4 : 1)) * 128; - int* swa_row = swa_indices + (query_idx_3 * swa_index_stride); - int* compressed_row = compressed_indices + (query_idx_3 * compressed_index_stride); - int valid_slot_1 = valid_iter_1 & 1; - int valid_empty_phase = valid_iter_1 >> 1 & 1 ^ 1; - mbarrier_wait(index_valid_empty_addr + (valid_slot_1) * 8, valid_empty_phase); - #pragma unroll - for (int index_pass = 0; index_pass < 4 / ((1) ? 4 : 1); index_pass++) { - int index_offset_4 = index_pass * 128 + lane * 4; - int combined_col = owner_col_base + index_offset_4; - int* index_src = ((combined_col < 128) ? swa_row + combined_col : compressed_row + (combined_col - 128)); - int values[4]; - #pragma unroll - for (int i = 0; i < 4; i++) { - values[i] = -1; - } - if (combined_col + 4 <= sparse_topk) { - int _vec_load_0[4]; - { - const int4* _ivptr_0 = reinterpret_cast(index_src + 0); - int4 _ivld_0; - _ivld_0 = *_ivptr_0; - _vec_load_0[0 + 0] = _ivld_0.x; - _vec_load_0[0 + 1] = _ivld_0.y; - _vec_load_0[0 + 2] = _ivld_0.z; - _vec_load_0[0 + 3] = _ivld_0.w; - } - #pragma unroll - for (int i_1 = 0; i_1 < 4; i_1++) { - values[i_1] = _vec_load_0[i_1]; - } - } - if (combined_col < sparse_topk && combined_col + 4 > sparse_topk) { - #pragma unroll - for (int i_2 = 0; i_2 < 4; i_2++) { - if (combined_col + i_2 < sparse_topk) { - values[i_2] = index_src[i_2]; - } - } - } - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(smem_indices_addr + (unsigned int)(index_offset_4 * 4)), "r"(values[0]), "r"(values[1]), "r"(values[2]), "r"(values[3]) : "memory"); - unsigned int valid_word = (((values[0] >= 0) ? 1 : 0) | ((values[1] >= 0) ? 1 : 0) << 1 | ((values[2] >= 0) ? 1 : 0) << 2 | ((values[3] >= 0) ? 1 : 0) << 3) << (lane & 7) * 4; - unsigned int _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, valid_word, 1); - valid_word = valid_word | _shfl_xor_0; - unsigned int _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, valid_word, 2); - valid_word = valid_word | _shfl_xor_1; - unsigned int _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, valid_word, 4); - valid_word = valid_word | _shfl_xor_2; - if ((lane & 7) == 0) { - smem_index_valid[valid_slot_1 * 36 + index_pass * 4 + (lane >> 3)] = valid_word; - } - } - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(index_valid_full_addr + (valid_slot_1) * 8); - asm volatile("barrier.sync 2, 160;" ::: "memory"); - asm volatile("barrier.sync 2, 160;" ::: "memory"); - valid_iter_1 = valid_iter_1 + 1; - } - } - } - // ---- Role: empty_warps ---- - if (warp >= 14 && warp <= 15) { - { // empty_warps_main - const int empty_dummy = 0; - #pragma unroll 1 - for (unsigned int work_idx_5 = cluster_id; work_idx_5 < total_work_items; work_idx_5 += num_clusters) { - asm volatile("tcgen05.fence::after_thread_sync;"); - } - } - } - - // Cleanup -} - -} // extern "C" diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_7a4b9556b3c84ebff6ce_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_7a4b9556b3c84ebff6ce_binding.cu deleted file mode 100644 index b7c794c30a4..00000000000 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_7a4b9556b3c84ebff6ce_binding.cu +++ /dev/null @@ -1,526 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -// Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_7a4b9556b3c84ebff6ce'. -#include -#include -#include -#include - -#include "tvm_ffi_utils.h" - -#include -#include -#include -#include -#include -#include -#include - -struct CakeTensorMap; -extern "C" __global__ void kernel_cake_dsv4_7a4b9556b3c84ebff6ce(CakeTensorMap const* tmap_q, CakeTensorMap const* tmap_swa_kv, CakeTensorMap const* tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int max_q_len, int batch_size, int has_sinks); - - -namespace cake_host_shim_1779b169b6ce3d43 { - -using tvm::ffi::Optional; -using tvm::ffi::TensorView; - -class ScopedCudaDevice { - public: - explicit ScopedCudaDevice(int device_id) { - cudaError_t error = cudaGetDevice(&previous_device_); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed before host-shim launch: cudaError=" - << static_cast(error); - if (previous_device_ != device_id) { - error = cudaSetDevice(device_id); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaSetDevice failed before host-shim launch for cuda:" - << device_id << ": cudaError=" << static_cast(error); - restore_ = true; - } - } - - ScopedCudaDevice(const ScopedCudaDevice&) = delete; - ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; - - ~ScopedCudaDevice() noexcept { - if (restore_) { - (void)cudaSetDevice(previous_device_); - } - } - - private: - int previous_device_ = -1; - bool restore_ = false; -}; - -inline int64_t CakeDeviceMultiprocessorCount(int device_id) { - constexpr int kMaxCachedCudaDevices = 64; - TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) - << "physical-SM-count cache does not cover cuda:" << device_id; - static std::atomic count_by_device[kMaxCachedCudaDevices]{}; - int cached = count_by_device[device_id].load(std::memory_order_acquire); - if (cached > 0) return cached; - - int count = 0; - cudaError_t error = cudaDeviceGetAttribute( - &count, cudaDevAttrMultiProcessorCount, device_id); - TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) - << "querying multiProcessorCount failed for cuda:" << device_id - << ": cudaError=" << static_cast(error) << ", count=" << count; - // Concurrent first launches may repeat the immutable device query, but all - // publication is atomic and every later hot-path lookup is one acquire load. - count_by_device[device_id].store(count, std::memory_order_release); - return count; -} - -inline void CheckCudaTensor(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) - << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; -} - -inline void CheckSameCudaDevice( - const TensorView& t, - const TensorView& reference, - const char* name, - const char* reference_name) { - TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) - << name << " must be on the same CUDA device as " << reference_name - << ": got cuda:" << t.device().device_id - << " versus cuda:" << reference.device().device_id; -} - -inline void CheckCurrentCudaDevice( - const TensorView& reference, - const char* reference_name) { - int current_device = -1; - cudaError_t error = cudaGetDevice(¤t_device); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed while validating " << reference_name - << ": cudaError=" << static_cast(error); - TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) - << "current CUDA device must match " << reference_name - << ": current=cuda:" << current_device - << ", tensor=cuda:" << reference.device().device_id; -} - -inline void CheckContiguous(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; -} - -inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { - DLDataType d = t.dtype(); - TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) - << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits - << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits - << ", lanes=" << (int)d.lanes << ")"; -} - -// A logical axis.outer(trailing) folds every source dim above the trailing -// dimensions. Shape products are independent of physical strides, so verify -// the leading dimensions form one dense row-major chain instead of inventing -// a "folded stride". The descriptor reads its exact adjacent physical step -// separately through stride[-(trailing + 1)]. -inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { - TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) - << name << " cannot fold leading dimensions above " << trailing - << " trailing dims from ndim=" << t.ndim(); - int outer_last = t.ndim() - trailing - 1; - if (outer_last <= 0) { - return; - } - int64_t step = t.stride(outer_last); - TVM_FFI_CHECK(step > 0, ValueError) - << name << " physical strides must be positive"; - int64_t expected = step; - for (int axis = outer_last - 1; axis >= 0; --axis) { - expected *= t.size(axis + 1); - if (t.size(axis) > 1) { - TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) - << name << " leading dims are not physically foldable above " << trailing - << " trailing dims: stride(" << axis << ")=" << t.stride(axis) - << ", expected " << expected; - } - } -} - -// Explicit immutable caller storage: initialize once, outside graph capture, -// exactly as the production TmaDeviceSlot path. The owner must retain the -// allocation and must never expose it as mutable scratch or rewrite its bytes. -template -static inline void PrepareImmutableCallerTmaWorkspace( - const tvm::ffi::Tensor& retained_workspace, - const CUtensorMap (&maps)[N], cudaStream_t stream) { - void* workspace = retained_workspace.data_ptr(); - CUcontext context = nullptr; - CUresult result = cuCtxGetCurrent(&context); - TVM_FFI_CHECK(result == CUDA_SUCCESS && context != nullptr, RuntimeError) - << "immutable TMA workspace requires an active CUDA context"; - CUcontext allocation_context = nullptr; - result = cuPointerGetAttribute( - &allocation_context, CU_POINTER_ATTRIBUTE_CONTEXT, - reinterpret_cast(workspace)); - TVM_FFI_CHECK(result == CUDA_SUCCESS && allocation_context == context, RuntimeError) - << "immutable TMA workspace must belong to the active CUDA context"; - unsigned long long buffer_id = 0; - result = cuPointerGetAttribute( - &buffer_id, CU_POINTER_ATTRIBUTE_BUFFER_ID, - reinterpret_cast(workspace)); - TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) - << "querying immutable TMA allocation identity failed"; - std::string key = std::to_string(reinterpret_cast(context)); - key += ":" + std::to_string(reinterpret_cast(workspace)); - key += ":" + std::to_string(buffer_id); - std::string descriptor_bytes(reinterpret_cast(maps), sizeof(maps)); - struct InitializedWorkspace { - tvm::ffi::Tensor owner; - std::string descriptor_bytes; - }; - static std::mutex mu; - // Keep the actual Tensor allocation alive: a caching allocator may reuse a - // freed suballocation without changing its enclosing CUDA buffer identity. - static auto* initialized = new std::unordered_map(); - std::lock_guard lock(mu); - auto it = initialized->find(key); - if (it != initialized->end()) { - TVM_FFI_CHECK(it->second.descriptor_bytes == descriptor_bytes, ValueError) - << "immutable TMA descriptors changed; provide new private descriptor storage"; - return; - } - CUstreamCaptureStatus capture_status = CU_STREAM_CAPTURE_STATUS_NONE; - result = cuStreamIsCapturing(reinterpret_cast(stream), &capture_status); - TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) - << "querying immutable TMA initialization capture state failed"; - TVM_FFI_CHECK(capture_status == CU_STREAM_CAPTURE_STATUS_NONE, RuntimeError) - << "immutable TMA storage must be initialized before CUDA Graph capture"; - // Synchronous host-to-device initialization completes before publishing the - // slot to any thread/stream, matching TmaDeviceSlot. The consumer retains its - // generated tensor-map acquire; the mutable upload's release fence is unchanged. - result = cuMemcpyHtoD(reinterpret_cast(workspace), maps, sizeof(maps)); - TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) - << "initializing immutable caller TMA descriptors failed"; - initialized->emplace( - std::move(key), InitializedWorkspace{retained_workspace, std::move(descriptor_bytes)}); -} - -static inline void* CallerTmaWorkspaceSlot( - const tvm::ffi::TensorView& workspace, size_t slot) { - return static_cast(workspace.data_ptr()) + slot * sizeof(CUtensorMap); -} - -// 4D TMA descriptor for buffer 'tmap_q' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 3, ValueError) - << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - int64_t d2 = t.size(t.ndim() - 2); - int64_t d3 = t.size(t.ndim() - 3); - TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) - << "TMA source 'tmap_q' trailing dims must be positive"; - TVM_FFI_CHECK(d1 % 128 == 0, ValueError) - << "TMA source 'tmap_q' extent " << d1 - << " must divide exactly by " << 128; - uint64_t global_dim[4] = {(uint64_t)(128), (uint64_t)(d2), (uint64_t)((d1 / 128)), (uint64_t)(d3)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; - TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) - << "TMA box (128, 8, 1, 1) exceeds resolved global dims for 'tmap_q'"; - int64_t carrier_stride_0 = d1; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; - int64_t carrier_stride_1 = 128; - TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; - TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; - TVM_FFI_CHECK((carrier_stride_1 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; - int64_t carrier_stride_2 = (d2 * d1); - TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; - TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; - TVM_FFI_CHECK((carrier_stride_2 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; - uint64_t global_strides[3] = { - (uint64_t)((carrier_stride_0 * 8) / 8), - (uint64_t)((carrier_stride_1 * 8) / 8), - (uint64_t)((carrier_stride_2 * 8) / 8), - }; - uint32_t box_dim[4] = {128u, 8u, 1u, 1u}; - uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_swa_kv' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_swa_kv' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; - TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (128, 1) exceeds resolved global dims for 'tmap_swa_kv'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 8) / 8), - }; - uint32_t box_dim[2] = {128u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_compressed_kv' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_compressed_kv(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_compressed_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_compressed_kv' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_compressed_kv' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_compressed_kv"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_compressed_kv' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved a non-positive global dim"; - TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (128, 1) exceeds resolved global dims for 'tmap_compressed_kv'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 8) / 8), - }; - uint32_t box_dim[2] = {128u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_kv') failed: CUresult=" << (int)r; - return tm; -} - -void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_compressed_kv, TensorView arg_O, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_seq_lens, TensorView arg_cum_seq_lens_q, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_max_q_len, int64_t arg_batch_size, int64_t arg_has_sinks, tvm::ffi::Tensor arg_tma_descriptor_workspace, int64_t grid_x, int64_t grid_y, int64_t grid_z) { - DLDevice dev = arg_tmap_q.device(); - ScopedCudaDevice device_guard(dev.device_id); - CheckCudaTensor(arg_tmap_q, "tmap_q"); - CheckDtype(arg_tmap_q, "tmap_q", 1, 8, 1); - CheckContiguous(arg_tmap_q, "tmap_q"); - CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); - CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 1, 8, 1); - CheckCudaTensor(arg_tmap_compressed_kv, "tmap_compressed_kv"); - CheckDtype(arg_tmap_compressed_kv, "tmap_compressed_kv", 1, 8, 1); - CheckCudaTensor(arg_O, "O"); - CheckDtype(arg_O, "O", 4, 16, 1); - CheckContiguous(arg_O, "O"); - CheckCudaTensor(arg_swa_indices, "swa_indices"); - CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); - CheckContiguous(arg_swa_indices, "swa_indices"); - CheckCudaTensor(arg_compressed_indices, "compressed_indices"); - CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); - CheckContiguous(arg_compressed_indices, "compressed_indices"); - CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); - CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckCudaTensor(arg_seq_lens, "seq_lens"); - CheckDtype(arg_seq_lens, "seq_lens", 0, 32, 1); - CheckContiguous(arg_seq_lens, "seq_lens"); - CheckCudaTensor(arg_cum_seq_lens_q, "cum_seq_lens_q"); - CheckDtype(arg_cum_seq_lens_q, "cum_seq_lens_q", 0, 32, 1); - CheckContiguous(arg_cum_seq_lens_q, "cum_seq_lens_q"); - CheckCudaTensor(arg_sinks, "sinks"); - CheckDtype(arg_sinks, "sinks", 2, 32, 1); - CheckContiguous(arg_sinks, "sinks"); - CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); - CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); - CheckContiguous(arg_bmm1_scale, "bmm1_scale"); - CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); - CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); - CheckContiguous(arg_bmm2_scale, "bmm2_scale"); - TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) - << "scalar 'swa_index_stride' value " << arg_swa_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) - << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) - << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) - << "scalar 'sparse_topk' value " << arg_sparse_topk - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_max_q_len >= -2147483648LL && arg_max_q_len <= 2147483647LL, ValueError) - << "scalar 'max_q_len' value " << arg_max_q_len - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_batch_size >= -2147483648LL && arg_batch_size <= 2147483647LL, ValueError) - << "scalar 'batch_size' value " << arg_batch_size - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) - << "scalar 'has_sinks' value " << arg_has_sinks - << " is outside i32 range [-2147483648, 2147483647]"; - CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); - CheckSameCudaDevice(arg_tmap_compressed_kv, arg_tmap_q, "tmap_compressed_kv", "tmap_q"); - CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); - CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); - CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); - CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); - CheckSameCudaDevice(arg_seq_lens, arg_tmap_q, "seq_lens", "tmap_q"); - CheckSameCudaDevice(arg_cum_seq_lens_q, arg_tmap_q, "cum_seq_lens_q", "tmap_q"); - CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); - CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); - CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); - CheckCudaTensor(arg_tma_descriptor_workspace, "tma_descriptor_workspace"); - CheckSameCudaDevice(arg_tma_descriptor_workspace, arg_tmap_q, "tma_descriptor_workspace", "tmap_q"); - CheckContiguous(arg_tma_descriptor_workspace, "tma_descriptor_workspace"); - CheckDtype(arg_tma_descriptor_workspace, "tma_descriptor_workspace", 1, 8, 1); - TVM_FFI_CHECK(arg_tma_descriptor_workspace.numel() >= 384, ValueError) - << "tma_descriptor_workspace requires at least 384 bytes"; - TVM_FFI_CHECK( - reinterpret_cast(arg_tma_descriptor_workspace.data_ptr()) % 128 == 0, - ValueError) - << "tma_descriptor_workspace must be 128-byte aligned"; - TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) - << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y - << ", " << grid_z << ")"; - - cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); - CUtensorMap h_tmap_q = EncodeTma_tmap_q(arg_tmap_q); - CakeTensorMap const* p_tmap_q = static_cast(CallerTmaWorkspaceSlot(arg_tma_descriptor_workspace, 0u)); - CUtensorMap h_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); - CakeTensorMap const* p_tmap_swa_kv = static_cast(CallerTmaWorkspaceSlot(arg_tma_descriptor_workspace, 1u)); - CUtensorMap h_tmap_compressed_kv = EncodeTma_tmap_compressed_kv(arg_tmap_compressed_kv); - CakeTensorMap const* p_tmap_compressed_kv = static_cast(CallerTmaWorkspaceSlot(arg_tma_descriptor_workspace, 2u)); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); - int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); - int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); - int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); - int* p_seq_lens = static_cast(arg_seq_lens.data_ptr()); - int* p_cum_seq_lens_q = static_cast(arg_cum_seq_lens_q.data_ptr()); - float* p_sinks = static_cast(arg_sinks.data_ptr()); - float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); - float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); - int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; - int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; - int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; - int32_t v_sparse_topk = (int32_t)arg_sparse_topk; - int32_t v_max_q_len = (int32_t)arg_max_q_len; - int32_t v_batch_size = (int32_t)arg_batch_size; - int32_t v_has_sinks = (int32_t)arg_has_sinks; - void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_tmap_compressed_kv, &p_O, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_seq_lens, &p_cum_seq_lens_q, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_max_q_len, &v_batch_size, &v_has_sinks}; - - static std::mutex smem_mu; - static auto* smem_status_by_device = - new std::unordered_map(); - cudaError_t smem_status = cudaSuccess; - { - std::lock_guard lock(smem_mu); - auto it = smem_status_by_device->find(dev.device_id); - if (it == smem_status_by_device->end()) { - int current_device = -1; - smem_status = cudaGetDevice(¤t_device); - if (smem_status == cudaSuccess && current_device != dev.device_id) { - smem_status = cudaErrorInvalidDevice; - } - if (smem_status == cudaSuccess) { - smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_7a4b9556b3c84ebff6ce), - cudaFuncAttributeMaxDynamicSharedMemorySize, - 147072); - } - smem_status_by_device->emplace(dev.device_id, smem_status); - } else { - smem_status = it->second; - } - } - TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_7a4b9556b3c84ebff6ce failed: " - << cudaGetErrorString(smem_status); - dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); - dim3 block(512u, 1u, 1u); - const CUtensorMap workspace_maps[] = {h_tmap_q, h_tmap_swa_kv, h_tmap_compressed_kv}; - PrepareImmutableCallerTmaWorkspace(arg_tma_descriptor_workspace, workspace_maps, stream); - - cudaLaunchConfig_t config{}; - config.gridDim = grid; - config.blockDim = block; - config.dynamicSmemBytes = 147072u; - config.stream = stream; - cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_7a4b9556b3c84ebff6ce), kargs); - TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_7a4b9556b3c84ebff6ce failed: " - << cudaGetErrorString(launch_status); -} - -} // namespace cake_host_shim_1779b169b6ce3d43 - -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_1779b169b6ce3d43::Run); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_7a4b9556b3c84ebff6ce_kernel.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_7a4b9556b3c84ebff6ce_kernel.cu deleted file mode 100644 index ca25505d652..00000000000 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_7a4b9556b3c84ebff6ce_kernel.cu +++ /dev/null @@ -1,2625 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -typedef signed char int8_t; -typedef unsigned char uint8_t; -typedef unsigned short uint16_t; -typedef unsigned int uint32_t; -#if defined(__CUDACC_RTC__) -typedef unsigned long long uint64_t; -#else -typedef unsigned long uint64_t; -#endif -static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); -typedef signed int int32_t; -typedef short int int16_t; -struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; -struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; -static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); -static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); -template -struct __align__(128) CakeTensorMapPack { CakeTensorMap maps[N]; }; - -#if defined(__CUDACC_RTC__) -typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; -#else -#include -#endif - -static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); -static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); -#include -#include - -#define CAKE_INF CUDART_INF_F -#define TMEM_NCOLS 512 -#define TMEM_SOURCE_EXACT_SCRATCH_OFFSET 0 -#define NUM_Q_PIPE_STAGES 1 -#define NUM_KV_PIPE_STAGES 8 -#define NUM_PAGE_PIPE_STAGES 6 -#define NUM_SCORE_PIPE_STAGES 2 -#define NUM_STATS_PIPE_STAGES 2 -#define NUM_P_PIPE_STAGES 2 -#define NUM_O_PIPE_STAGES 1 -#define NUM_WORK_PIPE_STAGES 2 -#define NUM_THROTTLE_PIPE_STAGES 2 -#define SMEM_SMEM_Q_OFF 1024 -#define SMEM_SMEM_Q_STAGE_BYTES 1024 -#define SMEM_SMEM_Q_STRIDE 1024 -#define SMEM_SMEM_KV_OFF 5120 -#define SMEM_SMEM_KV_STAGE_BYTES 16384 -#define SMEM_SMEM_KV_STRIDE 16384 -#define SMEM_SMEM_V_OFF 5120 -#define SMEM_SMEM_V_STAGE_BYTES 16384 -#define SMEM_SMEM_V_STRIDE 16384 -#define SMEM_SMEM_P_OFF 136192 -#define SMEM_SMEM_P_STAGE_BYTES 1024 -#define SMEM_SMEM_P_STRIDE 1024 -#define SMEM_SMEM_PAGE_OFFSETS_OFF 138240 -#define SMEM_SMEM_PAGE_OFFSETS_STAGE_BYTES 1024 -#define SMEM_SMEM_PAGE_OFFSETS_STRIDE 1024 -#define SMEM_SMEM_O_OFF 144384 -#define SMEM_SMEM_O_STAGE_BYTES 2048 -#define SMEM_SMEM_O_STRIDE 2048 -#define SMEM_SMEM_SOFTMAX_OFF 146432 -#define SMEM_SMEM_SOFTMAX_STAGE_BYTES 128 -#define SMEM_SMEM_SOFTMAX_STRIDE 128 -#define SMEM_SMEM_SOFTMAX_U32_OFF 146432 -#define SMEM_SMEM_SOFTMAX_U32_STAGE_BYTES 128 -#define SMEM_SMEM_SOFTMAX_U32_STRIDE 128 -#define SMEM_SMEM_SOFTMAX_UNUSED_OFF 146560 -#define SMEM_SMEM_SOFTMAX_UNUSED_STAGE_BYTES 128 -#define SMEM_SMEM_SOFTMAX_UNUSED_STRIDE 128 -#define SMEM_SMEM_CORR_OFF 146688 -#define SMEM_SMEM_CORR_STAGE_BYTES 128 -#define SMEM_SMEM_CORR_STRIDE 128 -#define SMEM_SMEM_STATS_OFF 146816 -#define SMEM_SMEM_STATS_STAGE_BYTES 96 -#define SMEM_SMEM_STATS_STRIDE 96 -#define SMEM_WORK_RESPONSE_OFF 147008 -#define SMEM_WORK_RESPONSE_STAGE_BYTES 16 -#define SMEM_WORK_RESPONSE_STRIDE 16 -#define SMEM_TOTAL 147072 -#define THREADS 512 - -#include - -__device__ __forceinline__ uint32_t elect_sync() { - uint32_t pred = 0; - asm volatile( - "{\n\t" - ".reg .pred %%px;\n\t" - "elect.sync _|%%px, %1;\n\t" - "@%%px mov.s32 %0, 1;\n\t" - "}\n" - : "+r"(pred) - : "r"(0xFFFFFFFF)); - return pred; -} - - -__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { - asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" - :: "r"(mbar_addr), "r"(count) : "memory"); -} - -__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { - asm volatile("mbarrier.init.b64 [%0], %1;" - :: "l"(mbar_addr), "r"(count) : "memory"); -} - - -__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - - -// CTA-local pipelines have short, resident producer/consumer edges. Omitting -// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the -// explicit loop still makes this helper blocking until acquire succeeds. -__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE;\n\t" - "bra.uni LAB_WAIT;\n\t" - "DONE:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Source-faithful relaxed CTA wait used only by a typed protocol that does -// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. -__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, 10000000;\n\t" - "@P1 bra.uni DONE_RELAXED;\n\t" - "bra.uni LAB_WAIT_RELAXED;\n\t" - "DONE_RELAXED:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Exact source ports may request the PTX suspendTimeHint operand explicitly. -// The hint is expressed in nanoseconds and is kept separate from the canonical -// no-hint CTA helper so unrelated schedules retain their existing retry path. -__device__ __forceinline__ void mbarrier_wait_suspend( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_SUSPEND:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_SUSPEND;\n\t" - "bra.uni LAB_WAIT_SUSPEND;\n\t" - "DONE_SUSPEND:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE_CLUSTER;\n\t" - "bra.uni LAB_WAIT_CLUSTER;\n\t" - "DONE_CLUSTER:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - ".reg .u32 WAIT_ADDR;\n\t" - "mov.u32 WAIT_ADDR, %0;\n\t" - "LAB_WAIT_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [WAIT_ADDR], %1, %2;\n\t" - "@P1 bra.uni DONE_HINT;\n\t" - "bra.uni LAB_WAIT_HINT;\n\t" - "DONE_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -// Exact unqualified CTA wait used by source schedules whose PTX intentionally -// omits the acquire qualifier while retaining a typed suspendTimeHint operand. -__device__ __forceinline__ void mbarrier_wait_relaxed_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED_HINT:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra DONE_RELAXED_HINT;\n\t" - "bra LAB_WAIT_RELAXED_HINT;\n\t" - "DONE_RELAXED_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); -} - -__device__ __forceinline__ void mbarrier_wait_cluster_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" - "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" - "DONE_CLUSTER_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_suspend( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait_cluster(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); - } -} - - -__device__ __forceinline__ void tcgen05_mma_f8f6f4( - int taddr, uint64_t a_desc, uint64_t b_desc, - uint32_t i_desc, int enable_input_d) { - asm volatile( - "{\n\t" - ".reg .pred p;\n\t" - "setp.ne.b32 p, %4, 0;\n\t" - "tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], %1, %2, %3, p;\n\t" - "}\n" - :: "r"(taddr), "l"(a_desc), "l"(b_desc), - "r"(i_desc), "r"(enable_input_d) - : "memory"); -} - - -__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { - return (x & 0x3FFFFULL) >> 4ULL; -} - - -union MmaSmemDesc { - uint64_t u64; - uint32_t u32[2]; -}; - -__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { - MmaSmemDesc tmp; - tmp.u64 = smem_desc; - tmp.u32[0] += offset; - smem_desc = tmp.u64; -} - - -__device__ __forceinline__ void elect_commit(int mbar_addr) { - asm volatile( - "{\n\t" - ".reg .pred leader;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" - ".shared::cluster.b64 [%0];\n\t" - "}\n" - :: "r"(mbar_addr)); -} - - -__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { - asm volatile( - "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" - :: "r"(mbar_addr), "r"(bytes) : "memory"); -} - - -__device__ __forceinline__ void tmem_ld_x4(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x4.b32" - " {%0, %1, %2, %3}, [%4];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { - asm volatile( - "tcgen05.st.sync.aligned.32x32b.x4.b32" - " [%0], {%1, %2, %3, %4};" - :: "r"(tmem_addr), - "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); -} - - -__device__ __forceinline__ float approx_exp2(float x) { - float y; - asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float approx_rcp(float x) { - float y; - asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float max_noftz(float a, float b) { - float c; - asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); - return c; -} - - -__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { - asm("mul.rn.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { - asm("mul.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { - asm("add.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { - asm("add.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { - asm("sub.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { - asm("sub.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("add.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { - float2 r; - asm("sub.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("sub.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ void fma_scale_x32( - float* sv, const float2* scale2, const float2* neg_max2) -{ - float2* sv_2 = reinterpret_cast(sv); - #pragma unroll - for (int j = 0; j < 16; j++) - fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); -} - -__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("mul.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - - -// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) - -__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("add.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("add.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("add.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("add.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("add.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("add.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("add.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - - -__device__ __forceinline__ void fence_async_shared() { - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); -} - - -__device__ __forceinline__ uint64_t make_smem_desc(int addr) { - const int SBO = 1024; - return desc_encode(addr) - | (desc_encode(SBO) << 32ULL) - | (1ULL << 46ULL) - | (2ULL << 61ULL); -} - - -__device__ __forceinline__ void tma_2d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3}], [%4];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_4d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.4d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5}], [%6];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr) { - // Canonical .shared::cta form for non-multicast gather4, matching - // trtllm-gen / cuda_ptx and the PTX ISA qualifier order - // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, - // .shared::cluster is reserved for the multicast variant (ctaMask). - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { - asm volatile( - "tcgen05.commit.cta_group::1.mbarrier::arrive::one" - ".shared::cluster.b64 [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tmem_ld_x4_wait(float* dst, int addr) { - tmem_ld_x4(dst, addr); - asm volatile("tcgen05.wait::ld.sync.aligned;"); -} - - -__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { - uint32_t result; - asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" - : "=r"(result) : "r"(val)); - return result; -} - -extern "C" { - -__global__ __launch_bounds__(512, 1) void -kernel_cake_dsv4_7a4b9556b3c84ebff6ce(CakeTensorMap const* tmap_q, CakeTensorMap const* tmap_swa_kv, CakeTensorMap const* tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int max_q_len, int batch_size, int has_sinks) -{ - const int tid = threadIdx.x; - const int warp = make_warp_uniform(tid / 32); - const int lane = tid % 32; - - extern __shared__ __align__(1024) char smem_raw[]; - int smem; - smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); - - const int mbar_base = smem; - #define q_full_addr (mbar_base + 0) - #define q_empty_addr (mbar_base + 8) - #define kv_full_addr (mbar_base + 16) - #define kv_empty_addr (mbar_base + 80) - #define page_full_addr (mbar_base + 144) - #define page_empty_addr (mbar_base + 192) - #define score_full_addr (mbar_base + 240) - #define score_empty_addr (mbar_base + 256) - #define stats_full_addr (mbar_base + 272) - #define stats_empty_addr (mbar_base + 288) - #define p_full_addr (mbar_base + 304) - #define p_empty_addr (mbar_base + 320) - #define o_full_addr (mbar_base + 336) - #define o_empty_addr (mbar_base + 344) - #define work_full_addr (mbar_base + 352) - #define work_empty_addr (mbar_base + 368) - #define throttle_full_addr (mbar_base + 384) - #define throttle_empty_addr (mbar_base + 400) - - const int bid = blockIdx.x; - const int num_bids = gridDim.x; - - const int cta_rank = 0; - if (tid == 0) { - asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_q)) : "memory"); - asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_swa_kv)) : "memory"); - asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_compressed_kv)) : "memory"); - } - __syncthreads(); - - - // Kernel setup ops - uint8_t* smem_q = reinterpret_cast(smem_raw + 1024); - const int smem_q_addr = smem + 1024; - uint8_t* smem_kv = reinterpret_cast(smem_raw + 5120); - const int smem_kv_addr = smem + 5120; - uint8_t* smem_v = reinterpret_cast(smem_raw + 5120); - const int smem_v_addr = smem + 5120; - uint8_t* smem_p = reinterpret_cast(smem_raw + 136192); - const int smem_p_addr = smem + 136192; - int* smem_page_offsets = reinterpret_cast(smem_raw + 138240); - const int smem_page_offsets_addr = smem + 138240; - __nv_bfloat16* smem_o = reinterpret_cast<__nv_bfloat16*>(smem_raw + 144384); - const int smem_o_addr = smem + 144384; - float* smem_softmax = reinterpret_cast(smem_raw + 146432); - const int smem_softmax_addr = smem + 146432; - unsigned int* smem_softmax_u32 = reinterpret_cast(smem_raw + 146432); - const int smem_softmax_u32_addr = smem + 146432; - float* smem_softmax_unused = reinterpret_cast(smem_raw + 146560); - const int smem_softmax_unused_addr = smem + 146560; - float* smem_corr = reinterpret_cast(smem_raw + 146688); - const int smem_corr_addr = smem + 146688; - float* smem_stats = reinterpret_cast(smem_raw + 146816); - const int smem_stats_addr = smem + 146816; - unsigned int* work_response = reinterpret_cast(smem_raw + 147008); - const int work_response_addr = smem + 147008; - asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)(tmap_q)) : "memory"); - asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)(tmap_swa_kv)) : "memory"); - asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)(tmap_compressed_kv)) : "memory"); - - // Mbarrier init (18 pipeline groups, 0 ordered-sequence groups, 52 barriers) - // Mbarriers at smem_raw[0..416) - - if (warp == 0) { - uint32_t leader = elect_sync(); - if (leader) { - // --- pipeline 'q_pipe' --- - // q_full: 1 barriers, init_count=1 - mbarrier_init(smem + 0, 1); - // q_empty: 1 barriers, init_count=1 - mbarrier_init(smem + 8, 1); - // --- pipeline 'kv_pipe' --- - // kv_full: 8 barriers, init_count=1 - mbarrier_init(smem + 16, 1); - mbarrier_init(smem + 24, 1); - mbarrier_init(smem + 32, 1); - mbarrier_init(smem + 40, 1); - mbarrier_init(smem + 48, 1); - mbarrier_init(smem + 56, 1); - mbarrier_init(smem + 64, 1); - mbarrier_init(smem + 72, 1); - // kv_empty: 8 barriers, init_count=1 - mbarrier_init(smem + 80, 1); - mbarrier_init(smem + 88, 1); - mbarrier_init(smem + 96, 1); - mbarrier_init(smem + 104, 1); - mbarrier_init(smem + 112, 1); - mbarrier_init(smem + 120, 1); - mbarrier_init(smem + 128, 1); - mbarrier_init(smem + 136, 1); - // --- pipeline 'page_pipe' --- - // page_full: 6 barriers, init_count=32 - mbarrier_init(smem + 144, 32); - mbarrier_init(smem + 152, 32); - mbarrier_init(smem + 160, 32); - mbarrier_init(smem + 168, 32); - mbarrier_init(smem + 176, 32); - mbarrier_init(smem + 184, 32); - // page_empty: 6 barriers, init_count=128 - mbarrier_init(smem + 192, 128); - mbarrier_init(smem + 200, 128); - mbarrier_init(smem + 208, 128); - mbarrier_init(smem + 216, 128); - mbarrier_init(smem + 224, 128); - mbarrier_init(smem + 232, 128); - // --- pipeline 'score_pipe' --- - // score_full: 2 barriers, init_count=1 - mbarrier_init(smem + 240, 1); - mbarrier_init(smem + 248, 1); - // score_empty: 2 barriers, init_count=128 - mbarrier_init(smem + 256, 128); - mbarrier_init(smem + 264, 128); - // --- pipeline 'stats_pipe' --- - // stats_full: 2 barriers, init_count=128 - mbarrier_init(smem + 272, 128); - mbarrier_init(smem + 280, 128); - // stats_empty: 2 barriers, init_count=128 - mbarrier_init(smem + 288, 128); - mbarrier_init(smem + 296, 128); - // --- pipeline 'p_pipe' --- - // p_full: 2 barriers, init_count=128 - mbarrier_init(smem + 304, 128); - mbarrier_init(smem + 312, 128); - // p_empty: 2 barriers, init_count=1 - mbarrier_init(smem + 320, 1); - mbarrier_init(smem + 328, 1); - // --- pipeline 'o_pipe' --- - // o_full: 1 barriers, init_count=1 - mbarrier_init(smem + 336, 1); - // o_empty: 1 barriers, init_count=128 - mbarrier_init(smem + 344, 128); - // --- pipeline 'work_pipe' --- - // work_full: 2 barriers, init_count=1 - mbarrier_init(smem + 352, 1); - mbarrier_init(smem + 360, 1); - // work_empty: 2 barriers, init_count=512 - mbarrier_init(smem + 368, 512); - mbarrier_init(smem + 376, 512); - // --- pipeline 'throttle_pipe' --- - // throttle_full: 2 barriers, init_count=128 - mbarrier_init(smem + 384, 128); - mbarrier_init(smem + 392, 128); - // throttle_empty: 2 barriers, init_count=32 - mbarrier_init(smem + 400, 32); - mbarrier_init(smem + 408, 32); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - - __syncwarp(); - - // TMEM alloc (512 columns, 512 used) - volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 416); - if (warp == 0) { - int _tmem_hold = smem + 416; - asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); - __syncwarp(); - asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); - } - - __syncthreads(); - asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); - - const int taddr = tmem_addr_storage[0]; - - // Kernel post-init ops - const int tmem_source_exact_scratch = taddr; - - // ---- Ordered hardware-WG register redistribution ---- - // Inc phase consumes the registers released above. - if (warp >= 8 && warp <= 11) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 136;"); - } - - // ---- Role: softmax ---- - if (warp <= 3) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 152;"); - { // softmax_main - unsigned int work_stage = 0; - unsigned int score_stage = 0; - unsigned int score_phase = 0; - unsigned int stats_stage = 0; - unsigned int stats_phase = 1; - unsigned int p_stage = 0; - unsigned int p_phase = 1; - unsigned int work_x = blockIdx.x; - unsigned int work_y = blockIdx.y; - unsigned int work_z = blockIdx.z; - const int warp_in_wg = warp; - const int wg_tid = warp_in_wg * 32 + lane; - const int col_pair_base = wg_tid % 4 * 2; - const int ldtm_row_base = warp_in_wg * 32 + lane / 4; - float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; - unsigned int _phase_work_full = 0; - #pragma unroll 1 - for (unsigned int _work = 0; _work < 120; _work++) { - int q_begin = cum_seq_lens_q[work_z]; - int q_end = cum_seq_lens_q[work_z + 1]; - int q_len = q_end - q_begin; - int query_idx = (unsigned int)q_begin + work_x; - int active_topk = 0; - if (work_x < (unsigned int)q_len) { - int _max_4 = ((sparse_topk_lens[query_idx] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx] + sparse_topk_lens_offset) : (0)); - int _min_3 = ((_max_4) < (sparse_topk) ? (_max_4) : (sparse_topk)); - active_topk = _min_3; - } - int num_steps = (active_topk + 128 - 1) / 128; - float row_max_pair[2]; - float row_sum_pair[2]; - row_max_pair[0] = -CAKE_INF; - row_max_pair[1] = -CAKE_INF; - row_sum_pair[0] = 0.0f; - row_sum_pair[1] = 0.0f; - asm volatile("barrier.sync 1, 128;" ::: "memory"); - uint32_t _amf_u_0 = __float_as_uint(-3.4028235e+38f); - uint32_t _amf_mask_0 = -int32_t(_amf_u_0 >> 31) | 0x80000000u; - unsigned int _amf_enc_0 = _amf_u_0 ^ _amf_mask_0; - if (wg_tid < 8) { - smem_softmax_u32[wg_tid] = _amf_enc_0; - } - asm volatile("barrier.sync 1, 128;" ::: "memory"); - #pragma unroll 1 - for (int tile = 0; tile < num_steps; tile++) { - mbarrier_wait(score_full_addr + (score_stage) * 8, score_phase); - int score_col = ((score_stage == 0) ? 0 : 8); - float scores[8]; - float score_lo[4]; - float score_hi[4]; - int tile_end = (tile + 1) * 128; - int is_full_tile = ((active_topk % 128 == 0) ? 1 : 0); - is_full_tile = is_full_tile | ((tile_end < active_topk) ? 1 : 0); - if (is_full_tile != 0) { - asm volatile( - "tcgen05.ld.sync.aligned.16x256b.x1.b32" - " {%0, %1, %2, %3}, [%4];" - : "=r"(*reinterpret_cast(&score_lo[0])), "=r"(*reinterpret_cast(&score_lo[1])), "=r"(*reinterpret_cast(&score_lo[2])), "=r"(*reinterpret_cast(&score_lo[3])) - : "r"(taddr + (unsigned int)score_col)); - asm volatile( - "tcgen05.ld.sync.aligned.16x256b.x1.b32" - " {%0, %1, %2, %3}, [%4];" - : "=r"(*reinterpret_cast(&score_hi[0])), "=r"(*reinterpret_cast(&score_hi[1])), "=r"(*reinterpret_cast(&score_hi[2])), "=r"(*reinterpret_cast(&score_hi[3])) - : "r"(taddr + (unsigned int)score_col + 1048576)); - #pragma unroll - for (int c = 0; c < 4; c++) { - scores[c] = score_lo[c]; - scores[c + 4] = score_hi[c]; - } - } else { - asm volatile("" ::: "memory"); - asm volatile( - "tcgen05.ld.sync.aligned.16x256b.x1.b32" - " {%0, %1, %2, %3}, [%4];" - : "=r"(*reinterpret_cast(&score_lo[0])), "=r"(*reinterpret_cast(&score_lo[1])), "=r"(*reinterpret_cast(&score_lo[2])), "=r"(*reinterpret_cast(&score_lo[3])) - : "r"(taddr + (unsigned int)score_col)); - asm volatile( - "tcgen05.ld.sync.aligned.16x256b.x1.b32" - " {%0, %1, %2, %3}, [%4];" - : "=r"(*reinterpret_cast(&score_hi[0])), "=r"(*reinterpret_cast(&score_hi[1])), "=r"(*reinterpret_cast(&score_hi[2])), "=r"(*reinterpret_cast(&score_hi[3])) - : "r"(taddr + (unsigned int)score_col + 1048576)); - #pragma unroll - for (int c_1 = 0; c_1 < 4; c_1++) { - scores[c_1] = score_lo[c_1]; - scores[c_1 + 4] = score_hi[c_1]; - } - if (active_topk <= tile * 128 + ldtm_row_base) { - scores[0] = -CAKE_INF; - scores[1] = -CAKE_INF; - } - if (active_topk <= tile * 128 + ldtm_row_base + 8) { - scores[2] = -CAKE_INF; - scores[3] = -CAKE_INF; - } - if (active_topk <= tile * 128 + ldtm_row_base + 16) { - scores[4] = -CAKE_INF; - scores[5] = -CAKE_INF; - } - if (active_topk <= tile * 128 + ldtm_row_base + 24) { - scores[6] = -CAKE_INF; - scores[7] = -CAKE_INF; - } - asm volatile("" ::: "memory"); - } - float pair_max[2]; - float _max_5 = max_noftz(scores[0], scores[2]); - float _max_6 = max_noftz(scores[4], scores[6]); - float _max_7 = max_noftz(_max_5, _max_6); - pair_max[0] = _max_7; - float _max_8 = max_noftz(scores[1], scores[3]); - float _max_9 = max_noftz(scores[5], scores[7]); - float _max_10 = max_noftz(_max_8, _max_9); - pair_max[1] = _max_10; - #pragma unroll - for (int c_2 = 0; c_2 < 2; c_2++) { - float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, pair_max[c_2], 16); - float _max_11 = max_noftz(pair_max[c_2], _shfl_xor_0); - pair_max[c_2] = _max_11; - float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, pair_max[c_2], 8); - float _max_12 = max_noftz(pair_max[c_2], _shfl_xor_1); - pair_max[c_2] = _max_12; - } - float old_max_pair[2]; - float new_max_pair[2]; - #pragma unroll - for (int c_3 = 0; c_3 < 2; c_3++) { - old_max_pair[c_3] = row_max_pair[c_3]; - float _max_13 = max_noftz(row_max_pair[c_3], pair_max[c_3]); - new_max_pair[c_3] = _max_13; - } - if (lane < 8) { - uint32_t _amf_u_1 = __float_as_uint(new_max_pair[0]); - uint32_t _amf_mask_1 = -int32_t(_amf_u_1 >> 31) | 0x80000000u; - unsigned int _amf_enc_1 = _amf_u_1 ^ _amf_mask_1; - uint32_t _amf_u_2 = __float_as_uint(new_max_pair[1]); - uint32_t _amf_mask_2 = -int32_t(_amf_u_2 >> 31) | 0x80000000u; - unsigned int _amf_enc_2 = _amf_u_2 ^ _amf_mask_2; - atomicMax(&smem_softmax_u32[col_pair_base], _amf_enc_1); - atomicMax(&smem_softmax_u32[col_pair_base + 1], _amf_enc_2); - } - asm volatile("barrier.sync 1, 128;" ::: "memory"); - uint32_t _amf_u_3 = smem_softmax_u32[col_pair_base]; - uint32_t _amf_mask_3 = ((_amf_u_3 >> 31) - 1u) | 0x80000000u; - float _amf_dec_0 = __uint_as_float(_amf_u_3 ^ _amf_mask_3); - new_max_pair[0] = _amf_dec_0; - uint32_t _amf_u_4 = smem_softmax_u32[col_pair_base + 1]; - uint32_t _amf_mask_4 = ((_amf_u_4 >> 31) - 1u) | 0x80000000u; - float _amf_dec_1 = __uint_as_float(_amf_u_4 ^ _amf_mask_4); - new_max_pair[1] = _amf_dec_1; - mbarrier_wait(stats_empty_addr + (stats_stage) * 8, stats_phase); - float stats_pair[4]; - stats_pair[0] = old_max_pair[0]; - stats_pair[1] = old_max_pair[1]; - stats_pair[2] = new_max_pair[0]; - stats_pair[3] = new_max_pair[1]; - int stats_col = ((stats_stage == 0) ? 16 : 48); - tmem_st_x4_f32(taddr + (unsigned int)stats_col, stats_pair); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(stats_full_addr + (stats_stage) * 8); - stats_stage += 1; - if (stats_stage == 2) { stats_stage = 0; stats_phase ^= 1; } - float acc_scale_pair[2]; - #pragma unroll - for (int c_4 = 0; c_4 < 2; c_4++) { - float delta = softmax_scale_log2 * (old_max_pair[c_4] - new_max_pair[c_4]); - float _exp2_0 = approx_exp2(delta); - acc_scale_pair[c_4] = ((old_max_pair[c_4] > -CAKE_INF) ? _exp2_0 : 1.0f); - row_max_pair[c_4] = new_max_pair[c_4]; - } - asm volatile("" ::: "memory"); - float exp_values[8]; - #pragma unroll - for (int c_5 = 0; c_5 < 8; c_5++) { - float safe_max = ((new_max_pair[c_5 % 2] == -CAKE_INF) ? 0.0f : new_max_pair[c_5 % 2]); - float _exp2_1 = approx_exp2(scores[c_5] * softmax_scale_log2 - safe_max * softmax_scale_log2); - exp_values[c_5] = _exp2_1; - } - mbarrier_wait(p_empty_addr + (p_stage) * 8, p_phase); - float p_values[8]; - #pragma unroll - for (int c_6 = 0; c_6 < 8; c_6++) { - p_values[c_6] = exp_values[c_6] * 448.0f; - } - unsigned int regs_p[2]; - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(p_values[0]), "f"(p_values[1]), - "f"(p_values[2]), "f"(p_values[3])); - regs_p[0] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(p_values[4]), "f"(p_values[5]), - "f"(p_values[6]), "f"(p_values[7])); - regs_p[1] = _packed; - } - int mtx_idx = lane / 8; - int thr_row_idx = lane % 8; - int seg_col_idx = warp_in_wg * 2 + mtx_idx ^ thr_row_idx; - int stsm_offset = thr_row_idx * 128 + seg_col_idx * 16; - const void* _stmatrix_b8_ptr_5 = reinterpret_cast(reinterpret_cast(smem_p) + (p_stage * 1024 + (unsigned int)stsm_offset)); - uint64_t _stmatrix_b8_addr64_5; - asm volatile("cvta.to.shared.u64 %0, %1;" : "=l"(_stmatrix_b8_addr64_5) : "l"(_stmatrix_b8_ptr_5)); - uint32_t _stmatrix_b8_addr_5; - asm volatile("cvt.u32.u64 %0, %1;" : "=r"(_stmatrix_b8_addr_5) : "l"(_stmatrix_b8_addr64_5)); - asm volatile("stmatrix.sync.aligned.m16n8.x2.trans.shared.b8 [%0], {%1, %2};\n" - :: "r"(_stmatrix_b8_addr_5), "r"(regs_p[0]), "r"(regs_p[1]) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(p_full_addr + (p_stage) * 8); - p_stage += 1; - if (p_stage == 2) { p_stage = 0; p_phase ^= 1; } - float2 _f2_0 = make_float2(exp_values[0], exp_values[1]); - float2 _f2_1 = make_float2(exp_values[2], exp_values[3]); - float2 _f2_2 = make_float2(exp_values[4], exp_values[5]); - float2 _f2_3 = make_float2(exp_values[6], exp_values[7]); - float2 _f2_4 = make_float2(row_sum_pair[0], row_sum_pair[1]); - float2 _f2_5 = make_float2(acc_scale_pair[0], acc_scale_pair[1]); - float2 sum_p0_pair = fma_f32x2_rn_ftz(_f2_4, _f2_5, _f2_0); - float2 sum_p1_pair = add_f32x2(sum_p0_pair, _f2_1); - float2 sum_p2_pair = add_f32x2(sum_p1_pair, _f2_2); - float2 next_sum_pair = add_f32x2(sum_p2_pair, _f2_3); - row_sum_pair[0] = next_sum_pair.x; - row_sum_pair[1] = next_sum_pair.y; - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - mbarrier_arrive(score_empty_addr + (score_stage) * 8); - score_stage += 1; - if (score_stage == 2) { score_stage = 0; score_phase ^= 1; } - } - if (num_steps > 0) { - mbarrier_wait(stats_empty_addr + (stats_stage) * 8, stats_phase); - float final_stats_pair[4]; - final_stats_pair[0] = row_sum_pair[0]; - final_stats_pair[1] = row_sum_pair[1]; - final_stats_pair[2] = row_max_pair[0]; - final_stats_pair[3] = row_max_pair[1]; - int final_stats_col = ((stats_stage == 0) ? 16 : 48); - tmem_st_x4_f32(taddr + (unsigned int)final_stats_col, final_stats_pair); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(stats_full_addr + (stats_stage) * 8); - stats_stage += 1; - if (stats_stage == 2) { stats_stage = 0; stats_phase ^= 1; } - } - mbarrier_wait(work_full_addr + (work_stage) * 8, _phase_work_full); - uint32_t _clc_valid_3 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_3) - : "r"(work_response_addr + work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_9 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_9) - : "r"(work_response_addr + work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_10 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_10) - : "r"(work_response_addr + work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_11 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_11) - : "r"(work_response_addr + work_stage * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(work_empty_addr + work_stage * 8), "r"(0) : "memory"); - work_stage += 1; - if (work_stage == 2) { work_stage = 0; _phase_work_full ^= 1; } - if (_clc_valid_3 == 0) { - break; - } - work_x = _clc_ctaid_9; - work_y = _clc_ctaid_10; - work_z = _clc_ctaid_11; - } - } - } - // ---- Role: correction ---- - if (warp >= 4 && warp <= 7) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 136;"); - { // correction_main - unsigned int work_stage_1 = 0; - unsigned int stats_stage_1 = 0; - unsigned int stats_phase_1 = 0; - unsigned int work_x_1 = blockIdx.x; - unsigned int work_y_1 = blockIdx.y; - unsigned int work_z_1 = blockIdx.z; - const int corr_rank = warp - 4; - const int corr_tid = corr_rank * 32 + lane; - const int col_pair_base_1 = corr_tid % 4 * 2; - float softmax_scale_log2_1 = bmm1_scale[0] * 1.4426950408889634f; - float output_scale = bmm2_scale[0]; - unsigned int _phase_o_full_0 = 0; - unsigned int _phase_work_full_1 = 0; - #pragma unroll 1 - for (unsigned int _work_1 = 0; _work_1 < 120; _work_1++) { - int q_begin_1 = cum_seq_lens_q[work_z_1]; - int q_end_1 = cum_seq_lens_q[work_z_1 + 1]; - int q_len_1 = q_end_1 - q_begin_1; - int query_idx_1 = (unsigned int)q_begin_1 + work_x_1; - int active_topk_1 = 0; - if (work_x_1 < (unsigned int)q_len_1) { - int _max_14 = ((sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) : (0)); - int _min_4 = ((_max_14) < (sparse_topk) ? (_max_14) : (sparse_topk)); - active_topk_1 = _min_4; - } - int num_steps_1 = (active_topk_1 + 128 - 1) / 128; - int head_group = work_y_1 / 4; - int value_quarter = work_y_1 % 4; - if (num_steps_1 > 0) { - mbarrier_wait(stats_full_addr + (stats_stage_1) * 8, stats_phase_1); - mbarrier_arrive(stats_empty_addr + (stats_stage_1) * 8); - stats_stage_1 += 1; - if (stats_stage_1 == 2) { stats_stage_1 = 0; stats_phase_1 ^= 1; } - } - #pragma unroll 1 - for (int tile_1 = 1; tile_1 < num_steps_1; tile_1++) { - mbarrier_wait(stats_full_addr + (stats_stage_1) * 8, stats_phase_1); - int stats_col_1 = ((stats_stage_1 == 0) ? 16 : 48); - float _tmem_load_0[4]; - tmem_ld_x4(&_tmem_load_0[0], taddr + (unsigned int)stats_col_1); - asm volatile("tcgen05.wait::ld.sync.aligned;"); - float scale_pair[2]; - #pragma unroll - for (int c_7 = 0; c_7 < 2; c_7++) { - float max_diff = _tmem_load_0[c_7] - _tmem_load_0[c_7 + 2]; - float _exp2_2 = approx_exp2(softmax_scale_log2_1 * max_diff); - scale_pair[c_7] = ((max_diff != 0.0f) ? _exp2_2 : 1.0f); - } - mbarrier_arrive(stats_empty_addr + (stats_stage_1) * 8); - stats_stage_1 += 1; - if (stats_stage_1 == 2) { stats_stage_1 = 0; stats_phase_1 ^= 1; } - mbarrier_wait(o_full_addr, _phase_o_full_0); - _phase_o_full_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - int needs_scale = ((scale_pair[0] != 1.0f) ? 1 : 0); - needs_scale = needs_scale | ((scale_pair[1] != 1.0f) ? 1 : 0); - int _vote_0 = __any_sync(0xFFFFFFFF, needs_scale != 0); - needs_scale = _vote_0; - if (needs_scale != 0) { - float values_lo[4]; - float values_hi[4]; - asm volatile( - "tcgen05.ld.sync.aligned.16x256b.x1.b32" - " {%0, %1, %2, %3}, [%4];" - : "=r"(*reinterpret_cast(&values_lo[0])), "=r"(*reinterpret_cast(&values_lo[1])), "=r"(*reinterpret_cast(&values_lo[2])), "=r"(*reinterpret_cast(&values_lo[3])) - : "r"(taddr + 80)); - asm volatile( - "tcgen05.ld.sync.aligned.16x256b.x1.b32" - " {%0, %1, %2, %3}, [%4];" - : "=r"(*reinterpret_cast(&values_hi[0])), "=r"(*reinterpret_cast(&values_hi[1])), "=r"(*reinterpret_cast(&values_hi[2])), "=r"(*reinterpret_cast(&values_hi[3])) - : "r"(taddr + 80 + 1048576)); - #pragma unroll - for (int pair = 0; pair < 4; pair++) { - asm volatile("" ::: "memory"); - values_lo[pair] = values_lo[pair] * scale_pair[pair % 2]; - values_hi[pair] = values_hi[pair] * scale_pair[pair % 2]; - } - asm volatile( - "tcgen05.st.sync.aligned.16x256b.x1.b32" - " [%0], {%1, %2, %3, %4};" - :: "r"(taddr + 80), "r"(*reinterpret_cast(&values_lo[0])), "r"(*reinterpret_cast(&values_lo[1])), "r"(*reinterpret_cast(&values_lo[2])), "r"(*reinterpret_cast(&values_lo[3]))); - asm volatile( - "tcgen05.st.sync.aligned.16x256b.x1.b32" - " [%0], {%1, %2, %3, %4};" - :: "r"(taddr + 80 + 1048576), "r"(*reinterpret_cast(&values_hi[0])), "r"(*reinterpret_cast(&values_hi[1])), "r"(*reinterpret_cast(&values_hi[2])), "r"(*reinterpret_cast(&values_hi[3]))); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - } - mbarrier_arrive(o_empty_addr); - } - if (num_steps_1 > 0) { - mbarrier_wait(stats_full_addr + (stats_stage_1) * 8, stats_phase_1); - int final_stats_col_1 = ((stats_stage_1 == 0) ? 16 : 48); - float _tmem_load_1[4]; - tmem_ld_x4(&_tmem_load_1[0], taddr + (unsigned int)final_stats_col_1); - asm volatile("tcgen05.wait::ld.sync.aligned;"); - mbarrier_arrive(stats_empty_addr + (stats_stage_1) * 8); - stats_stage_1 += 1; - if (stats_stage_1 == 2) { stats_stage_1 = 0; stats_phase_1 ^= 1; } - float pair_sum[2]; - pair_sum[0] = _tmem_load_1[0]; - pair_sum[1] = _tmem_load_1[1]; - float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, pair_sum[0], 16); - pair_sum[0] = pair_sum[0] + _shfl_xor_2; - float _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, pair_sum[1], 16); - pair_sum[1] = pair_sum[1] + _shfl_xor_3; - float _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, pair_sum[0], 8); - pair_sum[0] = pair_sum[0] + _shfl_xor_4; - float _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, pair_sum[1], 8); - pair_sum[1] = pair_sum[1] + _shfl_xor_5; - float _shfl_xor_6 = __shfl_xor_sync(0xFFFFFFFF, pair_sum[0], 4); - pair_sum[0] = pair_sum[0] + _shfl_xor_6; - float _shfl_xor_7 = __shfl_xor_sync(0xFFFFFFFF, pair_sum[1], 4); - pair_sum[1] = pair_sum[1] + _shfl_xor_7; - if (lane < 4) { - smem_corr[corr_rank * 8 + col_pair_base_1] = pair_sum[0]; - smem_corr[corr_rank * 8 + col_pair_base_1 + 1] = pair_sum[1]; - } - asm volatile("barrier.sync 2, 128;" ::: "memory"); - float sum0 = smem_corr[col_pair_base_1] + smem_corr[8 + col_pair_base_1] + smem_corr[16 + col_pair_base_1] + smem_corr[24 + col_pair_base_1]; - float sum1 = smem_corr[col_pair_base_1 + 1] + smem_corr[8 + col_pair_base_1 + 1] + smem_corr[16 + col_pair_base_1 + 1] + smem_corr[24 + col_pair_base_1 + 1]; - if (has_sinks != 0) { - float _exp2_3 = approx_exp2(sinks[head_group * 8 + col_pair_base_1] * 1.4426950408889634f - _tmem_load_1[2] * softmax_scale_log2_1); - sum0 = sum0 + _exp2_3; - float _exp2_4 = approx_exp2(sinks[head_group * 8 + col_pair_base_1 + 1] * 1.4426950408889634f - _tmem_load_1[3] * softmax_scale_log2_1); - sum1 = sum1 + _exp2_4; - } - float _rcp_0 = approx_rcp(sum0); - float scale0 = ((sum0 > 0.0f) ? output_scale * 0.002232142857142857f * _rcp_0 : 0.0f); - float _rcp_1 = approx_rcp(sum1); - float scale1 = ((sum1 > 0.0f) ? output_scale * 0.002232142857142857f * _rcp_1 : 0.0f); - mbarrier_wait(o_full_addr, _phase_o_full_0); - _phase_o_full_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - float output_lo[4]; - float output_hi[4]; - asm volatile( - "tcgen05.ld.sync.aligned.16x256b.x1.b32" - " {%0, %1, %2, %3}, [%4];" - : "=r"(*reinterpret_cast(&output_lo[0])), "=r"(*reinterpret_cast(&output_lo[1])), "=r"(*reinterpret_cast(&output_lo[2])), "=r"(*reinterpret_cast(&output_lo[3])) - : "r"(taddr + 80)); - asm volatile( - "tcgen05.ld.sync.aligned.16x256b.x1.b32" - " {%0, %1, %2, %3}, [%4];" - : "=r"(*reinterpret_cast(&output_hi[0])), "=r"(*reinterpret_cast(&output_hi[1])), "=r"(*reinterpret_cast(&output_hi[2])), "=r"(*reinterpret_cast(&output_hi[3])) - : "r"(taddr + 80 + 1048576)); - unsigned int packed_o[4]; - #pragma unroll - for (int pair_1 = 0; pair_1 < 4; pair_1++) { - asm volatile("" ::: "memory"); - output_lo[pair_1] = output_lo[pair_1] * ((pair_1 % 2 == 0) ? scale0 : scale1); - output_hi[pair_1] = output_hi[pair_1] * ((pair_1 % 2 == 0) ? scale0 : scale1); - } - float output_values[8]; - #pragma unroll - for (int c_8 = 0; c_8 < 4; c_8++) { - output_values[c_8] = output_lo[c_8]; - output_values[c_8 + 4] = output_hi[c_8]; - } - #pragma unroll - for (int _lp = 0; _lp < 4; _lp++) { - __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(output_values[_lp*2 + 0], output_values[_lp*2+1 + 0])); - packed_o[_lp] = *(uint32_t*)&_bf2; - } - int slice_idx = corr_rank / 2; - int warp_idx_in_slice = corr_rank % 2; - int mtx_idx_1 = lane / 8; - int thr_row_idx_1 = lane % 8; - int seg_col_idx_1 = warp_idx_in_slice * 4 + mtx_idx_1 ^ thr_row_idx_1; - int stsm_offset_1 = slice_idx * 8 * 128 + thr_row_idx_1 * 128 + seg_col_idx_1 * 16; - asm volatile("stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 [%0], {%1, %2, %3, %4};\n" - :: "l"(__cvta_generic_to_shared(reinterpret_cast(smem_o) + stsm_offset_1)), "r"(*reinterpret_cast(&packed_o[0])), "r"(*reinterpret_cast(&packed_o[1])), "r"(*reinterpret_cast(&packed_o[2])), "r"(*reinterpret_cast(&packed_o[3])) - : "memory"); - asm volatile("barrier.sync 2, 128;" ::: "memory"); - int copy_base = corr_tid * 16; - int smem_row = copy_base / 128; - int copy_smem_offset = copy_base ^ smem_row % 8 * 16; - int copy_row = smem_row % 8; - int copy_col = (smem_row / 8 * 128 + copy_base % 128) / 2; - unsigned int copy_vec[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(©_vec[0])), "=r"(*reinterpret_cast(©_vec[(0) + 1])), "=r"(*reinterpret_cast(©_vec[(0) + 2])), "=r"(*reinterpret_cast(©_vec[(0) + 3])) - : "r"(smem_o_addr + (unsigned int)copy_smem_offset)); - int head_groups = gridDim.y / 4; - int output_base = query_idx_1 * head_groups * 8 * 512 + head_group * 8 * 512 + value_quarter * 128; - reinterpret_cast(O + (output_base + copy_row * 512 + copy_col))[0] = reinterpret_cast(copy_vec)[0]; - mbarrier_arrive(o_empty_addr); - } - mbarrier_wait(work_full_addr + (work_stage_1) * 8, _phase_work_full_1); - uint32_t _clc_valid_4 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_4) - : "r"(work_response_addr + work_stage_1 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_12 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_12) - : "r"(work_response_addr + work_stage_1 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_13 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_13) - : "r"(work_response_addr + work_stage_1 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_14 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_14) - : "r"(work_response_addr + work_stage_1 * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(work_empty_addr + work_stage_1 * 8), "r"(0) : "memory"); - work_stage_1 += 1; - if (work_stage_1 == 2) { work_stage_1 = 0; _phase_work_full_1 ^= 1; } - if (_clc_valid_4 == 0) { - break; - } - work_x_1 = _clc_ctaid_12; - work_y_1 = _clc_ctaid_13; - work_z_1 = _clc_ctaid_14; - } - asm volatile("barrier.sync 11, 128;" ::: "memory"); - if (corr_rank == 0) { - int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); - asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); - } - } - } - // ---- Role: mma ---- - if (warp == 8) { - { // mma_main - unsigned int work_stage_2 = 0; - unsigned int q_stage = 0; - unsigned int q_phase = 0; - unsigned int kv_stage = 0; - unsigned int kv_phase = 0; - unsigned int score_stage_1 = 0; - unsigned int score_phase_1 = 1; - unsigned int p_stage_1 = 0; - unsigned int p_phase_1 = 0; - unsigned int work_x_2 = blockIdx.x; - unsigned int work_y_2 = blockIdx.y; - unsigned int work_z_2 = blockIdx.z; - unsigned int _phase_o_empty_0 = 1; - unsigned int _phase_work_full_2 = 0; - #pragma unroll 1 - for (unsigned int _work_2 = 0; _work_2 < 120; _work_2++) { - int q_begin_2 = cum_seq_lens_q[work_z_2]; - int q_end_2 = cum_seq_lens_q[work_z_2 + 1]; - int q_len_2 = q_end_2 - q_begin_2; - int query_idx_2 = (unsigned int)q_begin_2 + work_x_2; - int active_topk_2 = 0; - if (work_x_2 < (unsigned int)q_len_2) { - int _max_3 = ((sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) : (0)); - int _min_2 = ((_max_3) < (sparse_topk) ? (_max_3) : (sparse_topk)); - active_topk_2 = _min_2; - } - int num_steps_2 = (active_topk_2 + 128 - 1) / 128; - int value_quarter_1 = work_y_2 % 4; - if (num_steps_2 > 0) { - mbarrier_wait(q_full_addr + (q_stage) * 8, q_phase); - int first_output = 1; - mbarrier_wait(score_empty_addr + (score_stage_1) * 8, score_phase_1); - int score_col_1 = ((score_stage_1 == 0) ? 0 : 8); - #pragma unroll - for (int hs = 0; hs < 3; hs++) { - int hs_sel = ((value_quarter_1 > hs) ? hs : hs + 1); - mbarrier_wait(kv_full_addr + (kv_stage) * 8, kv_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - int _mma_a_lo_0 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage) * 1024); - int _mma_b_lo_0 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (hs_sel) * 64); - { - uint64_t _mma_ss_a_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_0); - uint64_t _mma_ss_b_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_0); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134348816, ((hs == 0) ? 0 : 1)); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134348816, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134348816, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134348816, 1); - } - } - elect_commit(kv_empty_addr + (kv_stage) * 8); - kv_stage += 1; - if (kv_stage == 8) { kv_stage = 0; kv_phase ^= 1; } - } - mbarrier_wait(kv_full_addr + (kv_stage) * 8, kv_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - int _mma_a_lo_1 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage) * 1024); - int _mma_b_lo_1 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (value_quarter_1) * 64); - { - uint64_t _mma_ss_a_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_1); - uint64_t _mma_ss_b_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_1); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134348816, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134348816, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134348816, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134348816, 1); - } - } - kv_stage += 1; - if (kv_stage == 8) { kv_stage = 0; kv_phase ^= 1; } - elect_commit(score_full_addr + (score_stage_1) * 8); - score_stage_1 += 1; - if (score_stage_1 == 2) { score_stage_1 = 0; score_phase_1 ^= 1; } - #pragma unroll 1 - for (int _tile = 0; _tile < num_steps_2 - 1; _tile++) { - mbarrier_wait(score_empty_addr + (score_stage_1) * 8, score_phase_1); - int score_col_0 = ((score_stage_1 == 0) ? 0 : 8); - #pragma unroll - for (int hs_1 = 0; hs_1 < 3; hs_1++) { - int hs_sel_1 = ((value_quarter_1 > hs_1) ? hs_1 : hs_1 + 1); - mbarrier_wait(kv_full_addr + (kv_stage) * 8, kv_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - int _mma_a_lo_2 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage) * 1024); - int _mma_b_lo_2 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (hs_sel_1) * 64); - { - uint64_t _mma_ss_a_desc_2 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_2); - uint64_t _mma_ss_b_desc_2 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_2); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134348816, ((hs_1 == 0) ? 0 : 1)); - } - incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134348816, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134348816, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134348816, 1); - } - } - elect_commit(kv_empty_addr + (kv_stage) * 8); - kv_stage += 1; - if (kv_stage == 8) { kv_stage = 0; kv_phase ^= 1; } - } - mbarrier_wait(kv_full_addr + (kv_stage) * 8, kv_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - int _mma_a_lo_3 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage) * 1024); - int _mma_b_lo_3 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (value_quarter_1) * 64); - { - uint64_t _mma_ss_a_desc_3 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_3); - uint64_t _mma_ss_b_desc_3 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_3); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134348816, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134348816, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134348816, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134348816, 1); - } - } - kv_stage += 1; - if (kv_stage == 8) { kv_stage = 0; kv_phase ^= 1; } - elect_commit(score_full_addr + (score_stage_1) * 8); - score_stage_1 += 1; - if (score_stage_1 == 2) { score_stage_1 = 0; score_phase_1 ^= 1; } - unsigned int v_cur = (kv_stage + 3) % 8; - mbarrier_wait(p_full_addr + (p_stage_1) * 8, p_phase_1); - mbarrier_wait(o_empty_addr, _phase_o_empty_0); - _phase_o_empty_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - int _mma_a_lo_4 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (v_cur) * 1024); - int _mma_b_lo_4 = make_warp_uniform(((((smem_p_addr) >> 4) & 0x3FFF) | 0x400000) + (p_stage_1) * 64); - { - uint64_t _mma_ss_a_desc_4 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_4); - uint64_t _mma_ss_b_desc_4 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_4); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134381584, ((first_output) ? 0 : 1)); - } - incr_smem_desc_lo(_mma_ss_a_desc_4, 256U); - incr_smem_desc_lo(_mma_ss_b_desc_4, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134381584, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_4, 256U); - incr_smem_desc_lo(_mma_ss_b_desc_4, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134381584, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_4, 256U); - incr_smem_desc_lo(_mma_ss_b_desc_4, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134381584, 1); - } - } - elect_commit(o_full_addr); - elect_commit(kv_empty_addr + (v_cur) * 8); - elect_commit(p_empty_addr + (p_stage_1) * 8); - p_stage_1 += 1; - if (p_stage_1 == 2) { p_stage_1 = 0; p_phase_1 ^= 1; } - first_output = 0; - } - unsigned int v_tail = (kv_stage + 7) % 8; - mbarrier_wait(p_full_addr + (p_stage_1) * 8, p_phase_1); - mbarrier_wait(o_empty_addr, _phase_o_empty_0); - _phase_o_empty_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - int _mma_a_lo_5 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (v_tail) * 1024); - int _mma_b_lo_5 = make_warp_uniform(((((smem_p_addr) >> 4) & 0x3FFF) | 0x400000) + (p_stage_1) * 64); - { - uint64_t _mma_ss_a_desc_5 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_5); - uint64_t _mma_ss_b_desc_5 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_5); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134381584, ((first_output) ? 0 : 1)); - } - incr_smem_desc_lo(_mma_ss_a_desc_5, 256U); - incr_smem_desc_lo(_mma_ss_b_desc_5, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134381584, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_5, 256U); - incr_smem_desc_lo(_mma_ss_b_desc_5, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134381584, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_5, 256U); - incr_smem_desc_lo(_mma_ss_b_desc_5, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134381584, 1); - } - } - elect_commit(o_full_addr); - elect_commit(kv_empty_addr + (v_tail) * 8); - elect_commit(p_empty_addr + (p_stage_1) * 8); - p_stage_1 += 1; - if (p_stage_1 == 2) { p_stage_1 = 0; p_phase_1 ^= 1; } - elect_commit(q_empty_addr + (q_stage) * 8); - q_phase ^= 1; - } - mbarrier_wait(work_full_addr + (work_stage_2) * 8, _phase_work_full_2); - uint32_t _clc_valid_2 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_2) - : "r"(work_response_addr + work_stage_2 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_6 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_6) - : "r"(work_response_addr + work_stage_2 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_7 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_7) - : "r"(work_response_addr + work_stage_2 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_8 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_8) - : "r"(work_response_addr + work_stage_2 * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(work_empty_addr + work_stage_2 * 8), "r"(0) : "memory"); - work_stage_2 += 1; - if (work_stage_2 == 2) { work_stage_2 = 0; _phase_work_full_2 ^= 1; } - if (_clc_valid_2 == 0) { - break; - } - work_x_2 = _clc_ctaid_6; - work_y_2 = _clc_ctaid_7; - work_z_2 = _clc_ctaid_8; - } - #pragma unroll - for (int tail = 0; tail < 2; tail++) { - elect_commit(score_full_addr + (tail) * 8); - } - } - } - // ---- Role: page_loader ---- - if (warp == 9) { - { // page_loader_main - unsigned int work_stage_3 = 0; - unsigned int page_stage = 0; - unsigned int page_phase = 1; - unsigned int work_x_3 = blockIdx.x; - unsigned int work_y_3 = blockIdx.y; - unsigned int work_z_3 = blockIdx.z; - unsigned int _phase_work_full_3 = 0; - #pragma unroll 1 - for (unsigned int _work_3 = 0; _work_3 < 120; _work_3++) { - int q_begin_3 = cum_seq_lens_q[work_z_3]; - int q_end_3 = cum_seq_lens_q[work_z_3 + 1]; - int q_len_3 = q_end_3 - q_begin_3; - int query_idx_3 = (unsigned int)q_begin_3 + work_x_3; - int active_topk_3 = 0; - if (work_x_3 < (unsigned int)q_len_3) { - int _max_0 = ((sparse_topk_lens[query_idx_3] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_3] + sparse_topk_lens_offset) : (0)); - int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); - active_topk_3 = _min_0; - } - int num_steps_3 = (active_topk_3 + 128 - 1) / 128; - int num_pairs = (num_steps_3 + 1) / 2; - int _max_1 = ((active_topk_3 - 1 & -4) > (0) ? (active_topk_3 - 1 & -4) : (0)); - int last_aligned = _max_1; - int* swa_row = swa_indices + (query_idx_3 * swa_index_stride); - int* compressed_row = compressed_indices + (query_idx_3 * compressed_index_stride); - #pragma unroll 1 - for (int pair_2 = 0; pair_2 < num_pairs; pair_2++) { - int lane_base = lane * 4; - #pragma unroll - for (int duplicate = 0; duplicate < 2; duplicate++) { - mbarrier_wait(page_empty_addr + (page_stage) * 8, page_phase); - int page0 = pair_2 * 2 * 128 + lane_base; - int page1 = page0 + 128; - int safe0 = ((page0 <= last_aligned) ? page0 : last_aligned); - int safe1 = ((page1 <= last_aligned) ? page1 : last_aligned); - int* src0 = ((safe0 < 128) ? swa_row + safe0 : compressed_row + (safe0 - 128)); - int* src1 = ((safe1 < 128) ? swa_row + safe1 : compressed_row + (safe1 - 128)); - asm volatile("cp.async.ca.shared::cta.global [%0], [%1], 16;" - :: "r"(smem_page_offsets_addr + page_stage * 1024 + (unsigned int)(lane_base * 4)), "l"(src0)); - asm volatile("cp.async.ca.shared::cta.global [%0], [%1], 16;" - :: "r"(smem_page_offsets_addr + page_stage * 1024 + (unsigned int)((128 + lane_base) * 4)), "l"(src1)); - asm volatile( - "{\n\t" - "cp.async.mbarrier.arrive.shared::cta.b64 [%0];\n\t" - "}" - :: "r"(page_full_addr + (page_stage) * 8) : "memory"); - mbarrier_arrive(page_full_addr + (page_stage) * 8); - page_stage += 1; - if (page_stage == 6) { page_stage = 0; page_phase ^= 1; } - } - } - mbarrier_wait(work_full_addr + (work_stage_3) * 8, _phase_work_full_3); - uint32_t _clc_valid_0 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_0) - : "r"(work_response_addr + work_stage_3 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_0 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_0) - : "r"(work_response_addr + work_stage_3 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_1 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_1) - : "r"(work_response_addr + work_stage_3 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_2 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_2) - : "r"(work_response_addr + work_stage_3 * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(work_empty_addr + work_stage_3 * 8), "r"(0) : "memory"); - work_stage_3 += 1; - if (work_stage_3 == 2) { work_stage_3 = 0; _phase_work_full_3 ^= 1; } - if (_clc_valid_0 == 0) { - break; - } - work_x_3 = _clc_ctaid_0; - work_y_3 = _clc_ctaid_1; - work_z_3 = _clc_ctaid_2; - } - } - } - // ---- Role: scheduler ---- - if (warp == 10) { - { // scheduler_main - asm volatile("griddepcontrol.wait;" ::: "memory"); - unsigned int work_stage_4 = 0; - unsigned int throttle_stage = 0; - unsigned int throttle_phase = 0; - unsigned int _phase_work_empty = 1; - unsigned int _phase_work_full_4 = 0; - #pragma unroll 1 - for (unsigned int _work_4 = 0; _work_4 < 120; _work_4++) { - mbarrier_wait(throttle_full_addr + (throttle_stage) * 8, throttle_phase); - mbarrier_arrive(throttle_empty_addr + (throttle_stage) * 8); - throttle_stage += 1; - if (throttle_stage == 2) { throttle_stage = 0; throttle_phase ^= 1; } - mbarrier_wait(work_empty_addr + (work_stage_4) * 8, _phase_work_empty); - if (elect_sync()) { - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [remAddr32], %2;\n\t" - "}" - :: "r"(work_full_addr + work_stage_4 * 8), "r"(0), "r"((uint32_t)(16)) : "memory"); - asm volatile( - "fence.proxy.async.shared::cta;\n\t" - "clusterlaunchcontrol.try_cancel.async.shared::cta" - ".mbarrier::complete_tx::bytes.b128" - " [%0], [%1];" - :: "r"(work_response_addr + work_stage_4 * 16 + 0 * 16), "r"(work_full_addr + work_stage_4 * 8) - : "memory"); - } - mbarrier_wait(work_full_addr + (work_stage_4) * 8, _phase_work_full_4); - uint32_t _clc_valid_5 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_5) - : "r"(work_response_addr + work_stage_4 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_15 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_15) - : "r"(work_response_addr + work_stage_4 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_16 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_16) - : "r"(work_response_addr + work_stage_4 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_17 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_17) - : "r"(work_response_addr + work_stage_4 * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(work_empty_addr + work_stage_4 * 8), "r"(0) : "memory"); - work_stage_4 += 1; - if (work_stage_4 == 2) { work_stage_4 = 0; _phase_work_empty ^= 1; _phase_work_full_4 ^= 1; } - if (_clc_valid_5 == 0) { - break; - } - } - } - } - // ---- Role: padding ---- - if (warp == 11) { - { // padding_main - unsigned int work_stage_5 = 0; - unsigned int work_x_4 = blockIdx.x; - unsigned int work_y_4 = blockIdx.y; - unsigned int work_z_4 = blockIdx.z; - unsigned int _phase_work_full_5 = 0; - #pragma unroll 1 - for (unsigned int _work_5 = 0; _work_5 < 120; _work_5++) { - mbarrier_wait(work_full_addr + (work_stage_5) * 8, _phase_work_full_5); - uint32_t _clc_valid_6 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_6) - : "r"(work_response_addr + work_stage_5 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_18 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_18) - : "r"(work_response_addr + work_stage_5 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_19 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_19) - : "r"(work_response_addr + work_stage_5 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_20 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_20) - : "r"(work_response_addr + work_stage_5 * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(work_empty_addr + work_stage_5 * 8), "r"(0) : "memory"); - work_stage_5 += 1; - if (work_stage_5 == 2) { work_stage_5 = 0; _phase_work_full_5 ^= 1; } - if (_clc_valid_6 == 0) { - break; - } - work_x_4 = _clc_ctaid_18; - work_y_4 = _clc_ctaid_19; - work_z_4 = _clc_ctaid_20; - } - } - } - // ---- Role: loader ---- - if (warp >= 12 && warp <= 15) { - asm volatile("setmaxnreg.dec.sync.aligned.u32 88;"); - { // loader_main - asm volatile("griddepcontrol.wait;" ::: "memory"); - const int load_rank = warp - 12; - unsigned int work_stage_6 = 0; - unsigned int throttle_stage_1 = 0; - unsigned int throttle_phase_1 = 1; - unsigned int q_stage_1 = 0; - unsigned int q_phase_1 = 1; - unsigned int kv_stage_1 = 0; - unsigned int kv_phase_1 = 1; - unsigned int page_k_stage = 0; - unsigned int page_k_phase = 0; - unsigned int page_v_stage = 1; - unsigned int page_v_phase = 0; - unsigned int work_x_5 = blockIdx.x; - unsigned int work_y_5 = blockIdx.y; - unsigned int work_z_5 = blockIdx.z; - unsigned int _phase_work_full_6 = 0; - #pragma unroll 1 - for (unsigned int _work_6 = 0; _work_6 < 120; _work_6++) { - int q_begin_4 = cum_seq_lens_q[work_z_5]; - int q_end_4 = cum_seq_lens_q[work_z_5 + 1]; - int q_len_4 = q_end_4 - q_begin_4; - int query_idx_4 = (unsigned int)q_begin_4 + work_x_5; - int active_topk_4 = 0; - if (work_x_5 < (unsigned int)q_len_4) { - int _max_2 = ((sparse_topk_lens[query_idx_4] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_4] + sparse_topk_lens_offset) : (0)); - int _min_1 = ((_max_2) < (sparse_topk) ? (_max_2) : (sparse_topk)); - active_topk_4 = _min_1; - } - int num_steps_4 = (active_topk_4 + 128 - 1) / 128; - int head_group_1 = work_y_5 / 4; - int value_quarter_2 = work_y_5 % 4; - mbarrier_wait(throttle_empty_addr + (throttle_stage_1) * 8, throttle_phase_1); - mbarrier_arrive(throttle_full_addr + (throttle_stage_1) * 8); - throttle_stage_1 += 1; - if (throttle_stage_1 == 2) { throttle_stage_1 = 0; throttle_phase_1 ^= 1; } - if (num_steps_4 > 0) { - mbarrier_wait(q_empty_addr + (q_stage_1) * 8, q_phase_1); - if (load_rank == 0) { - if (elect_sync()) { - mbarrier_arrive_expect_tx(q_full_addr + (q_stage_1) * 8, 4096); - #pragma unroll - for (int hs_2 = 0; hs_2 < 4; hs_2++) { - tma_4d_gmem2smem(smem_q_addr + (unsigned int)(hs_2 * 1024), tmap_q, 0, head_group_1 * 8, hs_2, query_idx_4, q_full_addr + (q_stage_1) * 8); - } - } - } - mbarrier_wait(page_full_addr + (page_k_stage) * 8, page_k_phase); - int page_base = 0; - #pragma unroll - for (int hs_3 = 0; hs_3 < 4; hs_3++) { - int hs_sel_2 = ((hs_3 == 3) ? value_quarter_2 : ((value_quarter_2 > hs_3) ? hs_3 : hs_3 + 1)); - mbarrier_wait(kv_empty_addr + (kv_stage_1) * 8, kv_phase_1); - if (load_rank == 0) { - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_1) * 8, 16384); - } - } - asm volatile("barrier.sync 9, 128;" ::: "memory"); - int rows_all[32]; - #pragma unroll - for (int group = 0; group < 8; group++) { - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&rows_all[group * 4])), "=r"(*reinterpret_cast(&rows_all[(group * 4) + 1])), "=r"(*reinterpret_cast(&rows_all[(group * 4) + 2])), "=r"(*reinterpret_cast(&rows_all[(group * 4) + 3])) - : "r"(smem_page_offsets_addr + page_k_stage * 1024 + (unsigned int)((page_base + load_rank * 4 + group * 16) * 4))); - } - if (elect_sync()) { - { - #pragma unroll - for (int group_1 = 0; group_1 < 8; group_1++) { - const int row_base = load_rank * 4 + group_1 * 16; - int dst_k = smem_kv_addr + kv_stage_1 * 16384 + (unsigned int)(row_base * 128); - tma_gather4_gmem2smem(dst_k, tmap_swa_kv, hs_sel_2 * 128, rows_all[group_1 * 4], rows_all[group_1 * 4 + 1], rows_all[group_1 * 4 + 2], rows_all[group_1 * 4 + 3], kv_full_addr + (kv_stage_1) * 8); - } - } - } - kv_stage_1 += 1; - if (kv_stage_1 == 8) { kv_stage_1 = 0; kv_phase_1 ^= 1; } - } - #pragma unroll 1 - for (int tile_2 = 0; tile_2 < num_steps_4 - 1; tile_2++) { - int next_tile = tile_2 + 1; - if ((next_tile & 1) == 0) { - mbarrier_wait(page_full_addr + (page_k_stage) * 8, page_k_phase); - } - int page_base_0 = (next_tile & 1) * 128; - #pragma unroll - for (int hs_4 = 0; hs_4 < 4; hs_4++) { - int hs_sel_3 = ((hs_4 == 3) ? value_quarter_2 : ((value_quarter_2 > hs_4) ? hs_4 : hs_4 + 1)); - mbarrier_wait(kv_empty_addr + (kv_stage_1) * 8, kv_phase_1); - if (load_rank == 0) { - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_1) * 8, 16384); - } - } - asm volatile("barrier.sync 9, 128;" ::: "memory"); - int rows_all_1[32]; - #pragma unroll - for (int group_2 = 0; group_2 < 8; group_2++) { - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&rows_all_1[group_2 * 4])), "=r"(*reinterpret_cast(&rows_all_1[(group_2 * 4) + 1])), "=r"(*reinterpret_cast(&rows_all_1[(group_2 * 4) + 2])), "=r"(*reinterpret_cast(&rows_all_1[(group_2 * 4) + 3])) - : "r"(smem_page_offsets_addr + page_k_stage * 1024 + (unsigned int)((page_base_0 + load_rank * 4 + group_2 * 16) * 4))); - } - if (elect_sync()) { - if (next_tile == 0) { - #pragma unroll - for (int group_3 = 0; group_3 < 8; group_3++) { - const int row_base_1 = load_rank * 4 + group_3 * 16; - int dst_k_1 = smem_kv_addr + kv_stage_1 * 16384 + (unsigned int)(row_base_1 * 128); - tma_gather4_gmem2smem(dst_k_1, tmap_swa_kv, hs_sel_3 * 128, rows_all_1[group_3 * 4], rows_all_1[group_3 * 4 + 1], rows_all_1[group_3 * 4 + 2], rows_all_1[group_3 * 4 + 3], kv_full_addr + (kv_stage_1) * 8); - } - } else { - #pragma unroll - for (int group_4 = 0; group_4 < 8; group_4++) { - const int row_base_2 = load_rank * 4 + group_4 * 16; - int dst_k_2 = smem_kv_addr + kv_stage_1 * 16384 + (unsigned int)(row_base_2 * 128); - tma_gather4_gmem2smem(dst_k_2, tmap_compressed_kv, hs_sel_3 * 128, rows_all_1[group_4 * 4], rows_all_1[group_4 * 4 + 1], rows_all_1[group_4 * 4 + 2], rows_all_1[group_4 * 4 + 3], kv_full_addr + (kv_stage_1) * 8); - } - } - } - kv_stage_1 += 1; - if (kv_stage_1 == 8) { kv_stage_1 = 0; kv_phase_1 ^= 1; } - } - if ((next_tile & 1) != 0 && next_tile < num_steps_4 - 1) { - mbarrier_arrive(page_empty_addr + (page_k_stage) * 8); - page_k_stage += 1; - if (page_k_stage == 6) { page_k_stage = 0; page_k_phase ^= 1; } - page_k_stage += 1; - if (page_k_stage == 6) { page_k_stage = 0; page_k_phase ^= 1; } - } - if ((tile_2 & 1) == 0) { - mbarrier_wait(page_full_addr + (page_v_stage) * 8, page_v_phase); - } - if ((tile_2 & 1) != 0) { - mbarrier_arrive(page_empty_addr + (page_v_stage) * 8); - page_v_stage += 1; - if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } - page_v_stage += 1; - if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } - } - } - int last_tile = num_steps_4 - 1; - mbarrier_arrive(page_empty_addr + (page_k_stage) * 8); - page_k_stage += 1; - if (page_k_stage == 6) { page_k_stage = 0; page_k_phase ^= 1; } - page_k_stage += 1; - if (page_k_stage == 6) { page_k_stage = 0; page_k_phase ^= 1; } - if ((last_tile & 1) == 0) { - mbarrier_wait(page_full_addr + (page_v_stage) * 8, page_v_phase); - } - mbarrier_arrive(page_empty_addr + (page_v_stage) * 8); - page_v_stage += 1; - if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } - page_v_stage += 1; - if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } - } - mbarrier_wait(work_full_addr + (work_stage_6) * 8, _phase_work_full_6); - uint32_t _clc_valid_1 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_1) - : "r"(work_response_addr + work_stage_6 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_3 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_3) - : "r"(work_response_addr + work_stage_6 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_4 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_4) - : "r"(work_response_addr + work_stage_6 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_5 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_5) - : "r"(work_response_addr + work_stage_6 * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(work_empty_addr + work_stage_6 * 8), "r"(0) : "memory"); - work_stage_6 += 1; - if (work_stage_6 == 2) { work_stage_6 = 0; _phase_work_full_6 ^= 1; } - if (_clc_valid_1 == 0) { - break; - } - work_x_5 = _clc_ctaid_3; - work_y_5 = _clc_ctaid_4; - work_z_5 = _clc_ctaid_5; - } - } - } - - // Cleanup -} - -} // extern "C" diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_9cf30aacb43b71034660_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_9cf30aacb43b71034660_binding.cu deleted file mode 100644 index b5754808d14..00000000000 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_9cf30aacb43b71034660_binding.cu +++ /dev/null @@ -1,385 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -// Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_9cf30aacb43b71034660'. -#include -#include -#include -#include - -#include "tvm_ffi_utils.h" - -#include -#include -#include -#include -#include -#include -#include - -extern "C" __global__ void kernel_cake_dsv4_9cf30aacb43b71034660(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks); - - -namespace cake_host_shim_8b6a5b581ef12c6c { - -using tvm::ffi::Optional; -using tvm::ffi::TensorView; - -class ScopedCudaDevice { - public: - explicit ScopedCudaDevice(int device_id) { - cudaError_t error = cudaGetDevice(&previous_device_); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed before host-shim launch: cudaError=" - << static_cast(error); - if (previous_device_ != device_id) { - error = cudaSetDevice(device_id); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaSetDevice failed before host-shim launch for cuda:" - << device_id << ": cudaError=" << static_cast(error); - restore_ = true; - } - } - - ScopedCudaDevice(const ScopedCudaDevice&) = delete; - ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; - - ~ScopedCudaDevice() noexcept { - if (restore_) { - (void)cudaSetDevice(previous_device_); - } - } - - private: - int previous_device_ = -1; - bool restore_ = false; -}; - -inline int64_t CakeDeviceMultiprocessorCount(int device_id) { - constexpr int kMaxCachedCudaDevices = 64; - TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) - << "physical-SM-count cache does not cover cuda:" << device_id; - static std::atomic count_by_device[kMaxCachedCudaDevices]{}; - int cached = count_by_device[device_id].load(std::memory_order_acquire); - if (cached > 0) return cached; - - int count = 0; - cudaError_t error = cudaDeviceGetAttribute( - &count, cudaDevAttrMultiProcessorCount, device_id); - TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) - << "querying multiProcessorCount failed for cuda:" << device_id - << ": cudaError=" << static_cast(error) << ", count=" << count; - // Concurrent first launches may repeat the immutable device query, but all - // publication is atomic and every later hot-path lookup is one acquire load. - count_by_device[device_id].store(count, std::memory_order_release); - return count; -} - -inline void CheckCudaTensor(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) - << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; -} - -inline void CheckSameCudaDevice( - const TensorView& t, - const TensorView& reference, - const char* name, - const char* reference_name) { - TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) - << name << " must be on the same CUDA device as " << reference_name - << ": got cuda:" << t.device().device_id - << " versus cuda:" << reference.device().device_id; -} - -inline void CheckCurrentCudaDevice( - const TensorView& reference, - const char* reference_name) { - int current_device = -1; - cudaError_t error = cudaGetDevice(¤t_device); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed while validating " << reference_name - << ": cudaError=" << static_cast(error); - TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) - << "current CUDA device must match " << reference_name - << ": current=cuda:" << current_device - << ", tensor=cuda:" << reference.device().device_id; -} - -inline void CheckContiguous(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; -} - -inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { - DLDataType d = t.dtype(); - TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) - << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits - << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits - << ", lanes=" << (int)d.lanes << ")"; -} - -// A logical axis.outer(trailing) folds every source dim above the trailing -// dimensions. Shape products are independent of physical strides, so verify -// the leading dimensions form one dense row-major chain instead of inventing -// a "folded stride". The descriptor reads its exact adjacent physical step -// separately through stride[-(trailing + 1)]. -inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { - TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) - << name << " cannot fold leading dimensions above " << trailing - << " trailing dims from ndim=" << t.ndim(); - int outer_last = t.ndim() - trailing - 1; - if (outer_last <= 0) { - return; - } - int64_t step = t.stride(outer_last); - TVM_FFI_CHECK(step > 0, ValueError) - << name << " physical strides must be positive"; - int64_t expected = step; - for (int axis = outer_last - 1; axis >= 0; --axis) { - expected *= t.size(axis + 1); - if (t.size(axis) > 1) { - TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) - << name << " leading dims are not physically foldable above " << trailing - << " trailing dims: stride(" << axis << ")=" << t.stride(axis) - << ", expected " << expected; - } - } -} - -// 4D TMA descriptor for buffer 'tmap_q' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 3, ValueError) - << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - int64_t d2 = t.size(t.ndim() - 2); - int64_t d3 = t.size(t.ndim() - 3); - TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) - << "TMA source 'tmap_q' trailing dims must be positive"; - TVM_FFI_CHECK(d1 % 64 == 0, ValueError) - << "TMA source 'tmap_q' extent " << d1 - << " must divide exactly by " << 64; - uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(d3)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) - << "TMA box (64, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; - int64_t carrier_stride_0 = d1; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; - int64_t carrier_stride_1 = 64; - TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; - TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; - TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; - int64_t carrier_stride_2 = (d2 * d1); - TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; - TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; - TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; - uint64_t global_strides[3] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - (uint64_t)((carrier_stride_1 * 16) / 8), - (uint64_t)((carrier_stride_2 * 16) / 8), - }; - uint32_t box_dim[4] = {64u, 64u, 1u, 1u}; - uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_swa_kv' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_swa_kv' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_kv'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; - return tm; -} - -void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_O, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t grid_x, int64_t grid_y, int64_t grid_z) { - DLDevice dev = arg_tmap_q.device(); - ScopedCudaDevice device_guard(dev.device_id); - CheckCudaTensor(arg_tmap_q, "tmap_q"); - CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); - CheckContiguous(arg_tmap_q, "tmap_q"); - CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); - CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 4, 16, 1); - CheckCudaTensor(arg_O, "O"); - CheckDtype(arg_O, "O", 4, 16, 1); - CheckContiguous(arg_O, "O"); - CheckCudaTensor(arg_swa_indices, "swa_indices"); - CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); - CheckContiguous(arg_swa_indices, "swa_indices"); - CheckCudaTensor(arg_compressed_indices, "compressed_indices"); - CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); - CheckContiguous(arg_compressed_indices, "compressed_indices"); - CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); - CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckCudaTensor(arg_sinks, "sinks"); - CheckDtype(arg_sinks, "sinks", 2, 32, 1); - CheckContiguous(arg_sinks, "sinks"); - CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); - CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); - CheckContiguous(arg_bmm1_scale, "bmm1_scale"); - CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); - CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); - CheckContiguous(arg_bmm2_scale, "bmm2_scale"); - TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) - << "scalar 'num_heads' value " << arg_num_heads - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) - << "scalar 'swa_index_stride' value " << arg_swa_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) - << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) - << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) - << "scalar 'sparse_topk' value " << arg_sparse_topk - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) - << "scalar 'has_sinks' value " << arg_has_sinks - << " is outside i32 range [-2147483648, 2147483647]"; - CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); - CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); - CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); - CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); - CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); - CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); - CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); - CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); - TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) - << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y - << ", " << grid_z << ")"; - - cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); - CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); - CUtensorMap p_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); - int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); - int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); - int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); - float* p_sinks = static_cast(arg_sinks.data_ptr()); - float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); - float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); - int32_t v_num_heads = (int32_t)arg_num_heads; - int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; - int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; - int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; - int32_t v_sparse_topk = (int32_t)arg_sparse_topk; - int32_t v_has_sinks = (int32_t)arg_has_sinks; - void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_O, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_has_sinks}; - - static std::mutex smem_mu; - static auto* smem_status_by_device = - new std::unordered_map(); - cudaError_t smem_status = cudaSuccess; - { - std::lock_guard lock(smem_mu); - auto it = smem_status_by_device->find(dev.device_id); - if (it == smem_status_by_device->end()) { - int current_device = -1; - smem_status = cudaGetDevice(¤t_device); - if (smem_status == cudaSuccess && current_device != dev.device_id) { - smem_status = cudaErrorInvalidDevice; - } - if (smem_status == cudaSuccess) { - smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_9cf30aacb43b71034660), - cudaFuncAttributeMaxDynamicSharedMemorySize, - 198528); - } - smem_status_by_device->emplace(dev.device_id, smem_status); - } else { - smem_status = it->second; - } - } - TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_9cf30aacb43b71034660 failed: " - << cudaGetErrorString(smem_status); - dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); - dim3 block(512u, 1u, 1u); - - cudaLaunchConfig_t config{}; - config.gridDim = grid; - config.blockDim = block; - config.dynamicSmemBytes = 198528u; - config.stream = stream; - cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_9cf30aacb43b71034660), kargs); - TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_9cf30aacb43b71034660 failed: " - << cudaGetErrorString(launch_status); -} - -} // namespace cake_host_shim_8b6a5b581ef12c6c - -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_8b6a5b581ef12c6c::Run); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_9cf30aacb43b71034660_kernel.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_9cf30aacb43b71034660_kernel.cu deleted file mode 100644 index 3398f54dc49..00000000000 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_9cf30aacb43b71034660_kernel.cu +++ /dev/null @@ -1,1634 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -typedef signed char int8_t; -typedef unsigned char uint8_t; -typedef unsigned short uint16_t; -typedef unsigned int uint32_t; -#if defined(__CUDACC_RTC__) -typedef unsigned long long uint64_t; -#else -typedef unsigned long uint64_t; -#endif -static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); -typedef signed int int32_t; -typedef short int int16_t; -struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; -struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; -static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); -static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); - -#if defined(__CUDACC_RTC__) -typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; -#else -#include -#endif - -static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); -static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); -#include -#include - -#define CAKE_INF CUDART_INF_F -#define TMEM_NCOLS 384 -#define TMEM_TMEM_OFFSET 0 -#define NUM_KV_PIPE_STAGES 4 -#define SMEM_SMEM_Q_OFF 1024 -#define SMEM_SMEM_Q_STAGE_BYTES 8192 -#define SMEM_SMEM_Q_STRIDE 8192 -#define SMEM_SMEM_KV_OFF 66560 -#define SMEM_SMEM_KV_STAGE_BYTES 32768 -#define SMEM_SMEM_KV_STRIDE 32768 -#define SMEM_SMEM_V_OFF 66560 -#define SMEM_SMEM_V_STAGE_BYTES 32768 -#define SMEM_SMEM_V_STRIDE 32768 -#define SMEM_SMEM_SUM_OFF 197632 -#define SMEM_SMEM_SUM_STAGE_BYTES 256 -#define SMEM_SMEM_SUM_STRIDE 256 -#define SMEM_SMEM_INDEX_FLAGS_OFF 198400 -#define SMEM_SMEM_INDEX_FLAGS_STAGE_BYTES 16 -#define SMEM_SMEM_INDEX_FLAGS_STRIDE 16 -#define SMEM_SMEM_INDICES_OFF 197888 -#define SMEM_SMEM_INDICES_STAGE_BYTES 512 -#define SMEM_SMEM_INDICES_STRIDE 512 -#define SMEM_TOTAL 198528 -#define THREADS 512 - -#include - -__device__ __forceinline__ uint32_t elect_sync() { - uint32_t pred = 0; - asm volatile( - "{\n\t" - ".reg .pred %%px;\n\t" - "elect.sync _|%%px, %1;\n\t" - "@%%px mov.s32 %0, 1;\n\t" - "}\n" - : "+r"(pred) - : "r"(0xFFFFFFFF)); - return pred; -} - - -__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { - asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" - :: "r"(mbar_addr), "r"(count) : "memory"); -} - -__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { - asm volatile("mbarrier.init.b64 [%0], %1;" - :: "l"(mbar_addr), "r"(count) : "memory"); -} - - -__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - - -// CTA-local pipelines have short, resident producer/consumer edges. Omitting -// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the -// explicit loop still makes this helper blocking until acquire succeeds. -__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE;\n\t" - "bra.uni LAB_WAIT;\n\t" - "DONE:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Source-faithful relaxed CTA wait used only by a typed protocol that does -// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. -__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, 10000000;\n\t" - "@P1 bra.uni DONE_RELAXED;\n\t" - "bra.uni LAB_WAIT_RELAXED;\n\t" - "DONE_RELAXED:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Exact source ports may request the PTX suspendTimeHint operand explicitly. -// The hint is expressed in nanoseconds and is kept separate from the canonical -// no-hint CTA helper so unrelated schedules retain their existing retry path. -__device__ __forceinline__ void mbarrier_wait_suspend( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_SUSPEND:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_SUSPEND;\n\t" - "bra.uni LAB_WAIT_SUSPEND;\n\t" - "DONE_SUSPEND:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE_CLUSTER;\n\t" - "bra.uni LAB_WAIT_CLUSTER;\n\t" - "DONE_CLUSTER:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - ".reg .u32 WAIT_ADDR;\n\t" - "mov.u32 WAIT_ADDR, %0;\n\t" - "LAB_WAIT_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [WAIT_ADDR], %1, %2;\n\t" - "@P1 bra.uni DONE_HINT;\n\t" - "bra.uni LAB_WAIT_HINT;\n\t" - "DONE_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -// Exact unqualified CTA wait used by source schedules whose PTX intentionally -// omits the acquire qualifier while retaining a typed suspendTimeHint operand. -__device__ __forceinline__ void mbarrier_wait_relaxed_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED_HINT:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra DONE_RELAXED_HINT;\n\t" - "bra LAB_WAIT_RELAXED_HINT;\n\t" - "DONE_RELAXED_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); -} - -__device__ __forceinline__ void mbarrier_wait_cluster_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" - "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" - "DONE_CLUSTER_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_suspend( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait_cluster(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); - } -} - - -__device__ __forceinline__ void tcgen05_mma_f16( - int taddr, uint64_t a_desc, uint64_t b_desc, - uint32_t i_desc, int enable_input_d) { - asm volatile( - "{\n\t" - ".reg .pred p;\n\t" - "setp.ne.b32 p, %4, 0;\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" - "}\n" - :: "r"(taddr), "l"(a_desc), "l"(b_desc), - "r"(i_desc), "r"(enable_input_d) - : "memory"); -} - - -__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { - return (x & 0x3FFFFULL) >> 4ULL; -} - - -union MmaSmemDesc { - uint64_t u64; - uint32_t u32[2]; -}; - -__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { - MmaSmemDesc tmp; - tmp.u64 = smem_desc; - tmp.u32[0] += offset; - smem_desc = tmp.u64; -} - - -__device__ __forceinline__ void elect_commit(int mbar_addr) { - asm volatile( - "{\n\t" - ".reg .pred leader;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" - ".shared::cluster.b64 [%0];\n\t" - "}\n" - :: "r"(mbar_addr)); -} - - -__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { - asm volatile( - "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" - :: "r"(mbar_addr), "r"(bytes) : "memory"); -} - - -__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7," - " %8, %9, %10, %11, %12, %13, %14, %15," - " %16, %17, %18, %19, %20, %21, %22, %23," - " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), - "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), - "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), - "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), - "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), - "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), - "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ float approx_exp2(float x) { - float y; - asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float approx_rcp(float x) { - float y; - asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float max_noftz(float a, float b) { - float c; - asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); - return c; -} - - -__device__ __forceinline__ float warp_reduce_max(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); - return val; -} - - -__device__ __forceinline__ float warp_reduce_sum(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val += __shfl_xor_sync(0xFFFFFFFF, val, offset); - return val; -} - - -__device__ __forceinline__ float row_max_reduce(float2 acc) { - return max_noftz(acc.x, acc.y); -} - - -__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { - #pragma unroll - for (int j = 0; j < 16; j++) { - if (j % 2 == 0) - acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); - else - acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); - } -} - - -__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { - asm("mul.rn.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { - asm("mul.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { - asm("add.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { - asm("add.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { - asm("sub.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { - asm("sub.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("add.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { - float2 r; - asm("sub.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("sub.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ void fma_scale_x32( - float* sv, const float2* scale2, const float2* neg_max2) -{ - float2* sv_2 = reinterpret_cast(sv); - #pragma unroll - for (int j = 0; j < 16; j++) - fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); -} - -__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("mul.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - - -// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) - -__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("add.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("add.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("add.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("add.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("add.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("add.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("add.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - - -__device__ __forceinline__ void fence_async_shared() { - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); -} - - -__device__ __forceinline__ uint64_t make_smem_desc(int addr) { - const int SBO = 1024; - return desc_encode(addr) - | (desc_encode(SBO) << 32ULL) - | (1ULL << 46ULL) - | (2ULL << 61ULL); -} - - -__device__ __forceinline__ void tma_2d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3}], [%4];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_4d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.4d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5}], [%6];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr) { - // Canonical .shared::cta form for non-multicast gather4, matching - // trtllm-gen / cuda_ptx and the PTX ISA qualifier order - // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, - // .shared::cluster is reserved for the multicast variant (ctaMask). - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { - asm volatile( - "tcgen05.commit.cta_group::1.mbarrier::arrive::one" - ".shared::cluster.b64 [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { - uint32_t result; - asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" - : "=r"(result) : "r"(val)); - return result; -} - -extern "C" { - -__global__ __launch_bounds__(512, 1) void -kernel_cake_dsv4_9cf30aacb43b71034660(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks) -{ - const int tid = threadIdx.x; - const int warp = make_warp_uniform(tid / 32); - const int lane = tid % 32; - - extern __shared__ __align__(1024) char smem_raw[]; - int smem; - smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); - - const int mbar_base = smem; - #define q_full_addr (mbar_base + 0) - #define kv_full_addr (mbar_base + 8) - #define kv_empty_addr (mbar_base + 40) - #define s_full_addr (mbar_base + 72) - #define p_full_addr (mbar_base + 80) - #define sum_ready_addr (mbar_base + 88) - #define o_done_addr (mbar_base + 96) - #define tmem_dealloc_addr (mbar_base + 104) - - const int bid = blockIdx.x; - const int num_bids = gridDim.x; - - const int cta_rank = 0; - - // Kernel setup ops - __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); - const int smem_q_addr = smem + 1024; - __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); - const int smem_kv_addr = smem + 66560; - __nv_bfloat16* smem_v = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); - const int smem_v_addr = smem + 66560; - float* smem_sum = reinterpret_cast(smem_raw + 197632); - const int smem_sum_addr = smem + 197632; - int* smem_index_flags = reinterpret_cast(smem_raw + 198400); - const int smem_index_flags_addr = smem + 198400; - int* smem_indices = reinterpret_cast(smem_raw + 197888); - const int smem_indices_addr = smem + 197888; - - // Mbarrier init (8 pipeline groups, 0 ordered-sequence groups, 14 barriers) - // Mbarriers at smem_raw[0..112) - - if (warp == 0) { - uint32_t leader = elect_sync(); - if (leader) { - // q_full: 1 barriers, init_count=1 - mbarrier_init(smem + 0, 1); - // --- pipeline 'kv_pipe' --- - // kv_full: 4 barriers, init_count=4 - mbarrier_init(smem + 8, 4); - mbarrier_init(smem + 16, 4); - mbarrier_init(smem + 24, 4); - mbarrier_init(smem + 32, 4); - // kv_empty: 4 barriers, init_count=1 - mbarrier_init(smem + 40, 1); - mbarrier_init(smem + 48, 1); - mbarrier_init(smem + 56, 1); - mbarrier_init(smem + 64, 1); - // s_full: 1 barriers, init_count=1 - mbarrier_init(smem + 72, 1); - // p_full: 1 barriers, init_count=128 - mbarrier_init(smem + 80, 128); - // sum_ready: 1 barriers, init_count=128 - mbarrier_init(smem + 88, 128); - // o_done: 1 barriers, init_count=1 - mbarrier_init(smem + 96, 1); - // tmem_dealloc: 1 barriers, init_count=256 - mbarrier_init(smem + 104, 256); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - - __syncwarp(); - - // TMEM alloc (512 columns, 384 used) - volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 112); - if (warp == 0) { - int _tmem_hold = smem + 112; - asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); - __syncwarp(); - asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); - } - - __syncthreads(); - asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); - - const int taddr = tmem_addr_storage[0]; - - // Kernel post-init ops - const int tmem_tmem = taddr; - - // ---- Ordered hardware-WG register redistribution ---- - // Dec phase frees registers before any WG attempts inc. - if (warp >= 8 && warp <= 11) { - asm volatile("setmaxnreg.dec.sync.aligned.u32 64;"); - } - - // ---- Role: softmax ---- - if (warp <= 3) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); - { // softmax_main - float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; - int query_idx = blockIdx.x >> 2; - int _max_0 = ((sparse_topk_lens[query_idx] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx] + sparse_topk_lens_offset) : (0)); - int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); - int active_topk = _min_0; - const int warp_in_compute = warp; - const int tmem_row_origin = warp_in_compute * 32; - const int logical_row_origin = warp_in_compute * 16; - const int my_row = logical_row_origin + lane % 16; - const int col_half = lane / 16; - unsigned int _phase_s_full_0 = 0; - mbarrier_wait(s_full_addr, _phase_s_full_0); - _phase_s_full_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - int score_addr = taddr + (unsigned int)(tmem_row_origin << 16); - float _tmem_load_0[64]; - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_0[0])), "=r"(*reinterpret_cast(&_tmem_load_0[1])), "=r"(*reinterpret_cast(&_tmem_load_0[2])), "=r"(*reinterpret_cast(&_tmem_load_0[3])), "=r"(*reinterpret_cast(&_tmem_load_0[4])), "=r"(*reinterpret_cast(&_tmem_load_0[5])), "=r"(*reinterpret_cast(&_tmem_load_0[6])), "=r"(*reinterpret_cast(&_tmem_load_0[7])), "=r"(*reinterpret_cast(&_tmem_load_0[8])), "=r"(*reinterpret_cast(&_tmem_load_0[9])), "=r"(*reinterpret_cast(&_tmem_load_0[10])), "=r"(*reinterpret_cast(&_tmem_load_0[11])), "=r"(*reinterpret_cast(&_tmem_load_0[12])), "=r"(*reinterpret_cast(&_tmem_load_0[13])), "=r"(*reinterpret_cast(&_tmem_load_0[14])), "=r"(*reinterpret_cast(&_tmem_load_0[15])), "=r"(*reinterpret_cast(&_tmem_load_0[16])), "=r"(*reinterpret_cast(&_tmem_load_0[17])), "=r"(*reinterpret_cast(&_tmem_load_0[18])), "=r"(*reinterpret_cast(&_tmem_load_0[19])), "=r"(*reinterpret_cast(&_tmem_load_0[20])), "=r"(*reinterpret_cast(&_tmem_load_0[21])), "=r"(*reinterpret_cast(&_tmem_load_0[22])), "=r"(*reinterpret_cast(&_tmem_load_0[23])), "=r"(*reinterpret_cast(&_tmem_load_0[24])), "=r"(*reinterpret_cast(&_tmem_load_0[25])), "=r"(*reinterpret_cast(&_tmem_load_0[26])), "=r"(*reinterpret_cast(&_tmem_load_0[27])), "=r"(*reinterpret_cast(&_tmem_load_0[28])), "=r"(*reinterpret_cast(&_tmem_load_0[29])), "=r"(*reinterpret_cast(&_tmem_load_0[30])), "=r"(*reinterpret_cast(&_tmem_load_0[31])) - : "r"(score_addr)); - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_0[32])), "=r"(*reinterpret_cast(&_tmem_load_0[33])), "=r"(*reinterpret_cast(&_tmem_load_0[34])), "=r"(*reinterpret_cast(&_tmem_load_0[35])), "=r"(*reinterpret_cast(&_tmem_load_0[36])), "=r"(*reinterpret_cast(&_tmem_load_0[37])), "=r"(*reinterpret_cast(&_tmem_load_0[38])), "=r"(*reinterpret_cast(&_tmem_load_0[39])), "=r"(*reinterpret_cast(&_tmem_load_0[40])), "=r"(*reinterpret_cast(&_tmem_load_0[41])), "=r"(*reinterpret_cast(&_tmem_load_0[42])), "=r"(*reinterpret_cast(&_tmem_load_0[43])), "=r"(*reinterpret_cast(&_tmem_load_0[44])), "=r"(*reinterpret_cast(&_tmem_load_0[45])), "=r"(*reinterpret_cast(&_tmem_load_0[46])), "=r"(*reinterpret_cast(&_tmem_load_0[47])), "=r"(*reinterpret_cast(&_tmem_load_0[48])), "=r"(*reinterpret_cast(&_tmem_load_0[49])), "=r"(*reinterpret_cast(&_tmem_load_0[50])), "=r"(*reinterpret_cast(&_tmem_load_0[51])), "=r"(*reinterpret_cast(&_tmem_load_0[52])), "=r"(*reinterpret_cast(&_tmem_load_0[53])), "=r"(*reinterpret_cast(&_tmem_load_0[54])), "=r"(*reinterpret_cast(&_tmem_load_0[55])), "=r"(*reinterpret_cast(&_tmem_load_0[56])), "=r"(*reinterpret_cast(&_tmem_load_0[57])), "=r"(*reinterpret_cast(&_tmem_load_0[58])), "=r"(*reinterpret_cast(&_tmem_load_0[59])), "=r"(*reinterpret_cast(&_tmem_load_0[60])), "=r"(*reinterpret_cast(&_tmem_load_0[61])), "=r"(*reinterpret_cast(&_tmem_load_0[62])), "=r"(*reinterpret_cast(&_tmem_load_0[63])) - : "r"(score_addr + 32)); - int valid_cols = active_topk - col_half * 64; - int causal_cols = smem_index_flags[col_half * 2 + 1]; - valid_cols = ((valid_cols < causal_cols) ? valid_cols : causal_cols); - uint32_t _slice_lo_mask_0; - { - int _lim_0 = valid_cols; - if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } - else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); - } - } - if (!(_slice_lo_mask_0 & (1u << 0))) _tmem_load_0[0] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 1))) _tmem_load_0[1] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 2))) _tmem_load_0[2] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 3))) _tmem_load_0[3] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 4))) _tmem_load_0[4] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 5))) _tmem_load_0[5] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 6))) _tmem_load_0[6] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 7))) _tmem_load_0[7] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 8))) _tmem_load_0[8] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 9))) _tmem_load_0[9] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 10))) _tmem_load_0[10] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 11))) _tmem_load_0[11] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 12))) _tmem_load_0[12] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 13))) _tmem_load_0[13] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 14))) _tmem_load_0[14] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 15))) _tmem_load_0[15] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 16))) _tmem_load_0[16] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 17))) _tmem_load_0[17] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 18))) _tmem_load_0[18] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 19))) _tmem_load_0[19] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 20))) _tmem_load_0[20] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 21))) _tmem_load_0[21] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 22))) _tmem_load_0[22] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 23))) _tmem_load_0[23] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 24))) _tmem_load_0[24] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 25))) _tmem_load_0[25] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 26))) _tmem_load_0[26] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 27))) _tmem_load_0[27] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 28))) _tmem_load_0[28] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 29))) _tmem_load_0[29] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 30))) _tmem_load_0[30] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 31))) _tmem_load_0[31] = -CAKE_INF; - uint32_t _slice_lo_mask_1; - { - int _lim_1 = valid_cols - 32; - if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } - else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); - } - } - if (!(_slice_lo_mask_1 & (1u << 0))) _tmem_load_0[32] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 1))) _tmem_load_0[33] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 2))) _tmem_load_0[34] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 3))) _tmem_load_0[35] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 4))) _tmem_load_0[36] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 5))) _tmem_load_0[37] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 6))) _tmem_load_0[38] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 7))) _tmem_load_0[39] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 8))) _tmem_load_0[40] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 9))) _tmem_load_0[41] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 10))) _tmem_load_0[42] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 11))) _tmem_load_0[43] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 12))) _tmem_load_0[44] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 13))) _tmem_load_0[45] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 14))) _tmem_load_0[46] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 15))) _tmem_load_0[47] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 16))) _tmem_load_0[48] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 17))) _tmem_load_0[49] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 18))) _tmem_load_0[50] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 19))) _tmem_load_0[51] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 20))) _tmem_load_0[52] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 21))) _tmem_load_0[53] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 22))) _tmem_load_0[54] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 23))) _tmem_load_0[55] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 24))) _tmem_load_0[56] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 25))) _tmem_load_0[57] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 26))) _tmem_load_0[58] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 27))) _tmem_load_0[59] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 28))) _tmem_load_0[60] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 29))) _tmem_load_0[61] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 30))) _tmem_load_0[62] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 31))) _tmem_load_0[63] = -CAKE_INF; - float2 _reg_reduce_max2_2 = {-CAKE_INF, -CAKE_INF}; - row_max_x32_accum(&_tmem_load_0[0], _reg_reduce_max2_2); - row_max_x32_accum(&_tmem_load_0[32], _reg_reduce_max2_2); - float _tmem_load_0_max = row_max_reduce(_reg_reduce_max2_2); - float row_max = _tmem_load_0_max; - float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, row_max, 16); - float _max_1 = max_noftz(row_max, _shfl_xor_0); - row_max = _max_1; - if (has_sinks != 0 && my_row < num_heads) { - float sink_unscaled = sinks[my_row] * 1.4426950408889634f / softmax_scale_log2; - float _max_2 = max_noftz(row_max, sink_unscaled); - row_max = _max_2; - } - bool empty_row = row_max == -CAKE_INF; - float safe_max = ((empty_row) ? 0.0f : row_max); - float max_scaled = safe_max * softmax_scale_log2; - float2 _f2_0 = make_float2(softmax_scale_log2, softmax_scale_log2); - float2 _f2_1 = make_float2(-max_scaled, -max_scaled); - float2 _f2_2 = make_float2(0.0f, 0.0f); - float2 sum0 = _f2_2; - float2 _f2_3 = make_float2(0.0f, 0.0f); - float2 sum1 = _f2_3; - float2 _f2_4 = make_float2(0.0f, 0.0f); - float2 sum2 = _f2_4; - float2 _f2_5 = make_float2(0.0f, 0.0f); - float2 sum3 = _f2_5; - #pragma unroll - for (int i = 0; i < 64; i += 8) { - float2 _f2_6 = make_float2(_tmem_load_0[i], _tmem_load_0[i + 1]); - float2 _f2_7 = make_float2(_tmem_load_0[i + 2], _tmem_load_0[i + 3]); - float2 _f2_8 = make_float2(_tmem_load_0[i + 4], _tmem_load_0[i + 5]); - float2 _f2_9 = make_float2(_tmem_load_0[i + 6], _tmem_load_0[i + 7]); - float2 affine0 = fma_f32x2_rn_ftz(_f2_6, _f2_0, _f2_1); - float2 affine1 = fma_f32x2_rn_ftz(_f2_7, _f2_0, _f2_1); - float2 affine2 = fma_f32x2_rn_ftz(_f2_8, _f2_0, _f2_1); - float2 affine3 = fma_f32x2_rn_ftz(_f2_9, _f2_0, _f2_1); - float _exp2_0 = approx_exp2(affine0.x); - float exp0 = _exp2_0; - float _exp2_1 = approx_exp2(affine0.y); - float exp1 = _exp2_1; - float _exp2_2 = approx_exp2(affine1.x); - float exp2 = _exp2_2; - float _exp2_3 = approx_exp2(affine1.y); - float exp3 = _exp2_3; - float _exp2_4 = approx_exp2(affine2.x); - float exp4 = _exp2_4; - float _exp2_5 = approx_exp2(affine2.y); - float exp5 = _exp2_5; - float _exp2_6 = approx_exp2(affine3.x); - float exp6 = _exp2_6; - float _exp2_7 = approx_exp2(affine3.y); - float exp7 = _exp2_7; - _tmem_load_0[i] = exp0; - _tmem_load_0[i + 1] = exp1; - _tmem_load_0[i + 2] = exp2; - _tmem_load_0[i + 3] = exp3; - _tmem_load_0[i + 4] = exp4; - _tmem_load_0[i + 5] = exp5; - _tmem_load_0[i + 6] = exp6; - _tmem_load_0[i + 7] = exp7; - float2 _f2_10 = make_float2(exp0, exp1); - sum0 = add_f32x2(sum0, _f2_10); - float2 _f2_11 = make_float2(exp2, exp3); - sum1 = add_f32x2(sum1, _f2_11); - float2 _f2_12 = make_float2(exp4, exp5); - sum2 = add_f32x2(sum2, _f2_12); - float2 _f2_13 = make_float2(exp6, exp7); - sum3 = add_f32x2(sum3, _f2_13); - } - float2 sum01 = add_f32x2(sum0, sum1); - float2 sum23 = add_f32x2(sum2, sum3); - float2 row_sum_pair = add_f32x2(sum01, sum23); - float row_sum = row_sum_pair.x + row_sum_pair.y; - float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, row_sum, 16); - row_sum = row_sum + _shfl_xor_1; - if (has_sinks != 0 && my_row < num_heads) { - float _exp2_8 = approx_exp2(sinks[my_row] * 1.4426950408889634f - max_scaled); - row_sum = row_sum + _exp2_8; - } - unsigned int packed_p[32]; - #pragma unroll - for (int _lp = 0; _lp < 32; _lp++) { - __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_0[_lp*2 + 0], _tmem_load_0[_lp*2+1 + 0])); - packed_p[_lp] = *(uint32_t*)&_bf2; - } - int p_addr = taddr + 128 + (unsigned int)(tmem_row_origin << 16); - asm volatile( - "tcgen05.st.sync.aligned.16x32bx2.x16.b32" - " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" - :: "r"(p_addr), "r"(*reinterpret_cast(&packed_p[0])), "r"(*reinterpret_cast(&packed_p[1])), "r"(*reinterpret_cast(&packed_p[2])), "r"(*reinterpret_cast(&packed_p[3])), "r"(*reinterpret_cast(&packed_p[4])), "r"(*reinterpret_cast(&packed_p[5])), "r"(*reinterpret_cast(&packed_p[6])), "r"(*reinterpret_cast(&packed_p[7])), "r"(*reinterpret_cast(&packed_p[8])), "r"(*reinterpret_cast(&packed_p[9])), "r"(*reinterpret_cast(&packed_p[10])), "r"(*reinterpret_cast(&packed_p[11])), "r"(*reinterpret_cast(&packed_p[12])), "r"(*reinterpret_cast(&packed_p[13])), "r"(*reinterpret_cast(&packed_p[14])), "r"(*reinterpret_cast(&packed_p[15]))); - asm volatile( - "tcgen05.st.sync.aligned.16x32bx2.x16.b32" - " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" - :: "r"(p_addr + 16), "r"(*reinterpret_cast(&(packed_p + 16)[0])), "r"(*reinterpret_cast(&(packed_p + 16)[1])), "r"(*reinterpret_cast(&(packed_p + 16)[2])), "r"(*reinterpret_cast(&(packed_p + 16)[3])), "r"(*reinterpret_cast(&(packed_p + 16)[4])), "r"(*reinterpret_cast(&(packed_p + 16)[5])), "r"(*reinterpret_cast(&(packed_p + 16)[6])), "r"(*reinterpret_cast(&(packed_p + 16)[7])), "r"(*reinterpret_cast(&(packed_p + 16)[8])), "r"(*reinterpret_cast(&(packed_p + 16)[9])), "r"(*reinterpret_cast(&(packed_p + 16)[10])), "r"(*reinterpret_cast(&(packed_p + 16)[11])), "r"(*reinterpret_cast(&(packed_p + 16)[12])), "r"(*reinterpret_cast(&(packed_p + 16)[13])), "r"(*reinterpret_cast(&(packed_p + 16)[14])), "r"(*reinterpret_cast(&(packed_p + 16)[15]))); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - mbarrier_arrive(p_full_addr); - if (col_half == 0) { - float _max_3 = max_noftz(row_sum, 1.1754943508222875e-38f); - float published_sum = _max_3; - smem_sum[my_row] = published_sum; - } - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(sum_ready_addr); - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(tmem_dealloc_addr); - } - } - // ---- Role: epilogue ---- - if (warp >= 4 && warp <= 7) { - asm volatile("setmaxnreg.dec.sync.aligned.u32 128;"); - { // epilogue_main - float output_scale = bmm2_scale[0]; - int work_idx = blockIdx.x; - int query_idx_1 = work_idx >> 2; - int v_chunk = work_idx & 3; - const int warp_in_role = warp - 4; - const int tmem_row_origin_1 = warp_in_role * 32; - const int logical_row_origin_1 = warp_in_role * 16; - const int my_row_1 = logical_row_origin_1 + lane % 16; - const int col_half_1 = lane / 16; - const int row_addr = tmem_row_origin_1 << 16; - unsigned int _phase_o_done_0 = 0; - mbarrier_wait(o_done_addr, _phase_o_done_0); - _phase_o_done_0 ^= 1; - unsigned int _phase_sum_ready_0 = 0; - mbarrier_wait(sum_ready_addr, _phase_sum_ready_0); - _phase_sum_ready_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - float row_sum_1 = smem_sum[my_row_1]; - float _rcp_0 = approx_rcp(row_sum_1); - float inv_sum = _rcp_0; - int output_base = (query_idx_1 * num_heads + my_row_1) * 512 + v_chunk * 128; - float _tmem_load_1[64]; - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_1[0])), "=r"(*reinterpret_cast(&_tmem_load_1[1])), "=r"(*reinterpret_cast(&_tmem_load_1[2])), "=r"(*reinterpret_cast(&_tmem_load_1[3])), "=r"(*reinterpret_cast(&_tmem_load_1[4])), "=r"(*reinterpret_cast(&_tmem_load_1[5])), "=r"(*reinterpret_cast(&_tmem_load_1[6])), "=r"(*reinterpret_cast(&_tmem_load_1[7])), "=r"(*reinterpret_cast(&_tmem_load_1[8])), "=r"(*reinterpret_cast(&_tmem_load_1[9])), "=r"(*reinterpret_cast(&_tmem_load_1[10])), "=r"(*reinterpret_cast(&_tmem_load_1[11])), "=r"(*reinterpret_cast(&_tmem_load_1[12])), "=r"(*reinterpret_cast(&_tmem_load_1[13])), "=r"(*reinterpret_cast(&_tmem_load_1[14])), "=r"(*reinterpret_cast(&_tmem_load_1[15])), "=r"(*reinterpret_cast(&_tmem_load_1[16])), "=r"(*reinterpret_cast(&_tmem_load_1[17])), "=r"(*reinterpret_cast(&_tmem_load_1[18])), "=r"(*reinterpret_cast(&_tmem_load_1[19])), "=r"(*reinterpret_cast(&_tmem_load_1[20])), "=r"(*reinterpret_cast(&_tmem_load_1[21])), "=r"(*reinterpret_cast(&_tmem_load_1[22])), "=r"(*reinterpret_cast(&_tmem_load_1[23])), "=r"(*reinterpret_cast(&_tmem_load_1[24])), "=r"(*reinterpret_cast(&_tmem_load_1[25])), "=r"(*reinterpret_cast(&_tmem_load_1[26])), "=r"(*reinterpret_cast(&_tmem_load_1[27])), "=r"(*reinterpret_cast(&_tmem_load_1[28])), "=r"(*reinterpret_cast(&_tmem_load_1[29])), "=r"(*reinterpret_cast(&_tmem_load_1[30])), "=r"(*reinterpret_cast(&_tmem_load_1[31])) - : "r"(taddr + 256 + (unsigned int)row_addr)); - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_1[32])), "=r"(*reinterpret_cast(&_tmem_load_1[33])), "=r"(*reinterpret_cast(&_tmem_load_1[34])), "=r"(*reinterpret_cast(&_tmem_load_1[35])), "=r"(*reinterpret_cast(&_tmem_load_1[36])), "=r"(*reinterpret_cast(&_tmem_load_1[37])), "=r"(*reinterpret_cast(&_tmem_load_1[38])), "=r"(*reinterpret_cast(&_tmem_load_1[39])), "=r"(*reinterpret_cast(&_tmem_load_1[40])), "=r"(*reinterpret_cast(&_tmem_load_1[41])), "=r"(*reinterpret_cast(&_tmem_load_1[42])), "=r"(*reinterpret_cast(&_tmem_load_1[43])), "=r"(*reinterpret_cast(&_tmem_load_1[44])), "=r"(*reinterpret_cast(&_tmem_load_1[45])), "=r"(*reinterpret_cast(&_tmem_load_1[46])), "=r"(*reinterpret_cast(&_tmem_load_1[47])), "=r"(*reinterpret_cast(&_tmem_load_1[48])), "=r"(*reinterpret_cast(&_tmem_load_1[49])), "=r"(*reinterpret_cast(&_tmem_load_1[50])), "=r"(*reinterpret_cast(&_tmem_load_1[51])), "=r"(*reinterpret_cast(&_tmem_load_1[52])), "=r"(*reinterpret_cast(&_tmem_load_1[53])), "=r"(*reinterpret_cast(&_tmem_load_1[54])), "=r"(*reinterpret_cast(&_tmem_load_1[55])), "=r"(*reinterpret_cast(&_tmem_load_1[56])), "=r"(*reinterpret_cast(&_tmem_load_1[57])), "=r"(*reinterpret_cast(&_tmem_load_1[58])), "=r"(*reinterpret_cast(&_tmem_load_1[59])), "=r"(*reinterpret_cast(&_tmem_load_1[60])), "=r"(*reinterpret_cast(&_tmem_load_1[61])), "=r"(*reinterpret_cast(&_tmem_load_1[62])), "=r"(*reinterpret_cast(&_tmem_load_1[63])) - : "r"(taddr + 256 + (unsigned int)row_addr + 32)); - int col_base = col_half_1 * 64; - if (my_row_1 < num_heads) { - #pragma unroll - for (int offset = 0; offset < 64; offset += 8) { - { - const float2 _prescale2_0 = {inv_sum * output_scale, inv_sum * output_scale}; - #if __CUDA_ARCH__ >= 1000 - #pragma unroll - for (int _ps = 0; _ps < 4; _ps++) - mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_1[offset])[_ps], _prescale2_0); - #else - #pragma unroll - for (int _ps = 0; _ps < 8; _ps++) - _tmem_load_1[offset + _ps] *= inv_sum * output_scale; - #endif - __nv_bfloat162 _pk[4]; - _pk[0] = __floats2bfloat162_rn(_tmem_load_1[offset + 0], _tmem_load_1[offset + 1]); - _pk[1] = __floats2bfloat162_rn(_tmem_load_1[offset + 2], _tmem_load_1[offset + 3]); - _pk[2] = __floats2bfloat162_rn(_tmem_load_1[offset + 4], _tmem_load_1[offset + 5]); - _pk[3] = __floats2bfloat162_rn(_tmem_load_1[offset + 6], _tmem_load_1[offset + 7]); - *reinterpret_cast(&((__nv_bfloat16*)(O + (output_base + col_base + offset)))[0]) = *reinterpret_cast(&_pk[0]); - } - } - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(tmem_dealloc_addr); - } - } - // ---- Role: mma_warp ---- - if (warp == 8) { - { // mma_warp_main - unsigned int _phase_q_full_0 = 0; - mbarrier_wait(q_full_addr, _phase_q_full_0); - _phase_q_full_0 ^= 1; - unsigned int _phase_kv_full = 0; - #pragma unroll - for (int k_stage = 0; k_stage < 4; k_stage++) { - mbarrier_wait(kv_full_addr + (k_stage) * 8, _phase_kv_full); - asm volatile("tcgen05.fence::after_thread_sync;"); - int _mma_a_lo_0 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (k_stage * 2) * 512); - int _mma_b_lo_0 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (k_stage) * 2048); - { - uint64_t _mma_ss_a_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_0); - uint64_t _mma_ss_b_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_0); - if (elect_sync()) { - tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, ((k_stage == 0) ? 0 : 1)); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); - } - } - int _mma_a_lo_1 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (k_stage * 2 + 1) * 512); - int _mma_b_lo_1 = make_warp_uniform((((smem_kv_addr + 16384) >> 4) & 0x3FFF) + (k_stage) * 2048); - { - uint64_t _mma_ss_a_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_1); - uint64_t _mma_ss_b_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_1); - if (elect_sync()) { - tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); - } - } - if (k_stage == 3) { - elect_commit(s_full_addr); - } - elect_commit(kv_empty_addr + (k_stage) * 8); - } - int _mma_prepared_b_lo_2 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (0) * 2048); - unsigned int _phase_p_full_0 = 0; - mbarrier_wait(p_full_addr, _phase_p_full_0); - _phase_p_full_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - mbarrier_wait(kv_full_addr, 1); - asm volatile("tcgen05.fence::after_thread_sync;"); - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 dhi, blo, ta, id;\n\t" - ".reg .b64 db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "" - "mov.b32 dhi, 0x40004040;\n\t" - "mov.b32 id, 69272720;\n\t" - "mov.b32 ta, %2;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p0;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "}\n" - :: "r"((tmem_tmem + (256))), "r"(_mma_prepared_b_lo_2), "r"(tmem_tmem + 128), "r"(0)); - elect_commit(o_done_addr); - elect_commit(kv_empty_addr); - unsigned int _phase_tmem_dealloc_0 = 0; - mbarrier_wait(tmem_dealloc_addr, _phase_tmem_dealloc_0); - _phase_tmem_dealloc_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); - asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); - } - } - // ---- Role: empty ---- - if (warp >= 9 && warp <= 11) { - // idle — no tasks assigned - } - // ---- Role: load_warp ---- - if (warp >= 12 && warp <= 15) { - asm volatile("setmaxnreg.dec.sync.aligned.u32 64;"); - { // load_warp_main - const int load_warp_rank = warp - 12; - int work_idx_1 = blockIdx.x; - int query_idx_2 = work_idx_1 >> 2; - int v_chunk_1 = work_idx_1 & 3; - int* row_ptr = swa_indices + (query_idx_2 * swa_index_stride); - int index_offset = load_warp_rank * 32 + lane; - int sparse_row = -1; - if (index_offset < sparse_topk) { - sparse_row = row_ptr[index_offset]; - } - if (load_warp_rank == 0) { - if (elect_sync()) { - mbarrier_arrive_expect_tx(q_full_addr, 65536); - #pragma unroll - for (int q_stage = 0; q_stage < 8; q_stage++) { - tma_4d_gmem2smem(smem_q_addr + (unsigned int)(q_stage * 8192), (&tmap_q), 0, 0, q_stage, query_idx_2, q_full_addr); - } - } - } - smem_indices[index_offset] = sparse_row; - unsigned int _vote_0 = __ballot_sync(0xFFFFFFFF, sparse_row < 0); - unsigned int invalid_indices = _vote_0; - int valid_prefix = 32; - if (invalid_indices != 0) { - int _ffs_0 = __ffs(invalid_indices); - valid_prefix = _ffs_0 - 1; - } - if (elect_sync()) { - smem_index_flags[load_warp_rank] = valid_prefix; - } - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - asm volatile("barrier.sync 8, 128;" ::: "memory"); - if ((load_warp_rank & 1) != 0) { - if (elect_sync()) { - int lower_prefix = smem_index_flags[load_warp_rank - 1]; - int upper_prefix = smem_index_flags[load_warp_rank]; - smem_index_flags[load_warp_rank] = ((lower_prefix < 32) ? lower_prefix : 32 + upper_prefix); - } - } - unsigned int _phase_kv_empty = 1; - #pragma unroll - for (int k_stage_1 = 0; k_stage_1 < 4; k_stage_1++) { - mbarrier_wait(kv_empty_addr + (k_stage_1) * 8, _phase_kv_empty); - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (k_stage_1) * 8, 8192); - } - #pragma unroll - for (int local_group = 0; local_group < 8; local_group++) { - int dst_k = smem_kv_addr + (unsigned int)(k_stage_1 * 32768); - int group = load_warp_rank * 8 + local_group; - int group_offset = group * 4; - int raw_rows[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows[0])), "=r"(*reinterpret_cast(&raw_rows[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows[(0) + 3])) - : "r"(smem_indices_addr + (unsigned int)(group_offset * 4))); - int raw0 = raw_rows[0]; - int raw1 = raw_rows[1]; - int raw2 = raw_rows[2]; - int raw3 = raw_rows[3]; - int row0 = ((raw0 >= 0) ? raw0 : 0); - int row1 = ((raw1 >= 0) ? raw1 : 0); - int row2 = ((raw2 >= 0) ? raw2 : 0); - int row3 = ((raw3 >= 0) ? raw3 : 0); - if (elect_sync()) { - tma_gather4_gmem2smem(dst_k + group * 512, (&tmap_swa_kv), k_stage_1 * 128, row0, row1, row2, row3, kv_full_addr + (k_stage_1) * 8); - tma_gather4_gmem2smem(dst_k + 16384 + group * 512, (&tmap_swa_kv), k_stage_1 * 128 + 64, row0, row1, row2, row3, kv_full_addr + (k_stage_1) * 8); - } - } - } - #pragma unroll - for (int v_stage = 0; v_stage < 1; v_stage++) { - mbarrier_wait(kv_empty_addr + (v_stage) * 8, 0); - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (v_stage) * 8, 8192); - } - #pragma unroll - for (int local_group_1 = 0; local_group_1 < 8; local_group_1++) { - int dst_v = smem_v_addr + (unsigned int)(v_stage * 32768); - int v_col = v_chunk_1 * 128 + v_stage * 64; - int group_1 = load_warp_rank * 8 + local_group_1; - int group_offset_1 = group_1 * 4; - int raw_rows_1[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_1[0])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 3])) - : "r"(smem_indices_addr + (unsigned int)(group_offset_1 * 4))); - int raw0_1 = raw_rows_1[0]; - int raw1_1 = raw_rows_1[1]; - int raw2_1 = raw_rows_1[2]; - int raw3_1 = raw_rows_1[3]; - int row0_1 = ((raw0_1 >= 0) ? raw0_1 : 0); - int row1_1 = ((raw1_1 >= 0) ? raw1_1 : 0); - int row2_1 = ((raw2_1 >= 0) ? raw2_1 : 0); - int row3_1 = ((raw3_1 >= 0) ? raw3_1 : 0); - if (elect_sync()) { - tma_gather4_gmem2smem(dst_v + group_1 * 512, (&tmap_swa_kv), v_col, row0_1, row1_1, row2_1, row3_1, kv_full_addr + (v_stage) * 8); - tma_gather4_gmem2smem(dst_v + 16384 + group_1 * 512, (&tmap_swa_kv), v_col + 64, row0_1, row1_1, row2_1, row3_1, kv_full_addr + (v_stage) * 8); - } - } - } - } - } - - // Cleanup -} - -} // extern "C" diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_9e1188b192c9c806d7c6_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_9e1188b192c9c806d7c6_binding.cu deleted file mode 100644 index b3d17f4a22c..00000000000 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_9e1188b192c9c806d7c6_binding.cu +++ /dev/null @@ -1,461 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -// Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_9e1188b192c9c806d7c6'. -#include -#include -#include -#include - -#include "tvm_ffi_utils.h" - -#include -#include -#include -#include -#include -#include -#include - -extern "C" __global__ void kernel_cake_dsv4_9e1188b192c9c806d7c6(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items); - - -namespace cake_host_shim_4aa7ddc0eb16149e { - -using tvm::ffi::Optional; -using tvm::ffi::TensorView; - -class ScopedCudaDevice { - public: - explicit ScopedCudaDevice(int device_id) { - cudaError_t error = cudaGetDevice(&previous_device_); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed before host-shim launch: cudaError=" - << static_cast(error); - if (previous_device_ != device_id) { - error = cudaSetDevice(device_id); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaSetDevice failed before host-shim launch for cuda:" - << device_id << ": cudaError=" << static_cast(error); - restore_ = true; - } - } - - ScopedCudaDevice(const ScopedCudaDevice&) = delete; - ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; - - ~ScopedCudaDevice() noexcept { - if (restore_) { - (void)cudaSetDevice(previous_device_); - } - } - - private: - int previous_device_ = -1; - bool restore_ = false; -}; - -inline int64_t CakeDeviceMultiprocessorCount(int device_id) { - constexpr int kMaxCachedCudaDevices = 64; - TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) - << "physical-SM-count cache does not cover cuda:" << device_id; - static std::atomic count_by_device[kMaxCachedCudaDevices]{}; - int cached = count_by_device[device_id].load(std::memory_order_acquire); - if (cached > 0) return cached; - - int count = 0; - cudaError_t error = cudaDeviceGetAttribute( - &count, cudaDevAttrMultiProcessorCount, device_id); - TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) - << "querying multiProcessorCount failed for cuda:" << device_id - << ": cudaError=" << static_cast(error) << ", count=" << count; - // Concurrent first launches may repeat the immutable device query, but all - // publication is atomic and every later hot-path lookup is one acquire load. - count_by_device[device_id].store(count, std::memory_order_release); - return count; -} - -inline void CheckCudaTensor(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) - << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; -} - -inline void CheckSameCudaDevice( - const TensorView& t, - const TensorView& reference, - const char* name, - const char* reference_name) { - TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) - << name << " must be on the same CUDA device as " << reference_name - << ": got cuda:" << t.device().device_id - << " versus cuda:" << reference.device().device_id; -} - -inline void CheckCurrentCudaDevice( - const TensorView& reference, - const char* reference_name) { - int current_device = -1; - cudaError_t error = cudaGetDevice(¤t_device); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed while validating " << reference_name - << ": cudaError=" << static_cast(error); - TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) - << "current CUDA device must match " << reference_name - << ": current=cuda:" << current_device - << ", tensor=cuda:" << reference.device().device_id; -} - -inline void CheckContiguous(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; -} - -inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { - DLDataType d = t.dtype(); - TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) - << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits - << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits - << ", lanes=" << (int)d.lanes << ")"; -} - -// A logical axis.outer(trailing) folds every source dim above the trailing -// dimensions. Shape products are independent of physical strides, so verify -// the leading dimensions form one dense row-major chain instead of inventing -// a "folded stride". The descriptor reads its exact adjacent physical step -// separately through stride[-(trailing + 1)]. -inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { - TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) - << name << " cannot fold leading dimensions above " << trailing - << " trailing dims from ndim=" << t.ndim(); - int outer_last = t.ndim() - trailing - 1; - if (outer_last <= 0) { - return; - } - int64_t step = t.stride(outer_last); - TVM_FFI_CHECK(step > 0, ValueError) - << name << " physical strides must be positive"; - int64_t expected = step; - for (int axis = outer_last - 1; axis >= 0; --axis) { - expected *= t.size(axis + 1); - if (t.size(axis) > 1) { - TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) - << name << " leading dims are not physically foldable above " << trailing - << " trailing dims: stride(" << axis << ")=" << t.stride(axis) - << ", expected " << expected; - } - } -} - -// 4D TMA descriptor for buffer 'tmap_q' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 3, ValueError) - << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - int64_t d2 = t.size(t.ndim() - 2); - int64_t d3 = t.size(t.ndim() - 3); - TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) - << "TMA source 'tmap_q' trailing dims must be positive"; - TVM_FFI_CHECK(d1 % 128 == 0, ValueError) - << "TMA source 'tmap_q' extent " << d1 - << " must divide exactly by " << 128; - uint64_t global_dim[4] = {(uint64_t)(128), (uint64_t)(d2), (uint64_t)((d1 / 128)), (uint64_t)(d3)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; - TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) - << "TMA box (128, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; - int64_t carrier_stride_0 = d1; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; - int64_t carrier_stride_1 = 128; - TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; - TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; - TVM_FFI_CHECK((carrier_stride_1 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; - int64_t carrier_stride_2 = (d2 * d1); - TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; - TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; - TVM_FFI_CHECK((carrier_stride_2 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; - uint64_t global_strides[3] = { - (uint64_t)((carrier_stride_0 * 8) / 8), - (uint64_t)((carrier_stride_1 * 8) / 8), - (uint64_t)((carrier_stride_2 * 8) / 8), - }; - uint32_t box_dim[4] = {128u, 64u, 1u, 1u}; - uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_swa_kv' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_swa_kv' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; - TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (128, 1) exceeds resolved global dims for 'tmap_swa_kv'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 8) / 8), - }; - uint32_t box_dim[2] = {128u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_compressed_kv' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_compressed_kv(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_compressed_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_compressed_kv' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_compressed_kv' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_compressed_kv"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_compressed_kv' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved a non-positive global dim"; - TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (128, 1) exceeds resolved global dims for 'tmap_compressed_kv'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 8) / 8), - }; - uint32_t box_dim[2] = {128u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_kv') failed: CUresult=" << (int)r; - return tm; -} - -void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_compressed_kv, TensorView arg_O, TensorView arg_partial_lse, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_num_query_tokens, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, int64_t grid_x, int64_t grid_y, int64_t grid_z) { - DLDevice dev = arg_tmap_q.device(); - ScopedCudaDevice device_guard(dev.device_id); - CheckCudaTensor(arg_tmap_q, "tmap_q"); - CheckDtype(arg_tmap_q, "tmap_q", 1, 8, 1); - CheckContiguous(arg_tmap_q, "tmap_q"); - CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); - CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 1, 8, 1); - CheckCudaTensor(arg_tmap_compressed_kv, "tmap_compressed_kv"); - CheckDtype(arg_tmap_compressed_kv, "tmap_compressed_kv", 1, 8, 1); - CheckCudaTensor(arg_O, "O"); - CheckDtype(arg_O, "O", 4, 16, 1); - CheckContiguous(arg_O, "O"); - CheckCudaTensor(arg_partial_lse, "partial_lse"); - CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); - CheckContiguous(arg_partial_lse, "partial_lse"); - CheckCudaTensor(arg_swa_indices, "swa_indices"); - CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); - CheckContiguous(arg_swa_indices, "swa_indices"); - CheckCudaTensor(arg_compressed_indices, "compressed_indices"); - CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); - CheckContiguous(arg_compressed_indices, "compressed_indices"); - CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); - CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckCudaTensor(arg_sinks, "sinks"); - CheckDtype(arg_sinks, "sinks", 2, 32, 1); - CheckContiguous(arg_sinks, "sinks"); - CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); - CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); - CheckContiguous(arg_bmm1_scale, "bmm1_scale"); - CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); - CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); - CheckContiguous(arg_bmm2_scale, "bmm2_scale"); - TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) - << "scalar 'num_heads' value " << arg_num_heads - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) - << "scalar 'swa_index_stride' value " << arg_swa_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) - << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) - << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_num_query_tokens >= -2147483648LL && arg_num_query_tokens <= 2147483647LL, ValueError) - << "scalar 'num_query_tokens' value " << arg_num_query_tokens - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) - << "scalar 'sparse_topk' value " << arg_sparse_topk - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) - << "scalar 'has_sinks' value " << arg_has_sinks - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) - << "scalar 'total_work_items' value " << arg_total_work_items - << " is outside i32 range [-2147483648, 2147483647]"; - CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); - CheckSameCudaDevice(arg_tmap_compressed_kv, arg_tmap_q, "tmap_compressed_kv", "tmap_q"); - CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); - CheckSameCudaDevice(arg_partial_lse, arg_tmap_q, "partial_lse", "tmap_q"); - CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); - CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); - CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); - CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); - CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); - CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); - TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) - << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y - << ", " << grid_z << ")"; - TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) - << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z - << ") must be divisible by cluster dims (2, 1, 1)"; - - cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); - CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); - CUtensorMap p_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); - CUtensorMap p_tmap_compressed_kv = EncodeTma_tmap_compressed_kv(arg_tmap_compressed_kv); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); - float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); - int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); - int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); - int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); - float* p_sinks = static_cast(arg_sinks.data_ptr()); - float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); - float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); - int32_t v_num_heads = (int32_t)arg_num_heads; - int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; - int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; - int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; - int32_t v_num_query_tokens = (int32_t)arg_num_query_tokens; - int32_t v_sparse_topk = (int32_t)arg_sparse_topk; - int32_t v_has_sinks = (int32_t)arg_has_sinks; - int32_t v_total_work_items = (int32_t)arg_total_work_items; - void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_tmap_compressed_kv, &p_O, &p_partial_lse, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_num_query_tokens, &v_sparse_topk, &v_has_sinks, &v_total_work_items}; - - static std::mutex smem_mu; - static auto* smem_status_by_device = - new std::unordered_map(); - cudaError_t smem_status = cudaSuccess; - { - std::lock_guard lock(smem_mu); - auto it = smem_status_by_device->find(dev.device_id); - if (it == smem_status_by_device->end()) { - int current_device = -1; - smem_status = cudaGetDevice(¤t_device); - if (smem_status == cudaSuccess && current_device != dev.device_id) { - smem_status = cudaErrorInvalidDevice; - } - if (smem_status == cudaSuccess) { - smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_9e1188b192c9c806d7c6), - cudaFuncAttributeMaxDynamicSharedMemorySize, - 223744); - } - smem_status_by_device->emplace(dev.device_id, smem_status); - } else { - smem_status = it->second; - } - } - TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_9e1188b192c9c806d7c6 failed: " - << cudaGetErrorString(smem_status); - dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); - dim3 block(512u, 1u, 1u); - - // Extended direct-source launch for a linked CUDA kernel symbol. - cudaLaunchAttribute attrs[2]{}; - int n = 0; - attrs[n].id = cudaLaunchAttributeClusterDimension; - attrs[n].val.clusterDim.x = 2u; - attrs[n].val.clusterDim.y = 1u; - attrs[n].val.clusterDim.z = 1u; - ++n; - attrs[n].id = cudaLaunchAttributeClusterSchedulingPolicyPreference; - attrs[n].val.clusterSchedulingPolicyPreference = cudaClusterSchedulingPolicySpread; - ++n; - cudaLaunchConfig_t config{}; - config.gridDim = grid; - config.blockDim = block; - config.dynamicSmemBytes = 223744u; - config.stream = stream; - config.attrs = attrs; - config.numAttrs = n; - cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_9e1188b192c9c806d7c6), kargs); - TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_9e1188b192c9c806d7c6 failed: " - << cudaGetErrorString(launch_status); -} - -} // namespace cake_host_shim_4aa7ddc0eb16149e - -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_4aa7ddc0eb16149e::Run); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_9e1188b192c9c806d7c6_kernel.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_9e1188b192c9c806d7c6_kernel.cu deleted file mode 100644 index b26d3631ccb..00000000000 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_9e1188b192c9c806d7c6_kernel.cu +++ /dev/null @@ -1,2802 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -typedef signed char int8_t; -typedef unsigned char uint8_t; -typedef unsigned short uint16_t; -typedef unsigned int uint32_t; -#if defined(__CUDACC_RTC__) -typedef unsigned long long uint64_t; -#else -typedef unsigned long uint64_t; -#endif -static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); -typedef signed int int32_t; -typedef short int int16_t; -struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; -struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; -static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); -static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); - -#if defined(__CUDACC_RTC__) -typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; -#else -#include -#endif - -static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); -static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); -#include -#include - -#define CAKE_INF CUDART_INF_F -#define TMEM_NCOLS 512 -#define TMEM_TMEM_SCRATCH_OFFSET 0 -#define NUM_K_PIPE_STAGES 3 -#define NUM_V_PIPE_STAGES 2 -#define NUM_KV_PIPE_STAGES 1 -#define NUM_INDEX_PIPE_STAGES 6 -#define SMEM_SMEM_Q_OFF 1024 -#define SMEM_SMEM_Q_STAGE_BYTES 8192 -#define SMEM_SMEM_Q_STRIDE 8192 -#define SMEM_SMEM_KV_OFF 33792 -#define SMEM_SMEM_KV_STAGE_BYTES 16384 -#define SMEM_SMEM_KV_STRIDE 16384 -#define SMEM_SMEM_V_OFF 33792 -#define SMEM_SMEM_V_STAGE_BYTES 16384 -#define SMEM_SMEM_V_STRIDE 16384 -#define SMEM_SMEM_Q_FULL_OFF 1024 -#define SMEM_SMEM_Q_FULL_STAGE_BYTES 32768 -#define SMEM_SMEM_Q_FULL_STRIDE 32768 -#define SMEM_SMEM_K_FULL_OFF 33792 -#define SMEM_SMEM_K_FULL_STAGE_BYTES 32768 -#define SMEM_SMEM_K_FULL_STRIDE 32768 -#define SMEM_SMEM_V_FULL_OFF 132096 -#define SMEM_SMEM_V_FULL_STAGE_BYTES 32768 -#define SMEM_SMEM_V_FULL_STRIDE 32768 -#define SMEM_SMEM_STATS_MAX_OFF 197632 -#define SMEM_SMEM_STATS_MAX_STAGE_BYTES 1024 -#define SMEM_SMEM_STATS_MAX_STRIDE 1024 -#define SMEM_SMEM_STATS_SUM_OFF 198656 -#define SMEM_SMEM_STATS_SUM_STAGE_BYTES 512 -#define SMEM_SMEM_STATS_SUM_STRIDE 512 -#define SMEM_SMEM_STATS_FINAL_MAX_OFF 199168 -#define SMEM_SMEM_STATS_FINAL_MAX_STAGE_BYTES 512 -#define SMEM_SMEM_STATS_FINAL_MAX_STRIDE 512 -#define SMEM_SMEM_SOFTMAX_WARP_PAIR_EXCHANGE_OFF 199680 -#define SMEM_SMEM_SOFTMAX_WARP_PAIR_EXCHANGE_STAGE_BYTES 1024 -#define SMEM_SMEM_SOFTMAX_WARP_PAIR_EXCHANGE_STRIDE 1024 -#define SMEM_SMEM_CORR_WARP_PAIR_EXCHANGE_OFF 200704 -#define SMEM_SMEM_CORR_WARP_PAIR_EXCHANGE_STAGE_BYTES 512 -#define SMEM_SMEM_CORR_WARP_PAIR_EXCHANGE_STRIDE 512 -#define SMEM_SMEM_SPARSE_INDICES_OFF 201216 -#define SMEM_SMEM_SPARSE_INDICES_STAGE_BYTES 1024 -#define SMEM_SMEM_SPARSE_INDICES_STRIDE 1024 -#define SMEM_SMEM_P_FP8_OFF 207360 -#define SMEM_SMEM_P_FP8_STAGE_BYTES 8192 -#define SMEM_SMEM_P_FP8_STRIDE 8192 -#define SMEM_TOTAL 223744 -#define THREADS 512 - -#include - -__device__ __forceinline__ uint32_t elect_sync() { - uint32_t pred = 0; - asm volatile( - "{\n\t" - ".reg .pred %%px;\n\t" - "elect.sync _|%%px, %1;\n\t" - "@%%px mov.s32 %0, 1;\n\t" - "}\n" - : "+r"(pred) - : "r"(0xFFFFFFFF)); - return pred; -} - - -__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { - asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" - :: "r"(mbar_addr), "r"(count) : "memory"); -} - -__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { - asm volatile("mbarrier.init.b64 [%0], %1;" - :: "l"(mbar_addr), "r"(count) : "memory"); -} - - -__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - - -// CTA-local pipelines have short, resident producer/consumer edges. Omitting -// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the -// explicit loop still makes this helper blocking until acquire succeeds. -__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE;\n\t" - "bra.uni LAB_WAIT;\n\t" - "DONE:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Source-faithful relaxed CTA wait used only by a typed protocol that does -// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. -__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, 10000000;\n\t" - "@P1 bra.uni DONE_RELAXED;\n\t" - "bra.uni LAB_WAIT_RELAXED;\n\t" - "DONE_RELAXED:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Exact source ports may request the PTX suspendTimeHint operand explicitly. -// The hint is expressed in nanoseconds and is kept separate from the canonical -// no-hint CTA helper so unrelated schedules retain their existing retry path. -__device__ __forceinline__ void mbarrier_wait_suspend( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_SUSPEND:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_SUSPEND;\n\t" - "bra.uni LAB_WAIT_SUSPEND;\n\t" - "DONE_SUSPEND:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE_CLUSTER;\n\t" - "bra.uni LAB_WAIT_CLUSTER;\n\t" - "DONE_CLUSTER:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - ".reg .u32 WAIT_ADDR;\n\t" - "mov.u32 WAIT_ADDR, %0;\n\t" - "LAB_WAIT_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [WAIT_ADDR], %1, %2;\n\t" - "@P1 bra.uni DONE_HINT;\n\t" - "bra.uni LAB_WAIT_HINT;\n\t" - "DONE_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -// Exact unqualified CTA wait used by source schedules whose PTX intentionally -// omits the acquire qualifier while retaining a typed suspendTimeHint operand. -__device__ __forceinline__ void mbarrier_wait_relaxed_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED_HINT:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra DONE_RELAXED_HINT;\n\t" - "bra LAB_WAIT_RELAXED_HINT;\n\t" - "DONE_RELAXED_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); -} - -__device__ __forceinline__ void mbarrier_wait_cluster_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" - "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" - "DONE_CLUSTER_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_suspend( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait_cluster(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); - } -} - - -__device__ __forceinline__ void tcgen05_mma_f8f6f4_cta2( - int taddr, uint64_t a_desc, uint64_t b_desc, - uint32_t i_desc, int enable_input_d) { - asm volatile( - "{\n\t" - ".reg .pred p;\n\t" - ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - "setp.ne.b32 p, %4, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" - "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "tcgen05.mma.cta_group::2.kind::f8f6f4 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" - "}\n" - :: "r"(taddr), "l"(a_desc), "l"(b_desc), - "r"(i_desc), "r"(enable_input_d) - : "memory"); -} - - -__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { - return (x & 0x3FFFFULL) >> 4ULL; -} - - -union MmaSmemDesc { - uint64_t u64; - uint32_t u32[2]; -}; - -__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { - MmaSmemDesc tmp; - tmp.u64 = smem_desc; - tmp.u32[0] += offset; - smem_desc = tmp.u64; -} - - -__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { - asm volatile( - "{\n\t" - ".reg .pred leader;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" - ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" - "}\n" - :: "r"(mbar_addr), "h"(cta_mask) : "memory"); -} - - -__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { - asm volatile( - "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" - :: "r"(mbar_addr), "r"(bytes) : "memory"); -} - - -__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7," - " %8, %9, %10, %11, %12, %13, %14, %15," - " %16, %17, %18, %19, %20, %21, %22, %23," - " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), - "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), - "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), - "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), - "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), - "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), - "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x16.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7," - " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), - "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), - "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ void tmem_ld_x4(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x4.b32" - " {%0, %1, %2, %3}, [%4];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { - asm volatile( - "tcgen05.st.sync.aligned.32x32b.x16.b32" - " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," - " %9, %10, %11, %12, %13, %14, %15, %16};" - :: "r"(tmem_addr), - "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), - "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), - "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), - "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); -} - - -__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { - asm volatile( - "tcgen05.st.sync.aligned.32x32b.x4.b32" - " [%0], {%1, %2, %3, %4};" - :: "r"(tmem_addr), - "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); -} - - -__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { - uint32_t addr; - asm("{\n\t" - ".reg .u64 u64addr;\n\t" - "cvta.to.shared.u64 u64addr, %1;\n\t" - "cvt.u32.u64 %0, u64addr;\n\t" - "}\n" : "=r"(addr) : "l"(ptr)); - return addr; -} - - -__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { - uint32_t remote; - asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" - : "=r"(remote) : "r"(local_addr), "r"(rank)); - return remote; -} - - -__device__ __forceinline__ float approx_exp2(float x) { - float y; - asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float approx_rcp(float x) { - float y; - asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float max_noftz(float a, float b) { - float c; - asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); - return c; -} - - -__device__ __forceinline__ float warp_reduce_max(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); - return val; -} - - -__device__ __forceinline__ float warp_reduce_sum(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val += __shfl_xor_sync(0xFFFFFFFF, val, offset); - return val; -} - - -__device__ __forceinline__ float row_max_reduce(float2 acc) { - return max_noftz(acc.x, acc.y); -} - - -__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { - #pragma unroll - for (int j = 0; j < 16; j++) { - if (j % 2 == 0) - acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); - else - acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); - } -} - - -__device__ __forceinline__ float2 ex2_emulation_f32x2_value(float2 value) { - const float c0 = 1.0f, c1 = 0.695146143436431884765625f; - const float c2 = 0.227564394474029541015625f, c3 = 0.077119089663028717041015625f; - const float magic = 12582912.0f; - float x0 = max_noftz(value.x, -127.0f), x1 = max_noftz(value.y, -127.0f); - float2 xc2 = make_float2(x0, x1), magic2 = make_float2(magic, magic); - float2 xr2; - asm("add.rm.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xr2) - : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&magic2)); - float2 c3_2 = make_float2(c3, c3), c2_2 = make_float2(c2, c2); - float2 c1_2 = make_float2(c1, c1), c0_2 = make_float2(c0, c0); - float2 xrb2, xfrac2; - asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xrb2) - : "l"(*(unsigned long long*)&xr2), "l"(*(unsigned long long*)&magic2)); - asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xfrac2) - : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&xrb2)); - float2 poly2; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&c3_2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c2_2)); - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c1_2)); - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c0_2)); - int x0r_i, x1r_i, p0_i, p1_i; - asm("mov.b64 {%0, %1}, %2;" : "=r"(x0r_i), "=r"(x1r_i) : "l"(*(unsigned long long*)&xr2)); - asm("mov.b64 {%0, %1}, %2;" : "=r"(p0_i), "=r"(p1_i) : "l"(*(unsigned long long*)&poly2)); - float r0, r1; - asm("mov.b32 %0, %1;" : "=f"(r0) : "r"((x0r_i << 23) + p0_i)); - asm("mov.b32 %0, %1;" : "=f"(r1) : "r"((x1r_i << 23) + p1_i)); - return make_float2(r0, r1); -} - -__device__ __forceinline__ void ex2_emulation_f32x2(float* x0_ptr, float* x1_ptr) { - float2 result = ex2_emulation_f32x2_value(make_float2(*x0_ptr, *x1_ptr)); - *x0_ptr = result.x; *x1_ptr = result.y; -} - -__device__ __forceinline__ void softmax_frag_exp2_cast( - float* sv, uint32_t* pv, int use_emu) -{ - #pragma unroll - for (int j = 0; j < 16; j++) { - if (use_emu && j >= 12) - ex2_emulation_f32x2(&sv[j*2], &sv[j*2+1]); - else { - sv[j*2] = approx_exp2(sv[j*2]); - sv[j*2+1] = approx_exp2(sv[j*2+1]); - } - } - #pragma unroll - for (int j = 0; j < 16; j++) { - __nv_bfloat162 bf = __float22bfloat162_rn({sv[j*2], sv[j*2+1]}); - pv[j] = reinterpret_cast(bf); - } -} - - - -__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { - const float2* sv2 = reinterpret_cast(sv); - #pragma unroll - for (int j = 0; j < 16; j++) { - asm("add.f32x2 %0, %1, %2;" - : "+l"(reinterpret_cast(*acc)) - : "l"(reinterpret_cast(*acc)), - "l"(reinterpret_cast(sv2[j]))); - } -} - - -__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { - asm("mul.rn.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { - asm("mul.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { - asm("add.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { - asm("add.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { - asm("sub.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { - asm("sub.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("add.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { - float2 r; - asm("sub.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("sub.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ void fma_scale_x32( - float* sv, const float2* scale2, const float2* neg_max2) -{ - float2* sv_2 = reinterpret_cast(sv); - #pragma unroll - for (int j = 0; j < 16; j++) - fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); -} - -__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("mul.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - - -// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) - -__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("add.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("add.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("add.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("add.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("add.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("add.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("add.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - - -__device__ __forceinline__ void fence_async_shared() { - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); -} - - -__device__ __forceinline__ uint64_t make_smem_desc(int addr) { - const int SBO = 1024; - return desc_encode(addr) - | (desc_encode(SBO) << 32ULL) - | (1ULL << 46ULL) - | (2ULL << 61ULL); -} - - -__device__ __forceinline__ void tma_2d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3}], [%4];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_4d_gmem2smem_cta2( - int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.4d.shared::cluster.global" - ".mbarrier::complete_tx::bytes.cta_group::2" - " [%0], [%1, {%2, %3, %4, %5}], [%6];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr) { - // Canonical .shared::cta form for non-multicast gather4, matching - // trtllm-gen / cuda_ptx and the PTX ISA qualifier order - // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, - // .shared::cluster is reserved for the multicast variant (ctaMask). - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem_cta2( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr) { - // Canonical .shared::cta form; see tma_gather4_gmem2smem above. - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" - ".mbarrier::complete_tx::bytes.cta_group::2" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem_mc( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr, unsigned short cta_mask) { - // Multicast variant: the PTX grammar ties the .shared::cluster - // destination to .multicast::cluster + ctaMask (cf. cuda_ptx / - // SM100_TMA_LOAD_MULTICAST_2D_GATHER4 in CUTLASS). - asm volatile( - "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" - ".mbarrier::complete_tx::bytes.multicast::cluster" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr), "h"(cta_mask) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem_mc_cta2( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr, unsigned short cta_mask) { - // Multicast + cta_group::2 variant; see tma_gather4_gmem2smem_mc. - asm volatile( - "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" - ".mbarrier::complete_tx::bytes.multicast::cluster.cta_group::2" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr), "h"(cta_mask) : "memory"); -} - - -__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { - asm volatile( - "{\n\t" - ".reg .b16 lo, hi;\n\t" - "mov.b32 {lo, hi}, %1;\n\t" - "tcgen05.commit.cta_group::2.mbarrier::arrive::one" - ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" - "}\n" - :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); -} - - -__device__ __forceinline__ void tmem_st_x8_u32(int addr, const uint32_t* src) { - asm volatile( - "tcgen05.st.sync.aligned.32x32b.x8.b32" - " [%0], {%1,%2,%3,%4,%5,%6,%7,%8};" - :: "r"(addr), - "r"(src[0]), "r"(src[1]), "r"(src[2]), "r"(src[3]), - "r"(src[4]), "r"(src[5]), "r"(src[6]), "r"(src[7])); -} - - -__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { - uint32_t result; - asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" - : "=r"(result) : "r"(val)); - return result; -} - -extern "C" { - -__global__ __launch_bounds__(512, 1) __cluster_dims__(2,1,1) void -kernel_cake_dsv4_9e1188b192c9c806d7c6(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items) -{ - const int tid = threadIdx.x; - const int warp = make_warp_uniform(tid / 32); - const int lane = tid % 32; - - extern __shared__ __align__(1024) char smem_raw[]; - int smem; - smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); - - const int mbar_base = smem; - #define q_full_addr (mbar_base + 0) - #define q_empty_addr (mbar_base + 8) - #define k_full_addr (mbar_base + 16) - #define k_empty_addr (mbar_base + 40) - #define v_full_addr (mbar_base + 64) - #define v_empty_addr (mbar_base + 80) - #define index_full_addr (mbar_base + 96) - #define index_empty_addr (mbar_base + 144) - #define s_full_addr (mbar_base + 192) - #define p_full_addr (mbar_base + 208) - #define corr_done_addr (mbar_base + 224) - #define stats_addr (mbar_base + 240) - #define sum_ready_addr (mbar_base + 256) - #define o_done_addr (mbar_base + 264) - #define pv_done_addr (mbar_base + 272) - #define s_seeded_addr (mbar_base + 280) - #define q_pair_ready_addr (mbar_base + 288) - #define kv_pair_ready_addr (mbar_base + 296) - #define pv_pair_ready_addr (mbar_base + 304) - #define tmem_dealloc_addr (mbar_base + 320) - #define tmem_dealloc_peer_addr (mbar_base + 328) - - const int bid = blockIdx.x; - const int num_bids = gridDim.x; - const unsigned int clusters_x = gridDim.x / 2; - const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; - const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; - - int cta_rank; - asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); - - // Kernel setup ops - uint8_t* smem_q = reinterpret_cast(smem_raw + 1024); - const int smem_q_addr = smem + 1024; - uint8_t* smem_kv = reinterpret_cast(smem_raw + 33792); - const int smem_kv_addr = smem + 33792; - uint8_t* smem_v = reinterpret_cast(smem_raw + 33792); - const int smem_v_addr = smem + 33792; - uint8_t* smem_q_full = reinterpret_cast(smem_raw + 1024); - const int smem_q_full_addr = smem + 1024; - uint8_t* smem_k_full = reinterpret_cast(smem_raw + 33792); - const int smem_k_full_addr = smem + 33792; - uint8_t* smem_v_full = reinterpret_cast(smem_raw + 132096); - const int smem_v_full_addr = smem + 132096; - float* smem_stats_max = reinterpret_cast(smem_raw + 197632); - const int smem_stats_max_addr = smem + 197632; - float* smem_stats_sum = reinterpret_cast(smem_raw + 198656); - const int smem_stats_sum_addr = smem + 198656; - float* smem_stats_final_max = reinterpret_cast(smem_raw + 199168); - const int smem_stats_final_max_addr = smem + 199168; - float* smem_softmax_warp_pair_exchange = reinterpret_cast(smem_raw + 199680); - const int smem_softmax_warp_pair_exchange_addr = smem + 199680; - float* smem_corr_warp_pair_exchange = reinterpret_cast(smem_raw + 200704); - const int smem_corr_warp_pair_exchange_addr = smem + 200704; - int* smem_sparse_indices = reinterpret_cast(smem_raw + 201216); - const int smem_sparse_indices_addr = smem + 201216; - uint8_t* smem_p_fp8 = reinterpret_cast(smem_raw + 207360); - const int smem_p_fp8_addr = smem + 207360; - asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); - asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); - - // Mbarrier init (21 pipeline groups, 0 ordered-sequence groups, 42 barriers) - // Mbarriers at smem_raw[0..336) - - if (warp == 12) { - uint32_t leader = elect_sync(); - if (leader) { - // q_full: 1 barriers, init_count=2 - mbarrier_init(smem + 0, 2); - // q_empty: 1 barriers, init_count=1 - mbarrier_init(smem + 8, 1); - // --- pipeline 'k_pipe' --- - // k_full: 3 barriers, init_count=2 - mbarrier_init(smem + 16, 2); - mbarrier_init(smem + 24, 2); - mbarrier_init(smem + 32, 2); - // k_empty: 3 barriers, init_count=1 - mbarrier_init(smem + 40, 1); - mbarrier_init(smem + 48, 1); - mbarrier_init(smem + 56, 1); - // --- pipeline 'v_pipe' --- - // v_full: 2 barriers, init_count=2 - mbarrier_init(smem + 64, 2); - mbarrier_init(smem + 72, 2); - // v_empty: 2 barriers, init_count=1 - mbarrier_init(smem + 80, 1); - mbarrier_init(smem + 88, 1); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - if (warp == 9) { - uint32_t leader = elect_sync(); - if (leader) { - // --- pipeline 'index_pipe' --- - // index_full: 6 barriers, init_count=32 - mbarrier_init(smem + 96, 32); - mbarrier_init(smem + 104, 32); - mbarrier_init(smem + 112, 32); - mbarrier_init(smem + 120, 32); - mbarrier_init(smem + 128, 32); - mbarrier_init(smem + 136, 32); - // index_empty: 6 barriers, init_count=256 - mbarrier_init(smem + 144, 256); - mbarrier_init(smem + 152, 256); - mbarrier_init(smem + 160, 256); - mbarrier_init(smem + 168, 256); - mbarrier_init(smem + 176, 256); - mbarrier_init(smem + 184, 256); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - if (warp == 0) { - uint32_t leader = elect_sync(); - if (leader) { - // s_full: stages (0,), init_count=1 - mbarrier_init(smem + 192, 1); - // o_done: 1 barriers, init_count=1 - mbarrier_init(smem + 264, 1); - // pv_done: 1 barriers, init_count=1 - mbarrier_init(smem + 272, 1); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - if (warp == 1) { - uint32_t leader = elect_sync(); - if (leader) { - // s_full: stages (1,), init_count=1 - mbarrier_init(smem + 200, 1); - // p_full: 2 barriers, init_count=128 - mbarrier_init(smem + 208, 128); - mbarrier_init(smem + 216, 128); - // corr_done: stages (0,), init_count=128 - mbarrier_init(smem + 224, 128); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - if (warp == 5) { - uint32_t leader = elect_sync(); - if (leader) { - // corr_done: stages (1,), init_count=128 - mbarrier_init(smem + 232, 128); - // stats: 2 barriers, init_count=128 - mbarrier_init(smem + 240, 128); - mbarrier_init(smem + 248, 128); - // sum_ready: 1 barriers, init_count=128 - mbarrier_init(smem + 256, 128); - // tmem_dealloc_peer: 1 barriers, init_count=32 - mbarrier_init(smem + 328, 32); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - if (warp == 2) { - uint32_t leader = elect_sync(); - if (leader) { - // s_seeded: 1 barriers, init_count=256 - mbarrier_init(smem + 280, 256); - // q_pair_ready: 1 barriers, init_count=64 - mbarrier_init(smem + 288, 64); - // kv_pair_ready: 1 barriers, init_count=64 - mbarrier_init(smem + 296, 64); - // pv_pair_ready: stages (0,), init_count=64 - mbarrier_init(smem + 304, 64); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - if (warp == 4) { - uint32_t leader = elect_sync(); - if (leader) { - // pv_pair_ready: stages (1,), init_count=64 - mbarrier_init(smem + 312, 64); - // tmem_dealloc: 1 barriers, init_count=416 - mbarrier_init(smem + 320, 416); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - - __syncwarp(); - - asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); - asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); - - // TMEM alloc (512 columns, 512 used) - volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 336); - if (warp == 0) { - int _tmem_hold = smem + 336; - asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); - __syncwarp(); - asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); - } - - __syncthreads(); - asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); - asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); - - const int taddr = tmem_addr_storage[0]; - - // Kernel post-init ops - const int tmem_tmem_scratch = taddr; - - // ---- Ordered hardware-WG register redistribution ---- - // Inc phase consumes the registers released above. - if (warp >= 8 && warp <= 11) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 136;"); - } - - // ---- Role: index_warp ---- - if (warp == 9) { - { // index_warp_main - const int index_dummy = 0; - unsigned int _phase_index_empty = 1; - unsigned int _phase_q_full_0 = 0; - { - int all_num_kv_tiles = (sparse_topk + 128 - 1) / 128; - unsigned int index_prod_stage = 0; - #pragma unroll - for (unsigned int static_work = 0; static_work < 1; static_work++) { - int split_idx = 0; - int query_idx = bid / 2; - int tiles_per_split = all_num_kv_tiles + 1 - 1; - int first_tile = split_idx * tiles_per_split; - int num_index_passes = (tiles_per_split + 1) / 2; - #pragma unroll 1 - for (int index_pass = 0; index_pass < num_index_passes; index_pass++) { - mbarrier_wait(index_empty_addr + (index_prod_stage) * 8, _phase_index_empty); - int index_stage_base = smem_sparse_indices_addr + index_prod_stage * 1024; - int stage_row_base = index_pass * 256; - #pragma unroll - for (int index_half = 0; index_half < 2; index_half++) { - int index_offset = index_half * 128 + lane * 4; - int tile_col_base = first_tile * 128 + stage_row_base + index_half * 128; - int* row_tile_ptr = ((tile_col_base < 128) ? swa_indices + (query_idx * swa_index_stride) : compressed_indices + (query_idx * compressed_index_stride + (tile_col_base - 128))); - int lane_col = tile_col_base + lane * 4; - if (lane_col + 4 <= sparse_topk) { - asm volatile("cp.async.cg.shared::cta.global.L2::128B [%0], [%1], 16;" - :: "r"(index_stage_base + index_offset * 4), "l"(row_tile_ptr + (lane * 4))); - } else { - int tail_rows[4]; - #pragma unroll - for (int row_i = 0; row_i < 4; row_i++) { - tail_rows[row_i] = -1; - } - #pragma unroll - for (int row_i_1 = 0; row_i_1 < 4; row_i_1++) { - if (lane_col + row_i_1 < sparse_topk) { - tail_rows[row_i_1] = row_tile_ptr[lane * 4 + row_i_1]; - } - } - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(smem_sparse_indices_addr + ((unsigned int)index_stage_base - smem_sparse_indices_addr + (unsigned int)(index_offset * 4))), "r"(tail_rows[0]), "r"(tail_rows[1]), "r"(tail_rows[2]), "r"(tail_rows[3]) : "memory"); - } - } - asm volatile("cp.async.commit_group;"); - asm volatile("cp.async.wait_group 0;"); - mbarrier_arrive(index_full_addr + (index_prod_stage) * 8); - index_prod_stage += 1; - if (index_prod_stage == 6) { index_prod_stage = 0; _phase_index_empty ^= 1; } - } - } - } - } - } - // ---- Role: load_warp ---- - if (warp >= 12 && warp <= 15) { - asm volatile("setmaxnreg.dec.sync.aligned.u32 72;"); - { // load_warp_main - const int wg2_dummy = 0; - const int load_warp_rank = warp - ((1) ? 12 : 9); - int all_num_kv_tiles_1 = (sparse_topk + 128 - 1) / 128; - unsigned int load_k_stage = 0; - unsigned int load_v_stage = 0; - unsigned int load_kv_stage = 0; - unsigned int load_k_index_stage = 0; - unsigned int load_v_index_stage = 0; - int k_cta_offset = bid % 2 * 64; - unsigned int _phase_q_empty_0 = 1; - unsigned int _phase_index_full = 0; - unsigned int _phase_k_empty = 1; - unsigned int _phase_k_empty_1 = 1; - unsigned int _phase_v_empty = 1; - #pragma unroll - for (unsigned int work_idx = 0; work_idx < 1; work_idx++) { - int split_idx_1 = 0; - int query_idx_1 = bid / 2 >> 1; - int v_chunk = bid / 2 & 1; - { - split_idx_1 = 0; - query_idx_1 = bid / 2; - v_chunk = bid % 2; - } - int tiles_per_split_1 = all_num_kv_tiles_1 + 1 - 1; - int first_tile_1 = split_idx_1 * tiles_per_split_1; - int num_kv_tiles = tiles_per_split_1; - int sparse_extent = num_kv_tiles * 128; - mbarrier_wait(q_empty_addr, _phase_q_empty_0); - _phase_q_empty_0 ^= 1; - if (load_warp_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); - #pragma unroll - for (int q_stage = 0; q_stage < 4; q_stage++) { - asm volatile( - "cp.async.bulk.tensor.4d.shared::cluster.global.mbarrier::complete_tx::bytes.multicast::cluster.cta_group::2" - " [%0], [%1, {%2, %3, %4, %5}], [%6], %7;" - :: "r"(smem_q_full_addr + (unsigned int)(q_stage * 8192)), "l"((&tmap_q)), "r"(0), "r"(bid % 2 * 64), "r"(q_stage), "r"(query_idx_1), - "r"(((q_full_addr) & 0xFEFFFFFF)), "h"((uint16_t)(1 << bid % 2)) : "memory"); - } - } - } - { - mbarrier_wait(index_full_addr + (load_k_index_stage) * 8, _phase_index_full); - asm volatile("tcgen05.fence::after_thread_sync;"); - } - int k_index_stage_base = smem_sparse_indices_addr + load_k_index_stage * 1024; - { - mbarrier_wait(k_empty_addr + (load_k_stage) * 8, _phase_k_empty); - if (load_warp_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); - } - } - int k_dst = smem_k_full_addr + load_k_stage * 32768 + (unsigned int)(load_warp_rank * 8192); - int group = lane; - if (group < 16) { - int group_offset = k_cta_offset + group * 4; - int raw_rows[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows[0])), "=r"(*reinterpret_cast(&raw_rows[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows[(0) + 3])) - : "r"(k_index_stage_base + group_offset * 4)); - int row0 = ((raw_rows[0] >= 0) ? raw_rows[0] : 0); - int row1 = ((raw_rows[1] >= 0) ? raw_rows[1] : 0); - int row2 = ((raw_rows[2] >= 0) ? raw_rows[2] : 0); - int row3 = ((raw_rows[3] >= 0) ? raw_rows[3] : 0); - if (first_tile_1 == 0) { - tma_gather4_gmem2smem_mc_cta2(k_dst + group * 512, (&tmap_swa_kv), load_warp_rank * 128, row0, row1, row2, row3, ((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } else { - tma_gather4_gmem2smem_mc_cta2(k_dst + group * 512, (&tmap_compressed_kv), load_warp_rank * 128, row0, row1, row2, row3, ((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } - } - load_k_stage += 1; - if (load_k_stage == 3) { load_k_stage = 0; _phase_k_empty ^= 1; } - } - #pragma unroll 1 - for (int tile = 1; tile < num_kv_tiles; tile++) { - if ((1 & (int)((tile & 1) == 0)) != 0) { - mbarrier_wait(index_full_addr + (load_k_index_stage) * 8, _phase_index_full); - asm volatile("tcgen05.fence::after_thread_sync;"); - } - k_index_stage_base = smem_sparse_indices_addr + load_k_index_stage * 1024; - { - mbarrier_wait(k_empty_addr + (load_k_stage) * 8, _phase_k_empty); - if (load_warp_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); - } - } - int k_dst_1 = smem_k_full_addr + load_k_stage * 32768 + (unsigned int)(load_warp_rank * 8192); - int group_1 = lane; - if (group_1 < 16) { - int group_offset_1 = (tile & 1) * 128 + k_cta_offset + group_1 * 4; - int raw_rows_1[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_1[0])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 3])) - : "r"(k_index_stage_base + group_offset_1 * 4)); - int row0_1 = ((raw_rows_1[0] >= 0) ? raw_rows_1[0] : 0); - int row1_1 = ((raw_rows_1[1] >= 0) ? raw_rows_1[1] : 0); - int row2_1 = ((raw_rows_1[2] >= 0) ? raw_rows_1[2] : 0); - int row3_1 = ((raw_rows_1[3] >= 0) ? raw_rows_1[3] : 0); - tma_gather4_gmem2smem_mc_cta2(k_dst_1 + group_1 * 512, (&tmap_compressed_kv), load_warp_rank * 128, row0_1, row1_1, row2_1, row3_1, ((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } - load_k_stage += 1; - if (load_k_stage == 3) { load_k_stage = 0; _phase_k_empty ^= 1; } - if ((tile & 1) != 0) { - load_k_index_stage += 1; - if (load_k_index_stage == 6) { load_k_index_stage = 0; _phase_index_full ^= 1; } - } - } - int prev_tile = tile - 1; - int global_prev_tile = first_tile_1 + prev_tile; - int v_index_stage_base = smem_sparse_indices_addr + load_v_index_stage * 1024; - { - mbarrier_wait(v_empty_addr + (load_v_stage) * 8, _phase_v_empty); - if (load_warp_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); - } - } - if (load_warp_rank < 2) { - int v_dst = smem_v_full_addr + load_v_stage * 32768 + (unsigned int)(load_warp_rank * 16384); - int v_col = v_chunk * 256 + load_warp_rank * 128; - int group_2 = lane; - int group_offset_2 = (prev_tile & 1) * 128 + group_2 * 4; - int raw_rows_2[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_2[0])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 3])) - : "r"(v_index_stage_base + group_offset_2 * 4)); - int row0_2 = ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : 0); - int row1_2 = ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : 0); - int row2_2 = ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : 0); - int row3_2 = ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : 0); - if (global_prev_tile == 0) { - tma_gather4_gmem2smem_mc_cta2(v_dst + group_2 * 512, (&tmap_swa_kv), v_col, row0_2, row1_2, row2_2, row3_2, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } else { - tma_gather4_gmem2smem_mc_cta2(v_dst + group_2 * 512, (&tmap_compressed_kv), v_col, row0_2, row1_2, row2_2, row3_2, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } - } - load_v_stage += 1; - if (load_v_stage == 2) { load_v_stage = 0; _phase_v_empty ^= 1; } - if ((prev_tile & 1) != 0) { - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(index_empty_addr + (load_v_index_stage) * 8); - load_v_index_stage = load_v_index_stage + 1; - if (load_v_index_stage == 6) { - load_v_index_stage = 0; - } - } - } - } - int last_tile = num_kv_tiles - 1; - int global_last_tile = first_tile_1 + last_tile; - { - mbarrier_wait(v_empty_addr + (load_v_stage) * 8, _phase_v_empty); - if (load_warp_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); - } - } - if (load_warp_rank < 2) { - int v_dst_1 = smem_v_full_addr + load_v_stage * 32768 + (unsigned int)(load_warp_rank * 16384); - int v_col_1 = v_chunk * 256 + load_warp_rank * 128; - int group_3 = lane; - int group_offset_3 = (last_tile & 1) * 128 + group_3 * 4; - int raw_rows_3[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_3[0])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 3])) - : "r"(smem_sparse_indices_addr + load_v_index_stage * 1024 + (unsigned int)(group_offset_3 * 4))); - int row0_3 = ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : 0); - int row1_3 = ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : 0); - int row2_3 = ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : 0); - int row3_3 = ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : 0); - if (global_last_tile == 0) { - tma_gather4_gmem2smem_mc_cta2(v_dst_1 + group_3 * 512, (&tmap_swa_kv), v_col_1, row0_3, row1_3, row2_3, row3_3, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } else { - tma_gather4_gmem2smem_mc_cta2(v_dst_1 + group_3 * 512, (&tmap_compressed_kv), v_col_1, row0_3, row1_3, row2_3, row3_3, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } - } - load_v_stage += 1; - if (load_v_stage == 2) { load_v_stage = 0; _phase_v_empty ^= 1; } - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(index_empty_addr + (load_v_index_stage) * 8); - load_v_index_stage = load_v_index_stage + 1; - if (load_v_index_stage == 6) { - load_v_index_stage = 0; - } - } - } - mbarrier_arrive(tmem_dealloc_addr); - } - } - // ---- Role: softmax_wg ---- - if (warp <= 3) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 152;"); - { // softmax_wg_main - float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; - const int wg_dummy_inc = 0; - int all_num_kv_tiles_2 = (sparse_topk + 128 - 1) / 128; - const int tmem_row_base = ((1) ? warp % 2 * 32 : warp % 4 * 32); - const int tmem_score_row_base = ((1) ? (int)(warp % 4 * 32) : tmem_row_base); - const int n_half = ((1) ? (int)(warp % 4 / 2) : 0); - const int my_row = tmem_row_base + lane; - const int stats_row = n_half * 64 + my_row; - int softmax_tile_cursor = 0; - unsigned int softmax_index_stage = 0; - unsigned int _phase_q_full_0_1 = 0; - #pragma unroll - for (unsigned int work_idx_1 = 0; work_idx_1 < 1; work_idx_1++) { - int split_idx_2 = 0; - int query_idx_2 = bid / 2 >> 1; - { - split_idx_2 = 0; - query_idx_2 = bid / 2; - } - int tiles_per_split_2 = all_num_kv_tiles_2 + 1 - 1; - int first_tile_2 = split_idx_2 * tiles_per_split_2; - int num_kv_tiles_1 = tiles_per_split_2; - int _max_0 = ((sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) : (0)); - int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); - int active_topk = _min_0; - float row_max_val = -CAKE_INF; - float row_sum_val = 0.0f; - int sink_head = ((1) ? bid % 2 * 64 + my_row : my_row); - if (has_sinks != 0 && sink_head < num_heads && split_idx_2 == 0) { - row_max_val = sinks[sink_head] * 1.4426950408889634f / softmax_scale_log2; - row_sum_val = 1.0f; - } - #pragma unroll 1 - for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { - int pipeline_tile = softmax_tile_cursor + tile_1; - int phase = pipeline_tile & 1; - int s_wait_phase = pipeline_tile >> 1 & 1; - mbarrier_wait(s_full_addr + (phase) * 8, s_wait_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - int s_off = ((phase != 0) ? 128 : 0); - int s_base = taddr + (unsigned int)s_off + (unsigned int)(tmem_score_row_base << 16); - float new_max = row_max_val; - int valid_sparse_cols = ((active_topk < sparse_topk) ? active_topk : sparse_topk); - valid_sparse_cols = valid_sparse_cols - (first_tile_2 + tile_1) * 128 - n_half * 64; - if (valid_sparse_cols < 0) { - valid_sparse_cols = 0; - } - if (valid_sparse_cols > 64) { - valid_sparse_cols = 64; - } - int index_tile_addr = smem_sparse_indices_addr + softmax_index_stage * 1024 + (unsigned int)((tile_1 & 1) * 512); - int index_lane_addr = index_tile_addr + (n_half * 64 + lane) * 4; - int staged_index[2]; - unsigned int invalid_cols[2]; - unsigned int any_invalid = 0; - #pragma unroll - for (int w = 0; w < 2; w++) { - asm volatile("ld.shared.b32 %0, [%1];" : "=r"(*reinterpret_cast(&staged_index[w])) : "r"(index_lane_addr + w * 128)); - unsigned int _vote_0 = __ballot_sync(0xFFFFFFFF, staged_index[w] < 0); - invalid_cols[w] = _vote_0; - any_invalid = any_invalid | invalid_cols[w]; - } - if ((tile_1 & 1) != 0 || tile_1 + 1 == num_kv_tiles_1) { - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(index_empty_addr + (softmax_index_stage) * 8); - softmax_index_stage = softmax_index_stage + 1; - if (softmax_index_stage == 6) { - softmax_index_stage = 0; - } - } - float _tmem_load_0[4]; - tmem_ld_x4(&_tmem_load_0[0], s_base); - float _tmem_load_1[64]; - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]) - : "r"(s_base)); - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]) - : "r"(s_base + 32)); - { - uint32_t _slice_lo_mask_0; - { - int _lim_0 = valid_sparse_cols; - if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } - else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); - } - } - if (!(_slice_lo_mask_0 & (1u << 0))) _tmem_load_1[0] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 1))) _tmem_load_1[1] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 2))) _tmem_load_1[2] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 3))) _tmem_load_1[3] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 4))) _tmem_load_1[4] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 5))) _tmem_load_1[5] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 6))) _tmem_load_1[6] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 7))) _tmem_load_1[7] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 8))) _tmem_load_1[8] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 9))) _tmem_load_1[9] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 10))) _tmem_load_1[10] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 11))) _tmem_load_1[11] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 12))) _tmem_load_1[12] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 13))) _tmem_load_1[13] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 14))) _tmem_load_1[14] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 15))) _tmem_load_1[15] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 16))) _tmem_load_1[16] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 17))) _tmem_load_1[17] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 18))) _tmem_load_1[18] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 19))) _tmem_load_1[19] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 20))) _tmem_load_1[20] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 21))) _tmem_load_1[21] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 22))) _tmem_load_1[22] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 23))) _tmem_load_1[23] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 24))) _tmem_load_1[24] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 25))) _tmem_load_1[25] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 26))) _tmem_load_1[26] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 27))) _tmem_load_1[27] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 28))) _tmem_load_1[28] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 29))) _tmem_load_1[29] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 30))) _tmem_load_1[30] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 31))) _tmem_load_1[31] = -CAKE_INF; - uint32_t _slice_lo_mask_1; - { - int _lim_1 = valid_sparse_cols - 32; - if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } - else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); - } - } - if (!(_slice_lo_mask_1 & (1u << 0))) _tmem_load_1[32] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 1))) _tmem_load_1[33] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 2))) _tmem_load_1[34] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 3))) _tmem_load_1[35] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 4))) _tmem_load_1[36] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 5))) _tmem_load_1[37] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 6))) _tmem_load_1[38] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 7))) _tmem_load_1[39] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 8))) _tmem_load_1[40] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 9))) _tmem_load_1[41] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 10))) _tmem_load_1[42] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 11))) _tmem_load_1[43] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 12))) _tmem_load_1[44] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 13))) _tmem_load_1[45] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 14))) _tmem_load_1[46] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 15))) _tmem_load_1[47] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 16))) _tmem_load_1[48] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 17))) _tmem_load_1[49] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 18))) _tmem_load_1[50] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 19))) _tmem_load_1[51] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 20))) _tmem_load_1[52] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 21))) _tmem_load_1[53] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 22))) _tmem_load_1[54] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 23))) _tmem_load_1[55] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 24))) _tmem_load_1[56] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 25))) _tmem_load_1[57] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 26))) _tmem_load_1[58] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 27))) _tmem_load_1[59] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 28))) _tmem_load_1[60] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 29))) _tmem_load_1[61] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 30))) _tmem_load_1[62] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 31))) _tmem_load_1[63] = -CAKE_INF; - if (any_invalid != 0) { - #pragma unroll - for (int w_1 = 0; w_1 < 2; w_1++) { - unsigned int mask_word = invalid_cols[w_1]; - _tmem_load_1[w_1 * 32] = (((mask_word & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32]); - _tmem_load_1[w_1 * 32 + 1] = (((mask_word >> 1 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 1]); - _tmem_load_1[w_1 * 32 + 2] = (((mask_word >> 2 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 2]); - _tmem_load_1[w_1 * 32 + 3] = (((mask_word >> 3 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 3]); - _tmem_load_1[w_1 * 32 + 4] = (((mask_word >> 4 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 4]); - _tmem_load_1[w_1 * 32 + 5] = (((mask_word >> 5 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 5]); - _tmem_load_1[w_1 * 32 + 6] = (((mask_word >> 6 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 6]); - _tmem_load_1[w_1 * 32 + 7] = (((mask_word >> 7 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 7]); - _tmem_load_1[w_1 * 32 + 8] = (((mask_word >> 8 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 8]); - _tmem_load_1[w_1 * 32 + 9] = (((mask_word >> 9 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 9]); - _tmem_load_1[w_1 * 32 + 10] = (((mask_word >> 10 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 10]); - _tmem_load_1[w_1 * 32 + 11] = (((mask_word >> 11 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 11]); - _tmem_load_1[w_1 * 32 + 12] = (((mask_word >> 12 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 12]); - _tmem_load_1[w_1 * 32 + 13] = (((mask_word >> 13 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 13]); - _tmem_load_1[w_1 * 32 + 14] = (((mask_word >> 14 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 14]); - _tmem_load_1[w_1 * 32 + 15] = (((mask_word >> 15 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 15]); - _tmem_load_1[w_1 * 32 + 16] = (((mask_word >> 16 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 16]); - _tmem_load_1[w_1 * 32 + 17] = (((mask_word >> 17 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 17]); - _tmem_load_1[w_1 * 32 + 18] = (((mask_word >> 18 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 18]); - _tmem_load_1[w_1 * 32 + 19] = (((mask_word >> 19 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 19]); - _tmem_load_1[w_1 * 32 + 20] = (((mask_word >> 20 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 20]); - _tmem_load_1[w_1 * 32 + 21] = (((mask_word >> 21 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 21]); - _tmem_load_1[w_1 * 32 + 22] = (((mask_word >> 22 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 22]); - _tmem_load_1[w_1 * 32 + 23] = (((mask_word >> 23 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 23]); - _tmem_load_1[w_1 * 32 + 24] = (((mask_word >> 24 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 24]); - _tmem_load_1[w_1 * 32 + 25] = (((mask_word >> 25 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 25]); - _tmem_load_1[w_1 * 32 + 26] = (((mask_word >> 26 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 26]); - _tmem_load_1[w_1 * 32 + 27] = (((mask_word >> 27 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 27]); - _tmem_load_1[w_1 * 32 + 28] = (((mask_word >> 28 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 28]); - _tmem_load_1[w_1 * 32 + 29] = (((mask_word >> 29 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 29]); - _tmem_load_1[w_1 * 32 + 30] = (((mask_word >> 30 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 30]); - _tmem_load_1[w_1 * 32 + 31] = (((mask_word >> 31 & 1) != 0) ? -CAKE_INF : _tmem_load_1[w_1 * 32 + 31]); - } - } - float2 _reg_reduce_max2_2 = {-CAKE_INF, -CAKE_INF}; - row_max_x32_accum(&_tmem_load_1[0], _reg_reduce_max2_2); - row_max_x32_accum(&_tmem_load_1[32], _reg_reduce_max2_2); - float _tmem_load_1_max = row_max_reduce(_reg_reduce_max2_2); - float _max_1 = max_noftz(new_max, _tmem_load_1_max); - new_max = _max_1; - smem_softmax_warp_pair_exchange[phase * 128 + stats_row] = new_max; - asm volatile("barrier.sync %0, 64;" :: "r"(2 + warp % 2) : "memory"); - float _max_2 = max_noftz(new_max, smem_softmax_warp_pair_exchange[phase * 128 + (stats_row ^ 64)]); - new_max = _max_2; - } - float _fma_0 = __fmaf_rn(row_max_val, softmax_scale_log2, (-new_max) * softmax_scale_log2); - float delta = _fma_0; - float _exp2_0 = approx_exp2(delta); - float exp_delta = _exp2_0; - float acc_scale = ((row_max_val > -CAKE_INF) ? exp_delta : 1.0f); - smem_stats_max[phase * 128 + stats_row] = acc_scale; - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(stats_addr + (phase) * 8); - row_max_val = new_max; - float safe_max = ((new_max == -CAKE_INF) ? 0.0f : new_max); - float max_scaled = safe_max * softmax_scale_log2; - float block_sum = 0.0f; - { - { - const float2 _fma_b2_3 = {softmax_scale_log2, softmax_scale_log2}; - const float2 _fma_c2_4 = {-max_scaled, -max_scaled}; - #pragma unroll - for (int _lf = 0; _lf < 32; _lf++) - fma_f32x2_inplace(&reinterpret_cast(_tmem_load_1)[_lf], _fma_b2_3, _fma_c2_4); - #pragma unroll - for (int _le = 0; _le < 64; _le++) { - _tmem_load_1[_le] = approx_exp2(_tmem_load_1[_le]); - } - uint32_t _fp8_0[16]; - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[0]), "f"(_tmem_load_1[1]), - "f"(_tmem_load_1[2]), "f"(_tmem_load_1[3])); - _fp8_0[0] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[4]), "f"(_tmem_load_1[5]), - "f"(_tmem_load_1[6]), "f"(_tmem_load_1[7])); - _fp8_0[1] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[8]), "f"(_tmem_load_1[9]), - "f"(_tmem_load_1[10]), "f"(_tmem_load_1[11])); - _fp8_0[2] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[12]), "f"(_tmem_load_1[13]), - "f"(_tmem_load_1[14]), "f"(_tmem_load_1[15])); - _fp8_0[3] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[16]), "f"(_tmem_load_1[17]), - "f"(_tmem_load_1[18]), "f"(_tmem_load_1[19])); - _fp8_0[4] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[20]), "f"(_tmem_load_1[21]), - "f"(_tmem_load_1[22]), "f"(_tmem_load_1[23])); - _fp8_0[5] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[24]), "f"(_tmem_load_1[25]), - "f"(_tmem_load_1[26]), "f"(_tmem_load_1[27])); - _fp8_0[6] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[28]), "f"(_tmem_load_1[29]), - "f"(_tmem_load_1[30]), "f"(_tmem_load_1[31])); - _fp8_0[7] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[32]), "f"(_tmem_load_1[33]), - "f"(_tmem_load_1[34]), "f"(_tmem_load_1[35])); - _fp8_0[8] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[36]), "f"(_tmem_load_1[37]), - "f"(_tmem_load_1[38]), "f"(_tmem_load_1[39])); - _fp8_0[9] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[40]), "f"(_tmem_load_1[41]), - "f"(_tmem_load_1[42]), "f"(_tmem_load_1[43])); - _fp8_0[10] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[44]), "f"(_tmem_load_1[45]), - "f"(_tmem_load_1[46]), "f"(_tmem_load_1[47])); - _fp8_0[11] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[48]), "f"(_tmem_load_1[49]), - "f"(_tmem_load_1[50]), "f"(_tmem_load_1[51])); - _fp8_0[12] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[52]), "f"(_tmem_load_1[53]), - "f"(_tmem_load_1[54]), "f"(_tmem_load_1[55])); - _fp8_0[13] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[56]), "f"(_tmem_load_1[57]), - "f"(_tmem_load_1[58]), "f"(_tmem_load_1[59])); - _fp8_0[14] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_1[60]), "f"(_tmem_load_1[61]), - "f"(_tmem_load_1[62]), "f"(_tmem_load_1[63])); - _fp8_0[15] = _packed; - } - int p_stage_base = smem_p_fp8_addr + (unsigned int)(phase * 8192); - #pragma unroll - for (int p_vec = 0; p_vec < 4; p_vec++) { - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_stage_base + (my_row * 128 + (n_half * 4 + p_vec) * 16 ^ (my_row * 128 + (n_half * 4 + p_vec) * 16 >> 7 & 7) << 4))), "r"(_fp8_0[p_vec * 4]), "r"(_fp8_0[p_vec * 4 + 1]), "r"(_fp8_0[p_vec * 4 + 2]), "r"(_fp8_0[p_vec * 4 + 3]) : "memory"); - } - } - } - { - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - } - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - { - asm volatile("" ::: "memory"); - } - mbarrier_arrive(p_full_addr + (phase) * 8); - { - { - asm volatile("" ::: "memory"); - const float2* _reg_reduce_src2_5 = reinterpret_cast(&_tmem_load_1[0]); - float2 _reg_reduce_sum2_5_0 = make_float2(0.0f, 0.0f); - float2 _reg_reduce_sum2_5_1 = make_float2(0.0f, 0.0f); - float2 _reg_reduce_sum2_5_2 = make_float2(0.0f, 0.0f); - float2 _reg_reduce_sum2_5_3 = make_float2(0.0f, 0.0f); - #pragma unroll - for (int _rr = 0; _rr < 8; _rr++) { - add_f32x2_inplace(&_reg_reduce_sum2_5_0, _reg_reduce_src2_5[_rr * 4 + 0]); - add_f32x2_inplace(&_reg_reduce_sum2_5_1, _reg_reduce_src2_5[_rr * 4 + 1]); - add_f32x2_inplace(&_reg_reduce_sum2_5_2, _reg_reduce_src2_5[_rr * 4 + 2]); - add_f32x2_inplace(&_reg_reduce_sum2_5_3, _reg_reduce_src2_5[_rr * 4 + 3]); - } - add_f32x2_inplace(&_reg_reduce_sum2_5_0, _reg_reduce_sum2_5_1); - add_f32x2_inplace(&_reg_reduce_sum2_5_2, _reg_reduce_sum2_5_3); - add_f32x2_inplace(&_reg_reduce_sum2_5_0, _reg_reduce_sum2_5_2); - float _tmem_load_1_sum = _reg_reduce_sum2_5_0.x + _reg_reduce_sum2_5_0.y; - block_sum = _tmem_load_1_sum; - float _fma_2 = __fmaf_rn(row_sum_val, acc_scale, block_sum); - row_sum_val = _fma_2; - } - } - } - smem_stats_sum[stats_row] = row_sum_val; - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(sum_ready_addr); - softmax_tile_cursor = softmax_tile_cursor + num_kv_tiles_1; - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(tmem_dealloc_addr); - } - } - // ---- Role: correction_wg ---- - if (warp >= 4 && warp <= 7) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 144;"); - { // correction_wg_main - float softmax_scale_log2_1 = bmm1_scale[0] * 1.4426950408889634f; - float output_scale = bmm2_scale[0]; - const int wg_dummy_inc_1 = 0; - int all_num_kv_tiles_3 = (sparse_topk + 128 - 1) / 128; - const int tmem_row_base_1 = ((1) ? warp % 2 * 32 : warp % 4 * 32); - const int n_half_1 = ((1) ? (int)(warp % 4 / 2) : 0); - const int my_row_1 = tmem_row_base_1 + lane; - const int stats_row_1 = n_half_1 * 64 + my_row_1; - const int corr_row = tmem_row_base_1 << 16; - int correction_tile_cursor = 0; - unsigned int _phase_q_full_0_2 = 0; - unsigned int _phase_pv_done_0 = 0; - unsigned int _phase_o_done_0 = 0; - unsigned int _phase_sum_ready_0 = 0; - #pragma unroll - for (unsigned int work_idx_2 = 0; work_idx_2 < 1; work_idx_2++) { - int split_idx_3 = 0; - int query_idx_3 = bid / 2 >> 1; - int v_chunk_1 = bid / 2 & 1; - { - split_idx_3 = 0; - query_idx_3 = bid / 2; - v_chunk_1 = bid % 2; - } - int tiles_per_split_3 = all_num_kv_tiles_3 + 1 - 1; - int first_tile_3 = split_idx_3 * tiles_per_split_3; - int num_kv_tiles_2 = tiles_per_split_3; - #pragma unroll 1 - for (int tile_2 = 0; tile_2 < num_kv_tiles_2; tile_2++) { - int pipeline_tile_1 = correction_tile_cursor + tile_2; - int phase_1 = pipeline_tile_1 & 1; - int stats_wait_phase = pipeline_tile_1 >> 1 & 1; - mbarrier_wait(stats_addr + (phase_1) * 8, stats_wait_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - float acc_scale_1 = smem_stats_max[phase_1 * 128 + stats_row_1]; - if (tile_2 > 0) { - mbarrier_wait(pv_done_addr, _phase_pv_done_0); - _phase_pv_done_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - } - if (tile_2 > 0) { - int _vote_1 = __any_sync(0xFFFFFFFF, acc_scale_1 < 1.0f); - int any_rescale = _vote_1; - if (any_rescale != 0) { - #pragma unroll - for (int vs = 0; vs < 2; vs++) { - int o_base = taddr + 256 + (unsigned int)(vs * 128) + (unsigned int)corr_row; - #pragma unroll - for (int c = 0; c < 128; c += 16) { - float _tmem_load_4[16]; - tmem_ld_x16(&_tmem_load_4[0], o_base + c); - const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; - #pragma unroll - for (int _ls = 0; _ls < 8; _ls++) - mul_f32x2_inplace(&reinterpret_cast(_tmem_load_4)[_ls], _scale2_0); - tmem_st_x16_f32(o_base + c, _tmem_load_4); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - } - } - } - } - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(corr_done_addr + (phase_1) * 8); - } - mbarrier_wait(o_done_addr, _phase_o_done_0); - _phase_o_done_0 ^= 1; - mbarrier_wait(sum_ready_addr, _phase_sum_ready_0); - _phase_sum_ready_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - float total_sum = smem_stats_sum[stats_row_1]; - float final_max = smem_stats_final_max[stats_row_1]; - { - smem_corr_warp_pair_exchange[stats_row_1] = total_sum; - asm volatile("barrier.sync %0, 64;" :: "r"(4 + warp % 2) : "memory"); - total_sum = total_sum + smem_corr_warp_pair_exchange[stats_row_1 ^ 64]; - } - float _rcp_0 = approx_rcp(total_sum); - float inv_sum = ((total_sum > 0.0f) ? _rcp_0 : 0.0f); - int head_idx = my_row_1; - { - head_idx = bid % 2 * 64 + my_row_1; - } - int direct_o_offset = (query_idx_3 * num_heads + head_idx) * 512 + v_chunk_1 * 256; - int partial_o_offset = (query_idx_3 * num_heads + head_idx + split_idx_3) * 512 + ((0) ? v_chunk_1 : n_half_1) * 256; - int o_offset = partial_o_offset; - #pragma unroll - for (int vs_1 = 0; vs_1 < 2; vs_1++) { - int o_base_epi = taddr + 256 + (unsigned int)(vs_1 * 128) + (unsigned int)corr_row; - #pragma unroll - for (int c_1 = 0; c_1 < 128; c_1 += 64) { - float _tmem_load_5[64]; - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x64.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, [%64];" - : "=f"(_tmem_load_5[0]), "=f"(_tmem_load_5[1]), "=f"(_tmem_load_5[2]), "=f"(_tmem_load_5[3]), "=f"(_tmem_load_5[4]), "=f"(_tmem_load_5[5]), "=f"(_tmem_load_5[6]), "=f"(_tmem_load_5[7]), "=f"(_tmem_load_5[8]), "=f"(_tmem_load_5[9]), "=f"(_tmem_load_5[10]), "=f"(_tmem_load_5[11]), "=f"(_tmem_load_5[12]), "=f"(_tmem_load_5[13]), "=f"(_tmem_load_5[14]), "=f"(_tmem_load_5[15]), "=f"(_tmem_load_5[16]), "=f"(_tmem_load_5[17]), "=f"(_tmem_load_5[18]), "=f"(_tmem_load_5[19]), "=f"(_tmem_load_5[20]), "=f"(_tmem_load_5[21]), "=f"(_tmem_load_5[22]), "=f"(_tmem_load_5[23]), "=f"(_tmem_load_5[24]), "=f"(_tmem_load_5[25]), "=f"(_tmem_load_5[26]), "=f"(_tmem_load_5[27]), "=f"(_tmem_load_5[28]), "=f"(_tmem_load_5[29]), "=f"(_tmem_load_5[30]), "=f"(_tmem_load_5[31]), "=f"(_tmem_load_5[32]), "=f"(_tmem_load_5[33]), "=f"(_tmem_load_5[34]), "=f"(_tmem_load_5[35]), "=f"(_tmem_load_5[36]), "=f"(_tmem_load_5[37]), "=f"(_tmem_load_5[38]), "=f"(_tmem_load_5[39]), "=f"(_tmem_load_5[40]), "=f"(_tmem_load_5[41]), "=f"(_tmem_load_5[42]), "=f"(_tmem_load_5[43]), "=f"(_tmem_load_5[44]), "=f"(_tmem_load_5[45]), "=f"(_tmem_load_5[46]), "=f"(_tmem_load_5[47]), "=f"(_tmem_load_5[48]), "=f"(_tmem_load_5[49]), "=f"(_tmem_load_5[50]), "=f"(_tmem_load_5[51]), "=f"(_tmem_load_5[52]), "=f"(_tmem_load_5[53]), "=f"(_tmem_load_5[54]), "=f"(_tmem_load_5[55]), "=f"(_tmem_load_5[56]), "=f"(_tmem_load_5[57]), "=f"(_tmem_load_5[58]), "=f"(_tmem_load_5[59]), "=f"(_tmem_load_5[60]), "=f"(_tmem_load_5[61]), "=f"(_tmem_load_5[62]), "=f"(_tmem_load_5[63]) - : "r"(o_base_epi + c_1)); - int gmem_base = o_offset + vs_1 * 128 + c_1; - if (bid % 2 == 0) { - #pragma unroll - for (int j = 0; j < 64; j += 16) { - { - const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; - #if __CUDA_ARCH__ >= 1000 - #pragma unroll - for (int _ps = 0; _ps < 8; _ps++) - mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_5[j])[_ps], _prescale2_1); - #else - #pragma unroll - for (int _ps = 0; _ps < 16; _ps++) - _tmem_load_5[j + _ps] *= inv_sum * output_scale; - #endif - { - __nv_bfloat162 _pk0 = __floats2bfloat162_rn(_tmem_load_5[j + 0], _tmem_load_5[j + 1]); - unsigned _pk_u0 = *reinterpret_cast(&_pk0); - __nv_bfloat162 _pk1 = __floats2bfloat162_rn(_tmem_load_5[j + 2], _tmem_load_5[j + 3]); - unsigned _pk_u1 = *reinterpret_cast(&_pk1); - __nv_bfloat162 _pk2 = __floats2bfloat162_rn(_tmem_load_5[j + 4], _tmem_load_5[j + 5]); - unsigned _pk_u2 = *reinterpret_cast(&_pk2); - __nv_bfloat162 _pk3 = __floats2bfloat162_rn(_tmem_load_5[j + 6], _tmem_load_5[j + 7]); - unsigned _pk_u3 = *reinterpret_cast(&_pk3); - __nv_bfloat162 _pk4 = __floats2bfloat162_rn(_tmem_load_5[j + 8], _tmem_load_5[j + 9]); - unsigned _pk_u4 = *reinterpret_cast(&_pk4); - __nv_bfloat162 _pk5 = __floats2bfloat162_rn(_tmem_load_5[j + 10], _tmem_load_5[j + 11]); - unsigned _pk_u5 = *reinterpret_cast(&_pk5); - __nv_bfloat162 _pk6 = __floats2bfloat162_rn(_tmem_load_5[j + 12], _tmem_load_5[j + 13]); - unsigned _pk_u6 = *reinterpret_cast(&_pk6); - __nv_bfloat162 _pk7 = __floats2bfloat162_rn(_tmem_load_5[j + 14], _tmem_load_5[j + 15]); - unsigned _pk_u7 = *reinterpret_cast(&_pk7); - asm volatile( - "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" - :: "l"((void*)(&((__nv_bfloat16*)(O + (gmem_base + j)))[0])), "r"(_pk_u0), "r"(_pk_u1), "r"(_pk_u2), "r"(_pk_u3), "r"(_pk_u4), "r"(_pk_u5), "r"(_pk_u6), "r"(_pk_u7) : "memory"); - } - } - } - } - } - } - correction_tile_cursor = correction_tile_cursor + num_kv_tiles_2; - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(tmem_dealloc_addr); - } - } - // ---- Role: mma_warp ---- - if (warp == 8) { - { // mma_warp_main - const int wg2_dummy_1 = 0; - int all_num_kv_tiles_4 = (sparse_topk + 128 - 1) / 128; - unsigned int mma_k_stage = 0; - unsigned int mma_v_stage = 0; - unsigned int mma_kv_stage = 0; - int mma_tile_cursor = 0; - unsigned int _phase_s_seeded_0 = 0; - unsigned int _phase_q_full_0_3 = 0; - unsigned int _phase_q_pair_ready_0 = 0; - unsigned int _phase_k_full = 0; - unsigned int _phase_k_full_1 = 0; - unsigned int _phase_kv_pair_ready = 0; - unsigned int _phase_v_full = 0; - #pragma unroll - for (unsigned int work_idx_3 = 0; work_idx_3 < 1; work_idx_3++) { - int split_idx_4 = 0; - { - split_idx_4 = 0; - } - int tiles_per_split_4 = all_num_kv_tiles_4 + 1 - 1; - int first_tile_4 = split_idx_4 * tiles_per_split_4; - int num_kv_tiles_3 = tiles_per_split_4; - { - if (cta_rank == 0) { - mbarrier_wait(q_full_addr, _phase_q_full_0_3); - _phase_q_full_0_3 ^= 1; - } - } - int first_pv = 1; - #pragma unroll 1 - for (int tile_3 = 0; tile_3 < num_kv_tiles_3; tile_3++) { - int pipeline_tile_2 = mma_tile_cursor + tile_3; - int phase_2 = pipeline_tile_2 & 1; - int score_col = ((phase_2 != 0) ? 128 : 0); - { - if (cta_rank == 0) { - mbarrier_wait(k_full_addr + (mma_k_stage) * 8, _phase_k_full); - int _mma_a_lo_0 = (((smem_q_full_addr) >> 4) & 0x3FFF) + (0) * 2048; - int _mma_b_lo_0 = (((smem_k_full_addr) >> 4) & 0x3FFF) + (mma_k_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 136314896;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (score_col))), "r"(0)); - int _mma_a_lo_1 = (((smem_q_full_addr + 8192) >> 4) & 0x3FFF) + (0) * 2048; - int _mma_b_lo_1 = (((smem_k_full_addr + 8192) >> 4) & 0x3FFF) + (mma_k_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 136314896;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); - int _mma_a_lo_2 = (((smem_q_full_addr + 16384) >> 4) & 0x3FFF) + (0) * 2048; - int _mma_b_lo_2 = (((smem_k_full_addr + 16384) >> 4) & 0x3FFF) + (mma_k_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 136314896;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); - int _mma_a_lo_3 = (((smem_q_full_addr + 24576) >> 4) & 0x3FFF) + (0) * 2048; - int _mma_b_lo_3 = (((smem_k_full_addr + 24576) >> 4) & 0x3FFF) + (mma_k_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 136314896;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); - if (cta_rank == 0) { - elect_commit_cg2_multicast(s_full_addr + (phase_2) * 8, (uint16_t)(3)); - } - if (cta_rank == 0) { - elect_commit_cg2_multicast(k_empty_addr + (mma_k_stage) * 8, (uint16_t)(3)); - } - } - mma_k_stage += 1; - if (mma_k_stage == 3) { mma_k_stage = 0; _phase_k_full ^= 1; } - } - if (tile_3 > 0) { - int prev_pipeline_tile = pipeline_tile_2 - 1; - int prev_phase = prev_pipeline_tile & 1; - int pv_wait_phase = prev_pipeline_tile >> 1 & 1; - mbarrier_wait(p_full_addr + (prev_phase) * 8, pv_wait_phase); - mbarrier_wait(corr_done_addr + (prev_phase) * 8, pv_wait_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - { - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cluster.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(pv_pair_ready_addr + (unsigned int)(prev_phase * 8)), "r"(0) : "memory"); - if (cta_rank == 0) { - mbarrier_wait_cluster_hint(pv_pair_ready_addr + (prev_phase) * 8, pv_wait_phase, 10000000); - } - } - int p_col = ((prev_phase != 0) ? 224 : 96); - { - if (cta_rank == 0) { - mbarrier_wait(v_full_addr + (mma_v_stage) * 8, _phase_v_full); - int _mma_a_lo_6 = (((smem_p_fp8_addr) >> 4) & 0x3FFF) + (prev_phase) * 512; - int _mma_b_lo_6 = ((((smem_v_full_addr) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 138477584;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_6), "r"(_mma_b_lo_6), "r"((tmem_tmem_scratch + (256))), "r"(((first_pv) ? 0 : 1))); - int _mma_a_lo_7 = (((smem_p_fp8_addr) >> 4) & 0x3FFF) + (prev_phase) * 512; - int _mma_b_lo_7 = ((((smem_v_full_addr + 16384) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 138477584;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_7), "r"(_mma_b_lo_7), "r"((tmem_tmem_scratch + (384))), "r"(((first_pv) ? 0 : 1))); - if (cta_rank == 0) { - elect_commit_cg2_multicast(v_empty_addr + (mma_v_stage) * 8, (uint16_t)(3)); - } - } - mma_v_stage += 1; - if (mma_v_stage == 2) { mma_v_stage = 0; _phase_v_full ^= 1; } - } - first_pv = 0; - if (cta_rank == 0) { - elect_commit_cg2_multicast(pv_done_addr, (uint16_t)(3)); - } - } - } - int last_pipeline_tile = mma_tile_cursor + num_kv_tiles_3 - 1; - int last_phase = last_pipeline_tile & 1; - int drain_wait_phase = last_pipeline_tile >> 1 & 1; - mbarrier_wait(p_full_addr + (last_phase) * 8, drain_wait_phase); - mbarrier_wait(corr_done_addr + (last_phase) * 8, drain_wait_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - { - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cluster.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(pv_pair_ready_addr + (unsigned int)(last_phase * 8)), "r"(0) : "memory"); - if (cta_rank == 0) { - mbarrier_wait_cluster_hint(pv_pair_ready_addr + (last_phase) * 8, drain_wait_phase, 10000000); - } - } - int p_col_last = ((last_phase != 0) ? 224 : 96); - { - if (cta_rank == 0) { - mbarrier_wait(v_full_addr + (mma_v_stage) * 8, _phase_v_full); - int _mma_a_lo_10 = (((smem_p_fp8_addr) >> 4) & 0x3FFF) + (last_phase) * 512; - int _mma_b_lo_10 = ((((smem_v_full_addr) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 138477584;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_10), "r"(_mma_b_lo_10), "r"((tmem_tmem_scratch + (256))), "r"(((first_pv) ? 0 : 1))); - int _mma_a_lo_11 = (((smem_p_fp8_addr) >> 4) & 0x3FFF) + (last_phase) * 512; - int _mma_b_lo_11 = ((((smem_v_full_addr + 16384) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 138477584;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_11), "r"(_mma_b_lo_11), "r"((tmem_tmem_scratch + (384))), "r"(((first_pv) ? 0 : 1))); - if (cta_rank == 0) { - elect_commit_cg2_multicast(v_empty_addr + (mma_v_stage) * 8, (uint16_t)(3)); - } - } - mma_v_stage += 1; - if (mma_v_stage == 2) { mma_v_stage = 0; _phase_v_full ^= 1; } - } - if (cta_rank == 0) { - elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); - } - if (cta_rank == 0) { - elect_commit_cg2_multicast(o_done_addr, (uint16_t)(3)); - } - mma_tile_cursor = mma_tile_cursor + num_kv_tiles_3; - } - mbarrier_arrive(tmem_dealloc_addr); - unsigned int _phase_tmem_dealloc_0 = 0; - mbarrier_wait(tmem_dealloc_addr, _phase_tmem_dealloc_0); - _phase_tmem_dealloc_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - { - int tmem_dealloc_peer_rank = bid % 2 ^ 1; - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cluster.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(tmem_dealloc_peer_addr), "r"(tmem_dealloc_peer_rank) : "memory"); - mbarrier_wait(tmem_dealloc_peer_addr, 0); - int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); - asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); - } - } - } - // ---- Role: empty1 ---- - if (warp >= 10 && warp <= 11) { - { // empty1_main - const int wg2_dummy_2 = 0; - unsigned int _phase_q_full_0_4 = 0; - } - } - - // Cleanup -} - -} // extern "C" diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_a2be51a1fd34ec0d9a3b_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_a2be51a1fd34ec0d9a3b_binding.cu deleted file mode 100644 index 035179e18f8..00000000000 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_a2be51a1fd34ec0d9a3b_binding.cu +++ /dev/null @@ -1,525 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -// Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_a2be51a1fd34ec0d9a3b'. -#include -#include -#include -#include - -#include "tvm_ffi_utils.h" - -#include -#include -#include -#include -#include -#include -#include - -struct CakeTensorMap; -extern "C" __global__ void kernel_cake_dsv4_a2be51a1fd34ec0d9a3b(CakeTensorMap const* tmap_q, CakeTensorMap const* tmap_swa_kv, CakeTensorMap const* tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items); - - -namespace cake_host_shim_5841b4c9e46c06fc { - -using tvm::ffi::Optional; -using tvm::ffi::TensorView; - -class ScopedCudaDevice { - public: - explicit ScopedCudaDevice(int device_id) { - cudaError_t error = cudaGetDevice(&previous_device_); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed before host-shim launch: cudaError=" - << static_cast(error); - if (previous_device_ != device_id) { - error = cudaSetDevice(device_id); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaSetDevice failed before host-shim launch for cuda:" - << device_id << ": cudaError=" << static_cast(error); - restore_ = true; - } - } - - ScopedCudaDevice(const ScopedCudaDevice&) = delete; - ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; - - ~ScopedCudaDevice() noexcept { - if (restore_) { - (void)cudaSetDevice(previous_device_); - } - } - - private: - int previous_device_ = -1; - bool restore_ = false; -}; - -inline int64_t CakeDeviceMultiprocessorCount(int device_id) { - constexpr int kMaxCachedCudaDevices = 64; - TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) - << "physical-SM-count cache does not cover cuda:" << device_id; - static std::atomic count_by_device[kMaxCachedCudaDevices]{}; - int cached = count_by_device[device_id].load(std::memory_order_acquire); - if (cached > 0) return cached; - - int count = 0; - cudaError_t error = cudaDeviceGetAttribute( - &count, cudaDevAttrMultiProcessorCount, device_id); - TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) - << "querying multiProcessorCount failed for cuda:" << device_id - << ": cudaError=" << static_cast(error) << ", count=" << count; - // Concurrent first launches may repeat the immutable device query, but all - // publication is atomic and every later hot-path lookup is one acquire load. - count_by_device[device_id].store(count, std::memory_order_release); - return count; -} - -inline void CheckCudaTensor(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) - << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; -} - -inline void CheckSameCudaDevice( - const TensorView& t, - const TensorView& reference, - const char* name, - const char* reference_name) { - TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) - << name << " must be on the same CUDA device as " << reference_name - << ": got cuda:" << t.device().device_id - << " versus cuda:" << reference.device().device_id; -} - -inline void CheckCurrentCudaDevice( - const TensorView& reference, - const char* reference_name) { - int current_device = -1; - cudaError_t error = cudaGetDevice(¤t_device); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed while validating " << reference_name - << ": cudaError=" << static_cast(error); - TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) - << "current CUDA device must match " << reference_name - << ": current=cuda:" << current_device - << ", tensor=cuda:" << reference.device().device_id; -} - -inline void CheckContiguous(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; -} - -inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { - DLDataType d = t.dtype(); - TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) - << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits - << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits - << ", lanes=" << (int)d.lanes << ")"; -} - -// A logical axis.outer(trailing) folds every source dim above the trailing -// dimensions. Shape products are independent of physical strides, so verify -// the leading dimensions form one dense row-major chain instead of inventing -// a "folded stride". The descriptor reads its exact adjacent physical step -// separately through stride[-(trailing + 1)]. -inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { - TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) - << name << " cannot fold leading dimensions above " << trailing - << " trailing dims from ndim=" << t.ndim(); - int outer_last = t.ndim() - trailing - 1; - if (outer_last <= 0) { - return; - } - int64_t step = t.stride(outer_last); - TVM_FFI_CHECK(step > 0, ValueError) - << name << " physical strides must be positive"; - int64_t expected = step; - for (int axis = outer_last - 1; axis >= 0; --axis) { - expected *= t.size(axis + 1); - if (t.size(axis) > 1) { - TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) - << name << " leading dims are not physically foldable above " << trailing - << " trailing dims: stride(" << axis << ")=" << t.stride(axis) - << ", expected " << expected; - } - } -} - -// Explicit immutable caller storage: initialize once, outside graph capture, -// exactly as the production TmaDeviceSlot path. The owner must retain the -// allocation and must never expose it as mutable scratch or rewrite its bytes. -template -static inline void PrepareImmutableCallerTmaWorkspace( - const tvm::ffi::Tensor& retained_workspace, - const CUtensorMap (&maps)[N], cudaStream_t stream) { - void* workspace = retained_workspace.data_ptr(); - CUcontext context = nullptr; - CUresult result = cuCtxGetCurrent(&context); - TVM_FFI_CHECK(result == CUDA_SUCCESS && context != nullptr, RuntimeError) - << "immutable TMA workspace requires an active CUDA context"; - CUcontext allocation_context = nullptr; - result = cuPointerGetAttribute( - &allocation_context, CU_POINTER_ATTRIBUTE_CONTEXT, - reinterpret_cast(workspace)); - TVM_FFI_CHECK(result == CUDA_SUCCESS && allocation_context == context, RuntimeError) - << "immutable TMA workspace must belong to the active CUDA context"; - unsigned long long buffer_id = 0; - result = cuPointerGetAttribute( - &buffer_id, CU_POINTER_ATTRIBUTE_BUFFER_ID, - reinterpret_cast(workspace)); - TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) - << "querying immutable TMA allocation identity failed"; - std::string key = std::to_string(reinterpret_cast(context)); - key += ":" + std::to_string(reinterpret_cast(workspace)); - key += ":" + std::to_string(buffer_id); - std::string descriptor_bytes(reinterpret_cast(maps), sizeof(maps)); - struct InitializedWorkspace { - tvm::ffi::Tensor owner; - std::string descriptor_bytes; - }; - static std::mutex mu; - // Keep the actual Tensor allocation alive: a caching allocator may reuse a - // freed suballocation without changing its enclosing CUDA buffer identity. - static auto* initialized = new std::unordered_map(); - std::lock_guard lock(mu); - auto it = initialized->find(key); - if (it != initialized->end()) { - TVM_FFI_CHECK(it->second.descriptor_bytes == descriptor_bytes, ValueError) - << "immutable TMA descriptors changed; provide new private descriptor storage"; - return; - } - CUstreamCaptureStatus capture_status = CU_STREAM_CAPTURE_STATUS_NONE; - result = cuStreamIsCapturing(reinterpret_cast(stream), &capture_status); - TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) - << "querying immutable TMA initialization capture state failed"; - TVM_FFI_CHECK(capture_status == CU_STREAM_CAPTURE_STATUS_NONE, RuntimeError) - << "immutable TMA storage must be initialized before CUDA Graph capture"; - // Synchronous host-to-device initialization completes before publishing the - // slot to any thread/stream, matching TmaDeviceSlot. The consumer retains its - // generated tensor-map acquire; the mutable upload's release fence is unchanged. - result = cuMemcpyHtoD(reinterpret_cast(workspace), maps, sizeof(maps)); - TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) - << "initializing immutable caller TMA descriptors failed"; - initialized->emplace( - std::move(key), InitializedWorkspace{retained_workspace, std::move(descriptor_bytes)}); -} - -static inline void* CallerTmaWorkspaceSlot( - const tvm::ffi::TensorView& workspace, size_t slot) { - return static_cast(workspace.data_ptr()) + slot * sizeof(CUtensorMap); -} - -// 4D TMA descriptor for buffer 'tmap_q' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 3, ValueError) - << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - int64_t d2 = t.size(t.ndim() - 2); - int64_t d3 = t.size(t.ndim() - 3); - TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) - << "TMA source 'tmap_q' trailing dims must be positive"; - TVM_FFI_CHECK(d1 % 128 == 0, ValueError) - << "TMA source 'tmap_q' extent " << d1 - << " must divide exactly by " << 128; - uint64_t global_dim[4] = {(uint64_t)(128), (uint64_t)(d2), (uint64_t)((d1 / 128)), (uint64_t)(d3)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; - TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) - << "TMA box (128, 128, 1, 1) exceeds resolved global dims for 'tmap_q'"; - int64_t carrier_stride_0 = d1; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; - int64_t carrier_stride_1 = 128; - TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; - TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; - TVM_FFI_CHECK((carrier_stride_1 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; - int64_t carrier_stride_2 = (d2 * d1); - TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; - TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; - TVM_FFI_CHECK((carrier_stride_2 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; - uint64_t global_strides[3] = { - (uint64_t)((carrier_stride_0 * 8) / 8), - (uint64_t)((carrier_stride_1 * 8) / 8), - (uint64_t)((carrier_stride_2 * 8) / 8), - }; - uint32_t box_dim[4] = {128u, 128u, 1u, 1u}; - uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_swa_kv' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_swa_kv' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; - TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (128, 1) exceeds resolved global dims for 'tmap_swa_kv'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 8) / 8), - }; - uint32_t box_dim[2] = {128u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_compressed_kv' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_compressed_kv(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_compressed_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_compressed_kv' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_compressed_kv' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_compressed_kv"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_compressed_kv' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved a non-positive global dim"; - TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (128, 1) exceeds resolved global dims for 'tmap_compressed_kv'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 8) / 8), - }; - uint32_t box_dim[2] = {128u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_kv') failed: CUresult=" << (int)r; - return tm; -} - -void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_compressed_kv, TensorView arg_O, TensorView arg_partial_lse, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_num_query_tokens, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, tvm::ffi::Tensor arg_tma_descriptor_workspace, int64_t grid_x, int64_t grid_y, int64_t grid_z) { - DLDevice dev = arg_tmap_q.device(); - ScopedCudaDevice device_guard(dev.device_id); - CheckCudaTensor(arg_tmap_q, "tmap_q"); - CheckDtype(arg_tmap_q, "tmap_q", 1, 8, 1); - CheckContiguous(arg_tmap_q, "tmap_q"); - CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); - CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 1, 8, 1); - CheckCudaTensor(arg_tmap_compressed_kv, "tmap_compressed_kv"); - CheckDtype(arg_tmap_compressed_kv, "tmap_compressed_kv", 1, 8, 1); - CheckCudaTensor(arg_O, "O"); - CheckDtype(arg_O, "O", 4, 16, 1); - CheckContiguous(arg_O, "O"); - CheckCudaTensor(arg_partial_lse, "partial_lse"); - CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); - CheckContiguous(arg_partial_lse, "partial_lse"); - CheckCudaTensor(arg_swa_indices, "swa_indices"); - CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); - CheckContiguous(arg_swa_indices, "swa_indices"); - CheckCudaTensor(arg_compressed_indices, "compressed_indices"); - CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); - CheckContiguous(arg_compressed_indices, "compressed_indices"); - CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); - CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckCudaTensor(arg_sinks, "sinks"); - CheckDtype(arg_sinks, "sinks", 2, 32, 1); - CheckContiguous(arg_sinks, "sinks"); - CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); - CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); - CheckContiguous(arg_bmm1_scale, "bmm1_scale"); - CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); - CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); - CheckContiguous(arg_bmm2_scale, "bmm2_scale"); - TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) - << "scalar 'num_heads' value " << arg_num_heads - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) - << "scalar 'swa_index_stride' value " << arg_swa_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) - << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) - << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_num_query_tokens >= -2147483648LL && arg_num_query_tokens <= 2147483647LL, ValueError) - << "scalar 'num_query_tokens' value " << arg_num_query_tokens - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) - << "scalar 'sparse_topk' value " << arg_sparse_topk - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) - << "scalar 'has_sinks' value " << arg_has_sinks - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) - << "scalar 'total_work_items' value " << arg_total_work_items - << " is outside i32 range [-2147483648, 2147483647]"; - CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); - CheckSameCudaDevice(arg_tmap_compressed_kv, arg_tmap_q, "tmap_compressed_kv", "tmap_q"); - CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); - CheckSameCudaDevice(arg_partial_lse, arg_tmap_q, "partial_lse", "tmap_q"); - CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); - CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); - CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); - CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); - CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); - CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); - CheckCudaTensor(arg_tma_descriptor_workspace, "tma_descriptor_workspace"); - CheckSameCudaDevice(arg_tma_descriptor_workspace, arg_tmap_q, "tma_descriptor_workspace", "tmap_q"); - CheckContiguous(arg_tma_descriptor_workspace, "tma_descriptor_workspace"); - CheckDtype(arg_tma_descriptor_workspace, "tma_descriptor_workspace", 1, 8, 1); - TVM_FFI_CHECK(arg_tma_descriptor_workspace.numel() >= 384, ValueError) - << "tma_descriptor_workspace requires at least 384 bytes"; - TVM_FFI_CHECK( - reinterpret_cast(arg_tma_descriptor_workspace.data_ptr()) % 128 == 0, - ValueError) - << "tma_descriptor_workspace must be 128-byte aligned"; - TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) - << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y - << ", " << grid_z << ")"; - - cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); - CUtensorMap h_tmap_q = EncodeTma_tmap_q(arg_tmap_q); - CakeTensorMap const* p_tmap_q = static_cast(CallerTmaWorkspaceSlot(arg_tma_descriptor_workspace, 0u)); - CUtensorMap h_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); - CakeTensorMap const* p_tmap_swa_kv = static_cast(CallerTmaWorkspaceSlot(arg_tma_descriptor_workspace, 1u)); - CUtensorMap h_tmap_compressed_kv = EncodeTma_tmap_compressed_kv(arg_tmap_compressed_kv); - CakeTensorMap const* p_tmap_compressed_kv = static_cast(CallerTmaWorkspaceSlot(arg_tma_descriptor_workspace, 2u)); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); - float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); - int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); - int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); - int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); - float* p_sinks = static_cast(arg_sinks.data_ptr()); - float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); - float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); - int32_t v_num_heads = (int32_t)arg_num_heads; - int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; - int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; - int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; - int32_t v_num_query_tokens = (int32_t)arg_num_query_tokens; - int32_t v_sparse_topk = (int32_t)arg_sparse_topk; - int32_t v_has_sinks = (int32_t)arg_has_sinks; - int32_t v_total_work_items = (int32_t)arg_total_work_items; - void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_tmap_compressed_kv, &p_O, &p_partial_lse, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_num_query_tokens, &v_sparse_topk, &v_has_sinks, &v_total_work_items}; - - static std::mutex smem_mu; - static auto* smem_status_by_device = - new std::unordered_map(); - cudaError_t smem_status = cudaSuccess; - { - std::lock_guard lock(smem_mu); - auto it = smem_status_by_device->find(dev.device_id); - if (it == smem_status_by_device->end()) { - int current_device = -1; - smem_status = cudaGetDevice(¤t_device); - if (smem_status == cudaSuccess && current_device != dev.device_id) { - smem_status = cudaErrorInvalidDevice; - } - if (smem_status == cudaSuccess) { - smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_a2be51a1fd34ec0d9a3b), - cudaFuncAttributeMaxDynamicSharedMemorySize, - 205824); - } - smem_status_by_device->emplace(dev.device_id, smem_status); - } else { - smem_status = it->second; - } - } - TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_a2be51a1fd34ec0d9a3b failed: " - << cudaGetErrorString(smem_status); - dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); - dim3 block(512u, 1u, 1u); - const CUtensorMap workspace_maps[] = {h_tmap_q, h_tmap_swa_kv, h_tmap_compressed_kv}; - PrepareImmutableCallerTmaWorkspace(arg_tma_descriptor_workspace, workspace_maps, stream); - - cudaLaunchConfig_t config{}; - config.gridDim = grid; - config.blockDim = block; - config.dynamicSmemBytes = 205824u; - config.stream = stream; - cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_a2be51a1fd34ec0d9a3b), kargs); - TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_a2be51a1fd34ec0d9a3b failed: " - << cudaGetErrorString(launch_status); -} - -} // namespace cake_host_shim_5841b4c9e46c06fc - -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_5841b4c9e46c06fc::Run); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_a2be51a1fd34ec0d9a3b_kernel.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_a2be51a1fd34ec0d9a3b_kernel.cu deleted file mode 100644 index 35598530ac2..00000000000 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_a2be51a1fd34ec0d9a3b_kernel.cu +++ /dev/null @@ -1,2400 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -typedef signed char int8_t; -typedef unsigned char uint8_t; -typedef unsigned short uint16_t; -typedef unsigned int uint32_t; -#if defined(__CUDACC_RTC__) -typedef unsigned long long uint64_t; -#else -typedef unsigned long uint64_t; -#endif -static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); -typedef signed int int32_t; -typedef short int int16_t; -struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; -struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; -static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); -static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); -template -struct __align__(128) CakeTensorMapPack { CakeTensorMap maps[N]; }; - -#if defined(__CUDACC_RTC__) -typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; -#else -#include -#endif - -static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); -static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); -#include -#include - -#define CAKE_INF CUDART_INF_F -#define TMEM_NCOLS 512 -#define TMEM_TMEM_SCRATCH_OFFSET 0 -#define NUM_K_PIPE_STAGES 1 -#define NUM_V_PIPE_STAGES 1 -#define NUM_KV_PIPE_STAGES 8 -#define NUM_INDEX_PIPE_STAGES 1 -#define SMEM_SMEM_Q_OFF 1024 -#define SMEM_SMEM_Q_STAGE_BYTES 16384 -#define SMEM_SMEM_Q_STRIDE 16384 -#define SMEM_SMEM_KV_OFF 66560 -#define SMEM_SMEM_KV_STAGE_BYTES 16384 -#define SMEM_SMEM_KV_STRIDE 16384 -#define SMEM_SMEM_V_OFF 66560 -#define SMEM_SMEM_V_STAGE_BYTES 16384 -#define SMEM_SMEM_V_STRIDE 16384 -#define SMEM_SMEM_Q_FULL_OFF 1024 -#define SMEM_SMEM_Q_FULL_STAGE_BYTES 32768 -#define SMEM_SMEM_Q_FULL_STRIDE 32768 -#define SMEM_SMEM_K_FULL_OFF 66560 -#define SMEM_SMEM_K_FULL_STAGE_BYTES 32768 -#define SMEM_SMEM_K_FULL_STRIDE 32768 -#define SMEM_SMEM_V_FULL_OFF 132096 -#define SMEM_SMEM_V_FULL_STAGE_BYTES 32768 -#define SMEM_SMEM_V_FULL_STRIDE 32768 -#define SMEM_SMEM_STATS_MAX_OFF 197632 -#define SMEM_SMEM_STATS_MAX_STAGE_BYTES 1024 -#define SMEM_SMEM_STATS_MAX_STRIDE 1024 -#define SMEM_SMEM_STATS_SUM_OFF 198656 -#define SMEM_SMEM_STATS_SUM_STAGE_BYTES 512 -#define SMEM_SMEM_STATS_SUM_STRIDE 512 -#define SMEM_SMEM_STATS_FINAL_MAX_OFF 199168 -#define SMEM_SMEM_STATS_FINAL_MAX_STAGE_BYTES 512 -#define SMEM_SMEM_STATS_FINAL_MAX_STRIDE 512 -#define SMEM_SMEM_SOFTMAX_WARP_PAIR_EXCHANGE_OFF 199680 -#define SMEM_SMEM_SOFTMAX_WARP_PAIR_EXCHANGE_STAGE_BYTES 1024 -#define SMEM_SMEM_SOFTMAX_WARP_PAIR_EXCHANGE_STRIDE 1024 -#define SMEM_SMEM_CORR_WARP_PAIR_EXCHANGE_OFF 200704 -#define SMEM_SMEM_CORR_WARP_PAIR_EXCHANGE_STAGE_BYTES 512 -#define SMEM_SMEM_CORR_WARP_PAIR_EXCHANGE_STRIDE 512 -#define SMEM_SMEM_SPARSE_INDICES_OFF 201216 -#define SMEM_SMEM_SPARSE_INDICES_STAGE_BYTES 4608 -#define SMEM_SMEM_SPARSE_INDICES_STRIDE 4608 -#define SMEM_SMEM_P_FP8_OFF 1024 -#define SMEM_SMEM_P_FP8_STAGE_BYTES 8192 -#define SMEM_SMEM_P_FP8_STRIDE 8192 -#define SMEM_TOTAL 205824 -#define THREADS 512 - -#include - -__device__ __forceinline__ uint32_t elect_sync() { - uint32_t pred = 0; - asm volatile( - "{\n\t" - ".reg .pred %%px;\n\t" - "elect.sync _|%%px, %1;\n\t" - "@%%px mov.s32 %0, 1;\n\t" - "}\n" - : "+r"(pred) - : "r"(0xFFFFFFFF)); - return pred; -} - - -__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { - asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" - :: "r"(mbar_addr), "r"(count) : "memory"); -} - -__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { - asm volatile("mbarrier.init.b64 [%0], %1;" - :: "l"(mbar_addr), "r"(count) : "memory"); -} - - -__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - - -__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { - uint32_t ticks = 0x989680; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE;\n\t" - "bra.uni LAB_WAIT;\n\t" - "DONE:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); -} - -// Source-faithful relaxed CTA wait used only by a typed protocol that does -// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. -__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, 10000000;\n\t" - "@P1 bra.uni DONE_RELAXED;\n\t" - "bra.uni LAB_WAIT_RELAXED;\n\t" - "DONE_RELAXED:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Exact source ports may request the PTX suspendTimeHint operand explicitly. -// The hint is expressed in nanoseconds and is kept separate from the canonical -// no-hint CTA helper so unrelated schedules retain their existing retry path. -__device__ __forceinline__ void mbarrier_wait_suspend( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_SUSPEND:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_SUSPEND;\n\t" - "bra.uni LAB_WAIT_SUSPEND;\n\t" - "DONE_SUSPEND:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE_CLUSTER;\n\t" - "bra.uni LAB_WAIT_CLUSTER;\n\t" - "DONE_CLUSTER:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - ".reg .u32 WAIT_ADDR;\n\t" - "mov.u32 WAIT_ADDR, %0;\n\t" - "LAB_WAIT_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [WAIT_ADDR], %1, %2;\n\t" - "@P1 bra.uni DONE_HINT;\n\t" - "bra.uni LAB_WAIT_HINT;\n\t" - "DONE_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -// Exact unqualified CTA wait used by source schedules whose PTX intentionally -// omits the acquire qualifier while retaining a typed suspendTimeHint operand. -__device__ __forceinline__ void mbarrier_wait_relaxed_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED_HINT:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra DONE_RELAXED_HINT;\n\t" - "bra LAB_WAIT_RELAXED_HINT;\n\t" - "DONE_RELAXED_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); -} - -__device__ __forceinline__ void mbarrier_wait_cluster_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" - "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" - "DONE_CLUSTER_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_suspend( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait_cluster(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); - } -} - - -__device__ __forceinline__ void tcgen05_mma_f8f6f4( - int taddr, uint64_t a_desc, uint64_t b_desc, - uint32_t i_desc, int enable_input_d) { - asm volatile( - "{\n\t" - ".reg .pred p;\n\t" - "setp.ne.b32 p, %4, 0;\n\t" - "tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], %1, %2, %3, p;\n\t" - "}\n" - :: "r"(taddr), "l"(a_desc), "l"(b_desc), - "r"(i_desc), "r"(enable_input_d) - : "memory"); -} - - -__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { - return (x & 0x3FFFFULL) >> 4ULL; -} - - -union MmaSmemDesc { - uint64_t u64; - uint32_t u32[2]; -}; - -__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { - MmaSmemDesc tmp; - tmp.u64 = smem_desc; - tmp.u32[0] += offset; - smem_desc = tmp.u64; -} - - -__device__ __forceinline__ void elect_commit(int mbar_addr) { - asm volatile( - "{\n\t" - ".reg .pred leader;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" - ".shared::cluster.b64 [%0];\n\t" - "}\n" - :: "r"(mbar_addr)); -} - - -__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { - asm volatile( - "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" - :: "r"(mbar_addr), "r"(bytes) : "memory"); -} - - -__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7," - " %8, %9, %10, %11, %12, %13, %14, %15," - " %16, %17, %18, %19, %20, %21, %22, %23," - " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), - "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), - "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), - "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), - "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), - "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), - "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x16.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7," - " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), - "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), - "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ void tmem_ld_x4(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x4.b32" - " {%0, %1, %2, %3}, [%4];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { - asm volatile( - "tcgen05.st.sync.aligned.32x32b.x16.b32" - " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," - " %9, %10, %11, %12, %13, %14, %15, %16};" - :: "r"(tmem_addr), - "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), - "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), - "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), - "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); -} - - -__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { - asm volatile( - "tcgen05.st.sync.aligned.32x32b.x4.b32" - " [%0], {%1, %2, %3, %4};" - :: "r"(tmem_addr), - "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); -} - - -__device__ __forceinline__ float approx_exp2(float x) { - float y; - asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float approx_rcp(float x) { - float y; - asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float max_noftz(float a, float b) { - float c; - asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); - return c; -} - - -__device__ __forceinline__ float warp_reduce_max(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); - return val; -} - - -__device__ __forceinline__ float warp_reduce_sum(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val += __shfl_xor_sync(0xFFFFFFFF, val, offset); - return val; -} - - -__device__ __forceinline__ float row_max_reduce(float2 acc) { - return max_noftz(acc.x, acc.y); -} - - -__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { - #pragma unroll - for (int j = 0; j < 16; j++) { - if (j % 2 == 0) - acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); - else - acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); - } -} - - -__device__ __forceinline__ float2 ex2_emulation_f32x2_value(float2 value) { - const float c0 = 1.0f, c1 = 0.695146143436431884765625f; - const float c2 = 0.227564394474029541015625f, c3 = 0.077119089663028717041015625f; - const float magic = 12582912.0f; - float x0 = max_noftz(value.x, -127.0f), x1 = max_noftz(value.y, -127.0f); - float2 xc2 = make_float2(x0, x1), magic2 = make_float2(magic, magic); - float2 xr2; - asm("add.rm.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xr2) - : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&magic2)); - float2 c3_2 = make_float2(c3, c3), c2_2 = make_float2(c2, c2); - float2 c1_2 = make_float2(c1, c1), c0_2 = make_float2(c0, c0); - float2 xrb2, xfrac2; - asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xrb2) - : "l"(*(unsigned long long*)&xr2), "l"(*(unsigned long long*)&magic2)); - asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xfrac2) - : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&xrb2)); - float2 poly2; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&c3_2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c2_2)); - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c1_2)); - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c0_2)); - int x0r_i, x1r_i, p0_i, p1_i; - asm("mov.b64 {%0, %1}, %2;" : "=r"(x0r_i), "=r"(x1r_i) : "l"(*(unsigned long long*)&xr2)); - asm("mov.b64 {%0, %1}, %2;" : "=r"(p0_i), "=r"(p1_i) : "l"(*(unsigned long long*)&poly2)); - float r0, r1; - asm("mov.b32 %0, %1;" : "=f"(r0) : "r"((x0r_i << 23) + p0_i)); - asm("mov.b32 %0, %1;" : "=f"(r1) : "r"((x1r_i << 23) + p1_i)); - return make_float2(r0, r1); -} - -__device__ __forceinline__ void ex2_emulation_f32x2(float* x0_ptr, float* x1_ptr) { - float2 result = ex2_emulation_f32x2_value(make_float2(*x0_ptr, *x1_ptr)); - *x0_ptr = result.x; *x1_ptr = result.y; -} - -__device__ __forceinline__ void softmax_frag_exp2_cast( - float* sv, uint32_t* pv, int use_emu) -{ - #pragma unroll - for (int j = 0; j < 16; j++) { - if (use_emu && j >= 12) - ex2_emulation_f32x2(&sv[j*2], &sv[j*2+1]); - else { - sv[j*2] = approx_exp2(sv[j*2]); - sv[j*2+1] = approx_exp2(sv[j*2+1]); - } - } - #pragma unroll - for (int j = 0; j < 16; j++) { - __nv_bfloat162 bf = __float22bfloat162_rn({sv[j*2], sv[j*2+1]}); - pv[j] = reinterpret_cast(bf); - } -} - - - -__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { - const float2* sv2 = reinterpret_cast(sv); - #pragma unroll - for (int j = 0; j < 16; j++) { - asm("add.f32x2 %0, %1, %2;" - : "+l"(reinterpret_cast(*acc)) - : "l"(reinterpret_cast(*acc)), - "l"(reinterpret_cast(sv2[j]))); - } -} - - -__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { - asm("mul.rn.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { - asm("mul.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { - asm("add.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { - asm("add.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { - asm("sub.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { - asm("sub.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("add.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { - float2 r; - asm("sub.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("sub.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ void fma_scale_x32( - float* sv, const float2* scale2, const float2* neg_max2) -{ - float2* sv_2 = reinterpret_cast(sv); - #pragma unroll - for (int j = 0; j < 16; j++) - fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); -} - -__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("mul.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - - -// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) - -__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("add.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("add.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("add.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("add.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("add.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("add.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("add.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - - -__device__ __forceinline__ void fence_async_shared() { - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); -} - - -__device__ __forceinline__ uint64_t make_smem_desc(int addr) { - const int SBO = 1024; - return desc_encode(addr) - | (desc_encode(SBO) << 32ULL) - | (1ULL << 46ULL) - | (2ULL << 61ULL); -} - - -__device__ __forceinline__ void tma_2d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3}], [%4];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_4d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.4d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5}], [%6];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr) { - // Canonical .shared::cta form for non-multicast gather4, matching - // trtllm-gen / cuda_ptx and the PTX ISA qualifier order - // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, - // .shared::cluster is reserved for the multicast variant (ctaMask). - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem_cta2( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr) { - // Canonical .shared::cta form; see tma_gather4_gmem2smem above. - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" - ".mbarrier::complete_tx::bytes.cta_group::2" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem_mc( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr, unsigned short cta_mask) { - // Multicast variant: the PTX grammar ties the .shared::cluster - // destination to .multicast::cluster + ctaMask (cf. cuda_ptx / - // SM100_TMA_LOAD_MULTICAST_2D_GATHER4 in CUTLASS). - asm volatile( - "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" - ".mbarrier::complete_tx::bytes.multicast::cluster" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr), "h"(cta_mask) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem_mc_cta2( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr, unsigned short cta_mask) { - // Multicast + cta_group::2 variant; see tma_gather4_gmem2smem_mc. - asm volatile( - "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" - ".mbarrier::complete_tx::bytes.multicast::cluster.cta_group::2" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr), "h"(cta_mask) : "memory"); -} - - -__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { - asm volatile( - "tcgen05.commit.cta_group::1.mbarrier::arrive::one" - ".shared::cluster.b64 [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tmem_st_x8_u32(int addr, const uint32_t* src) { - asm volatile( - "tcgen05.st.sync.aligned.32x32b.x8.b32" - " [%0], {%1,%2,%3,%4,%5,%6,%7,%8};" - :: "r"(addr), - "r"(src[0]), "r"(src[1]), "r"(src[2]), "r"(src[3]), - "r"(src[4]), "r"(src[5]), "r"(src[6]), "r"(src[7])); -} - - -__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { - uint32_t result; - asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" - : "=r"(result) : "r"(val)); - return result; -} - -extern "C" { - -__global__ __launch_bounds__(512, 1) void -kernel_cake_dsv4_a2be51a1fd34ec0d9a3b(CakeTensorMap const* tmap_q, CakeTensorMap const* tmap_swa_kv, CakeTensorMap const* tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items) -{ - const int tid = threadIdx.x; - const int warp = make_warp_uniform(tid / 32); - const int lane = tid % 32; - - extern __shared__ __align__(1024) char smem_raw[]; - int smem; - smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); - - const int mbar_base = smem; - #define q_full_addr (mbar_base + 0) - #define q_empty_addr (mbar_base + 8) - #define k_full_addr (mbar_base + 16) - #define k_empty_addr (mbar_base + 24) - #define v_full_addr (mbar_base + 32) - #define v_empty_addr (mbar_base + 40) - #define kv_full_addr (mbar_base + 48) - #define kv_empty_addr (mbar_base + 112) - #define index_full_addr (mbar_base + 176) - #define index_empty_addr (mbar_base + 184) - #define s_full_addr (mbar_base + 192) - #define p_full_addr (mbar_base + 208) - #define corr_done_addr (mbar_base + 224) - #define stats_addr (mbar_base + 240) - #define sum_ready_addr (mbar_base + 256) - #define o_done_addr (mbar_base + 264) - #define pv_done_addr (mbar_base + 272) - #define s_seeded_addr (mbar_base + 280) - #define q_pair_ready_addr (mbar_base + 288) - #define kv_pair_ready_addr (mbar_base + 296) - #define pv_pair_ready_addr (mbar_base + 360) - #define tmem_dealloc_addr (mbar_base + 376) - #define tmem_dealloc_peer_addr (mbar_base + 384) - - const int bid = blockIdx.x; - const int num_bids = gridDim.x; - - const int cta_rank = 0; - if (tid == 0) { - asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_q)) : "memory"); - asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_swa_kv)) : "memory"); - asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_compressed_kv)) : "memory"); - } - __syncthreads(); - - - // Kernel setup ops - uint8_t* smem_q = reinterpret_cast(smem_raw + 1024); - const int smem_q_addr = smem + 1024; - uint8_t* smem_kv = reinterpret_cast(smem_raw + 66560); - const int smem_kv_addr = smem + 66560; - uint8_t* smem_v = reinterpret_cast(smem_raw + 66560); - const int smem_v_addr = smem + 66560; - uint8_t* smem_q_full = reinterpret_cast(smem_raw + 1024); - const int smem_q_full_addr = smem + 1024; - uint8_t* smem_k_full = reinterpret_cast(smem_raw + 66560); - const int smem_k_full_addr = smem + 66560; - uint8_t* smem_v_full = reinterpret_cast(smem_raw + 132096); - const int smem_v_full_addr = smem + 132096; - float* smem_stats_max = reinterpret_cast(smem_raw + 197632); - const int smem_stats_max_addr = smem + 197632; - float* smem_stats_sum = reinterpret_cast(smem_raw + 198656); - const int smem_stats_sum_addr = smem + 198656; - float* smem_stats_final_max = reinterpret_cast(smem_raw + 199168); - const int smem_stats_final_max_addr = smem + 199168; - float* smem_softmax_warp_pair_exchange = reinterpret_cast(smem_raw + 199680); - const int smem_softmax_warp_pair_exchange_addr = smem + 199680; - float* smem_corr_warp_pair_exchange = reinterpret_cast(smem_raw + 200704); - const int smem_corr_warp_pair_exchange_addr = smem + 200704; - int* smem_sparse_indices = reinterpret_cast(smem_raw + 201216); - const int smem_sparse_indices_addr = smem + 201216; - uint8_t* smem_p_fp8 = reinterpret_cast(smem_raw + 1024); - const int smem_p_fp8_addr = smem + 1024; - - // Mbarrier init (23 pipeline groups, 0 ordered-sequence groups, 49 barriers) - // Mbarriers at smem_raw[0..392) - - if (warp == 0) { - uint32_t leader = elect_sync(); - if (leader) { - // q_full: 1 barriers, init_count=1 - mbarrier_init(smem + 0, 1); - // q_empty: 1 barriers, init_count=1 - mbarrier_init(smem + 8, 1); - // --- pipeline 'k_pipe' --- - // k_full: 1 barriers, init_count=1 - mbarrier_init(smem + 16, 1); - // k_empty: 1 barriers, init_count=1 - mbarrier_init(smem + 24, 1); - // --- pipeline 'v_pipe' --- - // v_full: 1 barriers, init_count=1 - mbarrier_init(smem + 32, 1); - // v_empty: 1 barriers, init_count=1 - mbarrier_init(smem + 40, 1); - // --- pipeline 'kv_pipe' --- - // kv_full: 8 barriers, init_count=1 - mbarrier_init(smem + 48, 1); - mbarrier_init(smem + 56, 1); - mbarrier_init(smem + 64, 1); - mbarrier_init(smem + 72, 1); - mbarrier_init(smem + 80, 1); - mbarrier_init(smem + 88, 1); - mbarrier_init(smem + 96, 1); - mbarrier_init(smem + 104, 1); - // kv_empty: 8 barriers, init_count=1 - mbarrier_init(smem + 112, 1); - mbarrier_init(smem + 120, 1); - mbarrier_init(smem + 128, 1); - mbarrier_init(smem + 136, 1); - mbarrier_init(smem + 144, 1); - mbarrier_init(smem + 152, 1); - mbarrier_init(smem + 160, 1); - mbarrier_init(smem + 168, 1); - // --- pipeline 'index_pipe' --- - // index_full: 1 barriers, init_count=1 - mbarrier_init(smem + 176, 1); - // index_empty: 1 barriers, init_count=1 - mbarrier_init(smem + 184, 1); - // s_full: 2 barriers, init_count=1 - mbarrier_init(smem + 192, 1); - mbarrier_init(smem + 200, 1); - // p_full: 2 barriers, init_count=128 - mbarrier_init(smem + 208, 128); - mbarrier_init(smem + 216, 128); - // corr_done: 2 barriers, init_count=128 - mbarrier_init(smem + 224, 128); - mbarrier_init(smem + 232, 128); - // stats: 2 barriers, init_count=128 - mbarrier_init(smem + 240, 128); - mbarrier_init(smem + 248, 128); - // sum_ready: 1 barriers, init_count=128 - mbarrier_init(smem + 256, 128); - // o_done: 1 barriers, init_count=1 - mbarrier_init(smem + 264, 1); - // pv_done: 1 barriers, init_count=1 - mbarrier_init(smem + 272, 1); - // s_seeded: 1 barriers, init_count=128 - mbarrier_init(smem + 280, 128); - // q_pair_ready: 1 barriers, init_count=64 - mbarrier_init(smem + 288, 64); - // --- pipeline 'kv_pipe' --- - // kv_pair_ready: 8 barriers, init_count=64 - mbarrier_init(smem + 296, 64); - mbarrier_init(smem + 304, 64); - mbarrier_init(smem + 312, 64); - mbarrier_init(smem + 320, 64); - mbarrier_init(smem + 328, 64); - mbarrier_init(smem + 336, 64); - mbarrier_init(smem + 344, 64); - mbarrier_init(smem + 352, 64); - // pv_pair_ready: 2 barriers, init_count=64 - mbarrier_init(smem + 360, 64); - mbarrier_init(smem + 368, 64); - // tmem_dealloc: 1 barriers, init_count=416 - mbarrier_init(smem + 376, 416); - // tmem_dealloc_peer: 1 barriers, init_count=32 - mbarrier_init(smem + 384, 32); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - - __syncwarp(); - - // TMEM alloc (512 columns, 512 used) - volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 392); - if (warp == 0) { - int _tmem_hold = smem + 392; - asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); - __syncwarp(); - asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); - } - - __syncthreads(); - asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); - - const int taddr = tmem_addr_storage[0]; - - // Kernel post-init ops - const int tmem_tmem_scratch = taddr; - - // ---- Ordered hardware-WG register redistribution ---- - // Dec phase frees registers before any WG attempts inc. - if (warp >= 8 && warp <= 15) { - asm volatile("setmaxnreg.dec.sync.aligned.u32 64;"); - } - - // ---- Role: index_warp ---- - if (warp == 13) { - { // index_warp_main - const int index_dummy = 0; - unsigned int _phase_index_empty = 1; - unsigned int _phase_q_full_0 = 0; - { - #pragma unroll 1 - for (unsigned int work_idx = bid; work_idx < total_work_items; work_idx += num_bids) { - mbarrier_wait(q_full_addr, _phase_q_full_0); - _phase_q_full_0 ^= 1; - } - } - } - } - // ---- Role: load_warp ---- - if (warp >= 9 && warp <= 12) { - { // load_warp_main - const int wg2_dummy = 0; - const int load_warp_rank = warp - ((0) ? 12 : 9); - int all_num_kv_tiles = (sparse_topk + 128 - 1) / 128; - unsigned int load_k_stage = 0; - unsigned int load_v_stage = 0; - unsigned int load_kv_stage = 0; - unsigned int load_k_index_stage = 0; - unsigned int load_v_index_stage = 0; - int k_cta_offset = 0; - unsigned int _phase_q_empty_0 = 1; - unsigned int _phase_index_full = 0; - unsigned int _phase_k_empty = 1; - unsigned int _phase_kv_empty = 1; - unsigned int _phase_v_empty = 1; - #pragma unroll 1 - for (unsigned int work_idx_1 = bid; work_idx_1 < total_work_items; work_idx_1 += num_bids) { - int split_idx = 0; - int query_idx = work_idx_1 >> 1; - int v_chunk = work_idx_1 & 1; - int tiles_per_split = all_num_kv_tiles; - int first_tile = split_idx * tiles_per_split; - int num_kv_tiles = tiles_per_split; - int sparse_extent = num_kv_tiles * 128; - mbarrier_wait(q_empty_addr, _phase_q_empty_0); - _phase_q_empty_0 ^= 1; - if (load_warp_rank == 0) { - if (elect_sync()) { - mbarrier_arrive_expect_tx(q_full_addr, 65536); - #pragma unroll - for (int q_stage = 0; q_stage < 4; q_stage++) { - tma_4d_gmem2smem(smem_q_addr + (unsigned int)(q_stage * 16384), tmap_q, 0, 0, q_stage, query_idx, q_full_addr); - } - } - } - { - int num_index_passes = sparse_extent / 128; - #pragma unroll 1 - for (int index_pass = load_warp_rank; index_pass < num_index_passes; index_pass += 4) { - int index_offset = index_pass * 128 + lane * 4; - int tile_col_base = (first_tile + index_pass) * 128; - int* row_tile_ptr = ((tile_col_base < 128) ? swa_indices + (query_idx * swa_index_stride) : compressed_indices + (query_idx * compressed_index_stride + (tile_col_base - 128))); - int lane_col = tile_col_base + lane * 4; - int sparse_rows[4]; - #pragma unroll - for (int row_i = 0; row_i < 4; row_i++) { - sparse_rows[row_i] = -1; - } - if (lane_col + 4 <= sparse_topk) { - int _vec_load_0[4]; - { - const int4* _ivptr_0 = reinterpret_cast(row_tile_ptr + (lane * 4) + 0); - int4 _ivld_0; - _ivld_0 = *_ivptr_0; - _vec_load_0[0 + 0] = _ivld_0.x; - _vec_load_0[0 + 1] = _ivld_0.y; - _vec_load_0[0 + 2] = _ivld_0.z; - _vec_load_0[0 + 3] = _ivld_0.w; - } - #pragma unroll - for (int row_i_1 = 0; row_i_1 < 4; row_i_1++) { - sparse_rows[row_i_1] = _vec_load_0[row_i_1]; - } - } else { - #pragma unroll - for (int row_i_2 = 0; row_i_2 < 4; row_i_2++) { - if (lane_col + row_i_2 < sparse_topk) { - sparse_rows[row_i_2] = row_tile_ptr[lane * 4 + row_i_2]; - } - } - } - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(smem_sparse_indices_addr + (unsigned int)(index_offset * 4)), "r"(sparse_rows[0]), "r"(sparse_rows[1]), "r"(sparse_rows[2]), "r"(sparse_rows[3]) : "memory"); - } - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - asm volatile("barrier.sync 8, 128;" ::: "memory"); - } - int k_index_stage_base = smem_sparse_indices_addr; - { - #pragma unroll - for (int qk_stage = 0; qk_stage < 4; qk_stage++) { - if (qk_stage == load_warp_rank) { - mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (load_kv_stage) * 8, 16384); - } - int k_dst = smem_kv_addr + load_kv_stage * 16384; - int group = lane; - if (group < 32) { - int group_offset = k_cta_offset + group * 4; - int raw_rows[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows[0])), "=r"(*reinterpret_cast(&raw_rows[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows[(0) + 3])) - : "r"(k_index_stage_base + group_offset * 4)); - int row0 = ((raw_rows[0] >= 0) ? raw_rows[0] : 0); - int row1 = ((raw_rows[1] >= 0) ? raw_rows[1] : 0); - int row2 = ((raw_rows[2] >= 0) ? raw_rows[2] : 0); - int row3 = ((raw_rows[3] >= 0) ? raw_rows[3] : 0); - if (first_tile == 0) { - tma_gather4_gmem2smem(k_dst + group * 512, tmap_swa_kv, qk_stage * 128, row0, row1, row2, row3, kv_full_addr + (load_kv_stage) * 8); - } else { - tma_gather4_gmem2smem(k_dst + group * 512, tmap_compressed_kv, qk_stage * 128, row0, row1, row2, row3, kv_full_addr + (load_kv_stage) * 8); - } - } - } - load_kv_stage += 1; - if (load_kv_stage == 8) { load_kv_stage = 0; _phase_kv_empty ^= 1; } - } - } - #pragma unroll 1 - for (int tile = 1; tile < num_kv_tiles; tile++) { - if ((0 & (int)((tile & 1) == 0)) != 0) { - mbarrier_wait(index_full_addr + (load_k_index_stage) * 8, _phase_index_full); - asm volatile("tcgen05.fence::after_thread_sync;"); - } - k_index_stage_base = smem_sparse_indices_addr; - { - #pragma unroll - for (int qk_stage_1 = 0; qk_stage_1 < 4; qk_stage_1++) { - if (qk_stage_1 == load_warp_rank) { - mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (load_kv_stage) * 8, 16384); - } - int k_dst_1 = smem_kv_addr + load_kv_stage * 16384; - int group_1 = lane; - if (group_1 < 32) { - int group_offset_1 = tile * 128 + k_cta_offset + group_1 * 4; - int raw_rows_1[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_1[0])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 3])) - : "r"(k_index_stage_base + group_offset_1 * 4)); - int row0_1 = ((raw_rows_1[0] >= 0) ? raw_rows_1[0] : 0); - int row1_1 = ((raw_rows_1[1] >= 0) ? raw_rows_1[1] : 0); - int row2_1 = ((raw_rows_1[2] >= 0) ? raw_rows_1[2] : 0); - int row3_1 = ((raw_rows_1[3] >= 0) ? raw_rows_1[3] : 0); - tma_gather4_gmem2smem(k_dst_1 + group_1 * 512, tmap_compressed_kv, qk_stage_1 * 128, row0_1, row1_1, row2_1, row3_1, kv_full_addr + (load_kv_stage) * 8); - } - } - load_kv_stage += 1; - if (load_kv_stage == 8) { load_kv_stage = 0; _phase_kv_empty ^= 1; } - } - } - int prev_tile = tile - 1; - int global_prev_tile = first_tile + prev_tile; - int v_index_stage_base = smem_sparse_indices_addr; - { - #pragma unroll - for (int pv_stage = 0; pv_stage < 2; pv_stage++) { - if (pv_stage == load_warp_rank) { - mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (load_kv_stage) * 8, 16384); - } - int v_dst = smem_kv_addr + load_kv_stage * 16384; - int v_col = v_chunk * 256 + pv_stage * 128; - int group_2 = lane; - int group_offset_2 = prev_tile * 128 + group_2 * 4; - int raw_rows_2[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_2[0])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 3])) - : "r"(v_index_stage_base + group_offset_2 * 4)); - int row0_2 = ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : 0); - int row1_2 = ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : 0); - int row2_2 = ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : 0); - int row3_2 = ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : 0); - if (global_prev_tile == 0) { - tma_gather4_gmem2smem(v_dst + group_2 * 512, tmap_swa_kv, v_col, row0_2, row1_2, row2_2, row3_2, kv_full_addr + (load_kv_stage) * 8); - } else { - tma_gather4_gmem2smem(v_dst + group_2 * 512, tmap_compressed_kv, v_col, row0_2, row1_2, row2_2, row3_2, kv_full_addr + (load_kv_stage) * 8); - } - } - load_kv_stage += 1; - if (load_kv_stage == 8) { load_kv_stage = 0; _phase_kv_empty ^= 1; } - } - } - } - int last_tile = num_kv_tiles - 1; - int global_last_tile = first_tile + last_tile; - { - #pragma unroll - for (int pv_stage_1 = 0; pv_stage_1 < 2; pv_stage_1++) { - if (pv_stage_1 == load_warp_rank) { - mbarrier_wait(kv_empty_addr + (load_kv_stage) * 8, _phase_kv_empty); - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (load_kv_stage) * 8, 16384); - } - int v_dst_1 = smem_kv_addr + load_kv_stage * 16384; - int v_col_1 = v_chunk * 256 + pv_stage_1 * 128; - int group_3 = lane; - int group_offset_3 = last_tile * 128 + group_3 * 4; - int raw_rows_3[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_3[0])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 3])) - : "r"(smem_sparse_indices_addr + (unsigned int)(group_offset_3 * 4))); - int row0_3 = ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : 0); - int row1_3 = ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : 0); - int row2_3 = ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : 0); - int row3_3 = ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : 0); - if (global_last_tile == 0) { - tma_gather4_gmem2smem(v_dst_1 + group_3 * 512, tmap_swa_kv, v_col_1, row0_3, row1_3, row2_3, row3_3, kv_full_addr + (load_kv_stage) * 8); - } else { - tma_gather4_gmem2smem(v_dst_1 + group_3 * 512, tmap_compressed_kv, v_col_1, row0_3, row1_3, row2_3, row3_3, kv_full_addr + (load_kv_stage) * 8); - } - } - load_kv_stage += 1; - if (load_kv_stage == 8) { load_kv_stage = 0; _phase_kv_empty ^= 1; } - } - } - } - mbarrier_arrive(tmem_dealloc_addr); - } - } - // ---- Role: softmax_wg ---- - if (warp <= 3) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 184;"); - { // softmax_wg_main - float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; - const int wg_dummy_inc = 0; - int all_num_kv_tiles_1 = (sparse_topk + 128 - 1) / 128; - const int tmem_row_base = ((0) ? warp % 2 * 32 : warp % 4 * 32); - const int tmem_score_row_base = ((0) ? (int)(warp % 4 * 32) : tmem_row_base); - const int n_half = ((0) ? (int)(warp % 4 / 2) : 0); - const int my_row = tmem_row_base + lane; - const int stats_row = n_half * 64 + my_row; - int softmax_tile_cursor = 0; - unsigned int softmax_index_stage = 0; - unsigned int _phase_q_full_0_1 = 0; - #pragma unroll 1 - for (unsigned int work_idx_2 = bid; work_idx_2 < total_work_items; work_idx_2 += num_bids) { - int split_idx_1 = 0; - int query_idx_1 = work_idx_2 >> 1; - int tiles_per_split_1 = all_num_kv_tiles_1; - int first_tile_1 = split_idx_1 * tiles_per_split_1; - int num_kv_tiles_1 = tiles_per_split_1; - int _max_0 = ((sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) : (0)); - int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); - int active_topk = _min_0; - { - float seed_zero[4]; - #pragma unroll - for (int seed_c4 = 0; seed_c4 < 4; seed_c4++) { - seed_zero[seed_c4] = 0.0f; - } - #pragma unroll - for (int seed_half = 0; seed_half < 2; seed_half++) { - #pragma unroll - for (int seed_c = 0; seed_c < 8; seed_c++) { - int seed_addr = taddr + (unsigned int)(((seed_half == 0) ? 96 : 224)) + (unsigned int)(seed_c * 4) + (unsigned int)(tmem_row_base << 16); - tmem_st_x4_f32(seed_addr, seed_zero); - } - } - { - #pragma unroll - for (int seed_half_1 = 0; seed_half_1 < 2; seed_half_1++) { - #pragma unroll - for (int seed_c_1 = 16; seed_c_1 < 24; seed_c_1++) { - int seed_addr_1 = taddr + (unsigned int)(((seed_half_1 == 0) ? 0 : 128)) + (unsigned int)(seed_c_1 * 4) + (unsigned int)(tmem_row_base << 16); - tmem_st_x4_f32(seed_addr_1, seed_zero); - } - } - } - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(s_seeded_addr); - } - { - mbarrier_wait(q_full_addr, _phase_q_full_0_1); - _phase_q_full_0_1 ^= 1; - } - float row_max_val = -CAKE_INF; - float row_sum_val = 0.0f; - int sink_head = ((0) ? my_row : my_row); - if (has_sinks != 0 && sink_head < num_heads && split_idx_1 == 0) { - row_max_val = sinks[sink_head] * 1.4426950408889634f / softmax_scale_log2; - row_sum_val = 1.0f; - } - #pragma unroll 1 - for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { - int pipeline_tile = softmax_tile_cursor + tile_1; - int phase = pipeline_tile & 1; - int s_wait_phase = pipeline_tile >> 1 & 1; - mbarrier_wait(s_full_addr + (phase) * 8, s_wait_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - int s_off = ((phase != 0) ? 128 : 0); - int s_base = taddr + (unsigned int)s_off + (unsigned int)(tmem_score_row_base << 16); - float new_max = row_max_val; - int valid_sparse_cols = ((active_topk < sparse_topk) ? active_topk : sparse_topk); - valid_sparse_cols = valid_sparse_cols - (first_tile_1 + tile_1) * 128 - n_half * 64; - if (valid_sparse_cols < 0) { - valid_sparse_cols = 0; - } - if (valid_sparse_cols > 128) { - valid_sparse_cols = 128; - } - int index_tile_addr = smem_sparse_indices_addr + (unsigned int)(tile_1 * 512); - int index_lane_addr = index_tile_addr + (n_half * 64 + lane) * 4; - int staged_index[4]; - unsigned int invalid_cols[4]; - unsigned int any_invalid = 0; - #pragma unroll - for (int w = 0; w < 4; w++) { - asm volatile("ld.shared.b32 %0, [%1];" : "=r"(*reinterpret_cast(&staged_index[w])) : "r"(index_lane_addr + w * 128)); - unsigned int _vote_0 = __ballot_sync(0xFFFFFFFF, staged_index[w] < 0); - invalid_cols[w] = _vote_0; - any_invalid = any_invalid | invalid_cols[w]; - } - float _tmem_load_0[4]; - tmem_ld_x4(&_tmem_load_0[0], s_base); - float _tmem_load_1[4]; - tmem_ld_x4(&_tmem_load_1[0], s_base); - { - #pragma unroll - for (int frag = 0; frag < 4; frag++) { - float _tmem_load_2[32]; - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]) - : "r"(s_base + frag * 32)); - int frag_valid = valid_sparse_cols - frag * 32; - uint32_t _slice_lo_mask_1; - { - int _lim_0 = frag_valid; - if (_lim_0 <= 0) { _slice_lo_mask_1 = 0u; } - else if (_lim_0 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_0)); - } - } - if (!(_slice_lo_mask_1 & (1u << 0))) _tmem_load_2[0] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 1))) _tmem_load_2[1] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 2))) _tmem_load_2[2] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 3))) _tmem_load_2[3] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 4))) _tmem_load_2[4] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 5))) _tmem_load_2[5] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 6))) _tmem_load_2[6] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 7))) _tmem_load_2[7] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 8))) _tmem_load_2[8] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 9))) _tmem_load_2[9] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 10))) _tmem_load_2[10] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 11))) _tmem_load_2[11] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 12))) _tmem_load_2[12] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 13))) _tmem_load_2[13] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 14))) _tmem_load_2[14] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 15))) _tmem_load_2[15] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 16))) _tmem_load_2[16] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 17))) _tmem_load_2[17] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 18))) _tmem_load_2[18] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 19))) _tmem_load_2[19] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 20))) _tmem_load_2[20] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 21))) _tmem_load_2[21] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 22))) _tmem_load_2[22] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 23))) _tmem_load_2[23] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 24))) _tmem_load_2[24] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 25))) _tmem_load_2[25] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 26))) _tmem_load_2[26] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 27))) _tmem_load_2[27] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 28))) _tmem_load_2[28] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 29))) _tmem_load_2[29] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 30))) _tmem_load_2[30] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 31))) _tmem_load_2[31] = -CAKE_INF; - if (any_invalid != 0) { - unsigned int frag_mask_word = invalid_cols[frag]; - _tmem_load_2[0] = (((frag_mask_word & 1) != 0) ? -CAKE_INF : _tmem_load_2[0]); - _tmem_load_2[1] = (((frag_mask_word >> 1 & 1) != 0) ? -CAKE_INF : _tmem_load_2[1]); - _tmem_load_2[2] = (((frag_mask_word >> 2 & 1) != 0) ? -CAKE_INF : _tmem_load_2[2]); - _tmem_load_2[3] = (((frag_mask_word >> 3 & 1) != 0) ? -CAKE_INF : _tmem_load_2[3]); - _tmem_load_2[4] = (((frag_mask_word >> 4 & 1) != 0) ? -CAKE_INF : _tmem_load_2[4]); - _tmem_load_2[5] = (((frag_mask_word >> 5 & 1) != 0) ? -CAKE_INF : _tmem_load_2[5]); - _tmem_load_2[6] = (((frag_mask_word >> 6 & 1) != 0) ? -CAKE_INF : _tmem_load_2[6]); - _tmem_load_2[7] = (((frag_mask_word >> 7 & 1) != 0) ? -CAKE_INF : _tmem_load_2[7]); - _tmem_load_2[8] = (((frag_mask_word >> 8 & 1) != 0) ? -CAKE_INF : _tmem_load_2[8]); - _tmem_load_2[9] = (((frag_mask_word >> 9 & 1) != 0) ? -CAKE_INF : _tmem_load_2[9]); - _tmem_load_2[10] = (((frag_mask_word >> 10 & 1) != 0) ? -CAKE_INF : _tmem_load_2[10]); - _tmem_load_2[11] = (((frag_mask_word >> 11 & 1) != 0) ? -CAKE_INF : _tmem_load_2[11]); - _tmem_load_2[12] = (((frag_mask_word >> 12 & 1) != 0) ? -CAKE_INF : _tmem_load_2[12]); - _tmem_load_2[13] = (((frag_mask_word >> 13 & 1) != 0) ? -CAKE_INF : _tmem_load_2[13]); - _tmem_load_2[14] = (((frag_mask_word >> 14 & 1) != 0) ? -CAKE_INF : _tmem_load_2[14]); - _tmem_load_2[15] = (((frag_mask_word >> 15 & 1) != 0) ? -CAKE_INF : _tmem_load_2[15]); - _tmem_load_2[16] = (((frag_mask_word >> 16 & 1) != 0) ? -CAKE_INF : _tmem_load_2[16]); - _tmem_load_2[17] = (((frag_mask_word >> 17 & 1) != 0) ? -CAKE_INF : _tmem_load_2[17]); - _tmem_load_2[18] = (((frag_mask_word >> 18 & 1) != 0) ? -CAKE_INF : _tmem_load_2[18]); - _tmem_load_2[19] = (((frag_mask_word >> 19 & 1) != 0) ? -CAKE_INF : _tmem_load_2[19]); - _tmem_load_2[20] = (((frag_mask_word >> 20 & 1) != 0) ? -CAKE_INF : _tmem_load_2[20]); - _tmem_load_2[21] = (((frag_mask_word >> 21 & 1) != 0) ? -CAKE_INF : _tmem_load_2[21]); - _tmem_load_2[22] = (((frag_mask_word >> 22 & 1) != 0) ? -CAKE_INF : _tmem_load_2[22]); - _tmem_load_2[23] = (((frag_mask_word >> 23 & 1) != 0) ? -CAKE_INF : _tmem_load_2[23]); - _tmem_load_2[24] = (((frag_mask_word >> 24 & 1) != 0) ? -CAKE_INF : _tmem_load_2[24]); - _tmem_load_2[25] = (((frag_mask_word >> 25 & 1) != 0) ? -CAKE_INF : _tmem_load_2[25]); - _tmem_load_2[26] = (((frag_mask_word >> 26 & 1) != 0) ? -CAKE_INF : _tmem_load_2[26]); - _tmem_load_2[27] = (((frag_mask_word >> 27 & 1) != 0) ? -CAKE_INF : _tmem_load_2[27]); - _tmem_load_2[28] = (((frag_mask_word >> 28 & 1) != 0) ? -CAKE_INF : _tmem_load_2[28]); - _tmem_load_2[29] = (((frag_mask_word >> 29 & 1) != 0) ? -CAKE_INF : _tmem_load_2[29]); - _tmem_load_2[30] = (((frag_mask_word >> 30 & 1) != 0) ? -CAKE_INF : _tmem_load_2[30]); - _tmem_load_2[31] = (((frag_mask_word >> 31 & 1) != 0) ? -CAKE_INF : _tmem_load_2[31]); - } - float2 _reg_reduce_max2_1 = {-CAKE_INF, -CAKE_INF}; - _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_2[0], _tmem_load_2[1])); - _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_2[2], _tmem_load_2[3])); - _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_2[4], _tmem_load_2[5])); - _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_2[6], _tmem_load_2[7])); - _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_2[8], _tmem_load_2[9])); - _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_2[10], _tmem_load_2[11])); - _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_2[12], _tmem_load_2[13])); - _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_2[14], _tmem_load_2[15])); - _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_2[16], _tmem_load_2[17])); - _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_2[18], _tmem_load_2[19])); - _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_2[20], _tmem_load_2[21])); - _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_2[22], _tmem_load_2[23])); - _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_2[24], _tmem_load_2[25])); - _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_2[26], _tmem_load_2[27])); - _reg_reduce_max2_1.x = max_noftz(_reg_reduce_max2_1.x, max_noftz(_tmem_load_2[28], _tmem_load_2[29])); - _reg_reduce_max2_1.y = max_noftz(_reg_reduce_max2_1.y, max_noftz(_tmem_load_2[30], _tmem_load_2[31])); - float _tmem_load_2_max = row_max_reduce(_reg_reduce_max2_1); - float _max_3 = max_noftz(new_max, _tmem_load_2_max); - new_max = _max_3; - } - } - float _fma_0 = __fmaf_rn(row_max_val, softmax_scale_log2, (-new_max) * softmax_scale_log2); - float delta = _fma_0; - float _exp2_0 = approx_exp2(delta); - float exp_delta = _exp2_0; - float acc_scale = ((row_max_val > -CAKE_INF) ? exp_delta : 1.0f); - smem_stats_max[phase * 128 + stats_row] = acc_scale; - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(stats_addr + (phase) * 8); - row_max_val = new_max; - float safe_max = ((new_max == -CAKE_INF) ? 0.0f : new_max); - float max_scaled = safe_max * softmax_scale_log2; - float block_sum = 0.0f; - { - int p_off = ((phase != 0) ? 224 : 96); - int p_base = taddr + (unsigned int)p_off + (unsigned int)(tmem_row_base << 16); - #pragma unroll - for (int frag_1 = 0; frag_1 < 4; frag_1++) { - float _tmem_load_3[32]; - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]) - : "r"(s_base + frag_1 * 32)); - int frag_valid_1 = valid_sparse_cols - frag_1 * 32; - uint32_t _slice_lo_mask_3; - { - int _lim_2 = frag_valid_1; - if (_lim_2 <= 0) { _slice_lo_mask_3 = 0u; } - else if (_lim_2 >= 32) { _slice_lo_mask_3 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_slice_lo_mask_3) : "r"(_lim_2)); - } - } - if (!(_slice_lo_mask_3 & (1u << 0))) _tmem_load_3[0] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 1))) _tmem_load_3[1] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 2))) _tmem_load_3[2] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 3))) _tmem_load_3[3] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 4))) _tmem_load_3[4] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 5))) _tmem_load_3[5] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 6))) _tmem_load_3[6] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 7))) _tmem_load_3[7] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 8))) _tmem_load_3[8] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 9))) _tmem_load_3[9] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 10))) _tmem_load_3[10] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 11))) _tmem_load_3[11] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 12))) _tmem_load_3[12] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 13))) _tmem_load_3[13] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 14))) _tmem_load_3[14] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 15))) _tmem_load_3[15] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 16))) _tmem_load_3[16] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 17))) _tmem_load_3[17] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 18))) _tmem_load_3[18] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 19))) _tmem_load_3[19] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 20))) _tmem_load_3[20] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 21))) _tmem_load_3[21] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 22))) _tmem_load_3[22] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 23))) _tmem_load_3[23] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 24))) _tmem_load_3[24] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 25))) _tmem_load_3[25] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 26))) _tmem_load_3[26] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 27))) _tmem_load_3[27] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 28))) _tmem_load_3[28] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 29))) _tmem_load_3[29] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 30))) _tmem_load_3[30] = -CAKE_INF; - if (!(_slice_lo_mask_3 & (1u << 31))) _tmem_load_3[31] = -CAKE_INF; - if (any_invalid != 0) { - unsigned int frag_mask_word_1 = invalid_cols[frag_1]; - _tmem_load_3[0] = (((frag_mask_word_1 & 1) != 0) ? -CAKE_INF : _tmem_load_3[0]); - _tmem_load_3[1] = (((frag_mask_word_1 >> 1 & 1) != 0) ? -CAKE_INF : _tmem_load_3[1]); - _tmem_load_3[2] = (((frag_mask_word_1 >> 2 & 1) != 0) ? -CAKE_INF : _tmem_load_3[2]); - _tmem_load_3[3] = (((frag_mask_word_1 >> 3 & 1) != 0) ? -CAKE_INF : _tmem_load_3[3]); - _tmem_load_3[4] = (((frag_mask_word_1 >> 4 & 1) != 0) ? -CAKE_INF : _tmem_load_3[4]); - _tmem_load_3[5] = (((frag_mask_word_1 >> 5 & 1) != 0) ? -CAKE_INF : _tmem_load_3[5]); - _tmem_load_3[6] = (((frag_mask_word_1 >> 6 & 1) != 0) ? -CAKE_INF : _tmem_load_3[6]); - _tmem_load_3[7] = (((frag_mask_word_1 >> 7 & 1) != 0) ? -CAKE_INF : _tmem_load_3[7]); - _tmem_load_3[8] = (((frag_mask_word_1 >> 8 & 1) != 0) ? -CAKE_INF : _tmem_load_3[8]); - _tmem_load_3[9] = (((frag_mask_word_1 >> 9 & 1) != 0) ? -CAKE_INF : _tmem_load_3[9]); - _tmem_load_3[10] = (((frag_mask_word_1 >> 10 & 1) != 0) ? -CAKE_INF : _tmem_load_3[10]); - _tmem_load_3[11] = (((frag_mask_word_1 >> 11 & 1) != 0) ? -CAKE_INF : _tmem_load_3[11]); - _tmem_load_3[12] = (((frag_mask_word_1 >> 12 & 1) != 0) ? -CAKE_INF : _tmem_load_3[12]); - _tmem_load_3[13] = (((frag_mask_word_1 >> 13 & 1) != 0) ? -CAKE_INF : _tmem_load_3[13]); - _tmem_load_3[14] = (((frag_mask_word_1 >> 14 & 1) != 0) ? -CAKE_INF : _tmem_load_3[14]); - _tmem_load_3[15] = (((frag_mask_word_1 >> 15 & 1) != 0) ? -CAKE_INF : _tmem_load_3[15]); - _tmem_load_3[16] = (((frag_mask_word_1 >> 16 & 1) != 0) ? -CAKE_INF : _tmem_load_3[16]); - _tmem_load_3[17] = (((frag_mask_word_1 >> 17 & 1) != 0) ? -CAKE_INF : _tmem_load_3[17]); - _tmem_load_3[18] = (((frag_mask_word_1 >> 18 & 1) != 0) ? -CAKE_INF : _tmem_load_3[18]); - _tmem_load_3[19] = (((frag_mask_word_1 >> 19 & 1) != 0) ? -CAKE_INF : _tmem_load_3[19]); - _tmem_load_3[20] = (((frag_mask_word_1 >> 20 & 1) != 0) ? -CAKE_INF : _tmem_load_3[20]); - _tmem_load_3[21] = (((frag_mask_word_1 >> 21 & 1) != 0) ? -CAKE_INF : _tmem_load_3[21]); - _tmem_load_3[22] = (((frag_mask_word_1 >> 22 & 1) != 0) ? -CAKE_INF : _tmem_load_3[22]); - _tmem_load_3[23] = (((frag_mask_word_1 >> 23 & 1) != 0) ? -CAKE_INF : _tmem_load_3[23]); - _tmem_load_3[24] = (((frag_mask_word_1 >> 24 & 1) != 0) ? -CAKE_INF : _tmem_load_3[24]); - _tmem_load_3[25] = (((frag_mask_word_1 >> 25 & 1) != 0) ? -CAKE_INF : _tmem_load_3[25]); - _tmem_load_3[26] = (((frag_mask_word_1 >> 26 & 1) != 0) ? -CAKE_INF : _tmem_load_3[26]); - _tmem_load_3[27] = (((frag_mask_word_1 >> 27 & 1) != 0) ? -CAKE_INF : _tmem_load_3[27]); - _tmem_load_3[28] = (((frag_mask_word_1 >> 28 & 1) != 0) ? -CAKE_INF : _tmem_load_3[28]); - _tmem_load_3[29] = (((frag_mask_word_1 >> 29 & 1) != 0) ? -CAKE_INF : _tmem_load_3[29]); - _tmem_load_3[30] = (((frag_mask_word_1 >> 30 & 1) != 0) ? -CAKE_INF : _tmem_load_3[30]); - _tmem_load_3[31] = (((frag_mask_word_1 >> 31 & 1) != 0) ? -CAKE_INF : _tmem_load_3[31]); - } - const float2 _fma_b2_3 = {softmax_scale_log2, softmax_scale_log2}; - const float2 _fma_c2_4 = {-max_scaled, -max_scaled}; - float2 _fma_pair_5 = fma_f32x2(make_float2(_tmem_load_3[0], _tmem_load_3[1]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[0] = _fma_pair_5.x; - _tmem_load_3[1] = _fma_pair_5.y; - float2 _fma_pair_6 = fma_f32x2(make_float2(_tmem_load_3[2], _tmem_load_3[3]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[2] = _fma_pair_6.x; - _tmem_load_3[3] = _fma_pair_6.y; - float2 _fma_pair_7 = fma_f32x2(make_float2(_tmem_load_3[4], _tmem_load_3[5]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[4] = _fma_pair_7.x; - _tmem_load_3[5] = _fma_pair_7.y; - float2 _fma_pair_8 = fma_f32x2(make_float2(_tmem_load_3[6], _tmem_load_3[7]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[6] = _fma_pair_8.x; - _tmem_load_3[7] = _fma_pair_8.y; - float2 _fma_pair_9 = fma_f32x2(make_float2(_tmem_load_3[8], _tmem_load_3[9]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[8] = _fma_pair_9.x; - _tmem_load_3[9] = _fma_pair_9.y; - float2 _fma_pair_10 = fma_f32x2(make_float2(_tmem_load_3[10], _tmem_load_3[11]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[10] = _fma_pair_10.x; - _tmem_load_3[11] = _fma_pair_10.y; - float2 _fma_pair_11 = fma_f32x2(make_float2(_tmem_load_3[12], _tmem_load_3[13]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[12] = _fma_pair_11.x; - _tmem_load_3[13] = _fma_pair_11.y; - float2 _fma_pair_12 = fma_f32x2(make_float2(_tmem_load_3[14], _tmem_load_3[15]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[14] = _fma_pair_12.x; - _tmem_load_3[15] = _fma_pair_12.y; - float2 _fma_pair_13 = fma_f32x2(make_float2(_tmem_load_3[16], _tmem_load_3[17]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[16] = _fma_pair_13.x; - _tmem_load_3[17] = _fma_pair_13.y; - float2 _fma_pair_14 = fma_f32x2(make_float2(_tmem_load_3[18], _tmem_load_3[19]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[18] = _fma_pair_14.x; - _tmem_load_3[19] = _fma_pair_14.y; - float2 _fma_pair_15 = fma_f32x2(make_float2(_tmem_load_3[20], _tmem_load_3[21]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[20] = _fma_pair_15.x; - _tmem_load_3[21] = _fma_pair_15.y; - float2 _fma_pair_16 = fma_f32x2(make_float2(_tmem_load_3[22], _tmem_load_3[23]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[22] = _fma_pair_16.x; - _tmem_load_3[23] = _fma_pair_16.y; - float2 _fma_pair_17 = fma_f32x2(make_float2(_tmem_load_3[24], _tmem_load_3[25]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[24] = _fma_pair_17.x; - _tmem_load_3[25] = _fma_pair_17.y; - float2 _fma_pair_18 = fma_f32x2(make_float2(_tmem_load_3[26], _tmem_load_3[27]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[26] = _fma_pair_18.x; - _tmem_load_3[27] = _fma_pair_18.y; - float2 _fma_pair_19 = fma_f32x2(make_float2(_tmem_load_3[28], _tmem_load_3[29]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[28] = _fma_pair_19.x; - _tmem_load_3[29] = _fma_pair_19.y; - float2 _fma_pair_20 = fma_f32x2(make_float2(_tmem_load_3[30], _tmem_load_3[31]), _fma_b2_3, _fma_c2_4); - _tmem_load_3[30] = _fma_pair_20.x; - _tmem_load_3[31] = _fma_pair_20.y; - #pragma unroll - for (int _le = 0; _le < 32; _le++) { - _tmem_load_3[_le] = approx_exp2(_tmem_load_3[_le]); - } - float2 _reg_reduce_sum2_21 = make_float2(0.0f, 0.0f); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[0], _tmem_load_3[1])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[2], _tmem_load_3[3])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[4], _tmem_load_3[5])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[6], _tmem_load_3[7])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[8], _tmem_load_3[9])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[10], _tmem_load_3[11])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[12], _tmem_load_3[13])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[14], _tmem_load_3[15])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[16], _tmem_load_3[17])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[18], _tmem_load_3[19])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[20], _tmem_load_3[21])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[22], _tmem_load_3[23])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[24], _tmem_load_3[25])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[26], _tmem_load_3[27])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[28], _tmem_load_3[29])); - _reg_reduce_sum2_21 = add_f32x2(_reg_reduce_sum2_21, make_float2(_tmem_load_3[30], _tmem_load_3[31])); - float _tmem_load_3_sum = _reg_reduce_sum2_21.x + _reg_reduce_sum2_21.y; - block_sum = block_sum + _tmem_load_3_sum; - { - uint32_t _pv_packed[8]; - #pragma unroll - for (int _j = 0; _j < 8; _j++) { - uint32_t _pk; - asm("{\n\t" - ".reg .b16 _lo, _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}\n" - : "=r"(_pk) : "f"(_tmem_load_3[0 + _j * 4]), "f"(_tmem_load_3[0 + _j * 4 + 1]), - "f"(_tmem_load_3[0 + _j * 4 + 2]), "f"(_tmem_load_3[0 + _j * 4 + 3])); - _pv_packed[_j] = _pk; - } - tmem_st_x8_u32(p_base + frag_1 * 8, _pv_packed); - } - } - } - { - float _fma_1 = __fmaf_rn(row_sum_val, acc_scale, block_sum); - row_sum_val = _fma_1; - } - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(p_full_addr + (phase) * 8); - } - smem_stats_sum[stats_row] = row_sum_val; - smem_stats_final_max[stats_row] = row_max_val; - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(sum_ready_addr); - softmax_tile_cursor = softmax_tile_cursor + num_kv_tiles_1; - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(tmem_dealloc_addr); - } - } - // ---- Role: correction_wg ---- - if (warp >= 4 && warp <= 7) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 184;"); - { // correction_wg_main - float softmax_scale_log2_1 = bmm1_scale[0] * 1.4426950408889634f; - float output_scale = bmm2_scale[0]; - const int wg_dummy_inc_1 = 0; - int all_num_kv_tiles_2 = (sparse_topk + 128 - 1) / 128; - const int tmem_row_base_1 = ((0) ? warp % 2 * 32 : warp % 4 * 32); - const int n_half_1 = ((0) ? (int)(warp % 4 / 2) : 0); - const int my_row_1 = tmem_row_base_1 + lane; - const int stats_row_1 = n_half_1 * 64 + my_row_1; - const int corr_row = tmem_row_base_1 << 16; - int correction_tile_cursor = 0; - unsigned int _phase_q_full_0_2 = 0; - unsigned int _phase_pv_done_0 = 0; - unsigned int _phase_o_done_0 = 0; - unsigned int _phase_sum_ready_0 = 0; - #pragma unroll 1 - for (unsigned int work_idx_3 = bid; work_idx_3 < total_work_items; work_idx_3 += num_bids) { - int split_idx_2 = 0; - int query_idx_2 = work_idx_3 >> 1; - int v_chunk_1 = work_idx_3 & 1; - int tiles_per_split_2 = all_num_kv_tiles_2; - int first_tile_2 = split_idx_2 * tiles_per_split_2; - int num_kv_tiles_2 = tiles_per_split_2; - { - mbarrier_wait(q_full_addr, _phase_q_full_0_2); - _phase_q_full_0_2 ^= 1; - } - #pragma unroll 1 - for (int tile_2 = 0; tile_2 < num_kv_tiles_2; tile_2++) { - int pipeline_tile_1 = correction_tile_cursor + tile_2; - int phase_1 = pipeline_tile_1 & 1; - int stats_wait_phase = pipeline_tile_1 >> 1 & 1; - mbarrier_wait(stats_addr + (phase_1) * 8, stats_wait_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - float acc_scale_1 = smem_stats_max[phase_1 * 128 + stats_row_1]; - if (tile_2 > 0) { - mbarrier_wait(pv_done_addr, _phase_pv_done_0); - _phase_pv_done_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - } - if (tile_2 > 0) { - int _vote_1 = __any_sync(0xFFFFFFFF, acc_scale_1 < 1.0f); - int any_rescale = _vote_1; - if (any_rescale != 0) { - #pragma unroll - for (int vs = 0; vs < 2; vs++) { - int o_base = taddr + 256 + (unsigned int)(vs * 128) + (unsigned int)corr_row; - #pragma unroll - for (int c = 0; c < 128; c += 16) { - float _tmem_load_4[16]; - tmem_ld_x16(&_tmem_load_4[0], o_base + c); - const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; - #pragma unroll - for (int _ls = 0; _ls < 8; _ls++) - mul_f32x2_inplace(&reinterpret_cast(_tmem_load_4)[_ls], _scale2_0); - tmem_st_x16_f32(o_base + c, _tmem_load_4); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - } - } - } - } - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(corr_done_addr + (phase_1) * 8); - } - mbarrier_wait(o_done_addr, _phase_o_done_0); - _phase_o_done_0 ^= 1; - mbarrier_wait(sum_ready_addr, _phase_sum_ready_0); - _phase_sum_ready_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - float total_sum = smem_stats_sum[stats_row_1]; - float final_max = smem_stats_final_max[stats_row_1]; - float _rcp_0 = approx_rcp(total_sum); - float inv_sum = ((total_sum > 0.0f) ? _rcp_0 : 0.0f); - int head_idx = my_row_1; - int direct_o_offset = (query_idx_2 * num_heads + head_idx) * 512 + v_chunk_1 * 256; - int partial_o_offset = (query_idx_2 * num_heads + head_idx + split_idx_2) * 512 + ((0) ? v_chunk_1 : n_half_1) * 256; - int o_offset = direct_o_offset; - if ((0 & (int)(n_half_1 == 0) & (int)(head_idx < num_heads)) != 0) { - int stat_offset = query_idx_2 * num_heads + head_idx + split_idx_2; - float _log2_0; - asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_0) : "f"(total_sum)); - partial_lse[stat_offset] = ((total_sum > 0.0f) ? final_max * softmax_scale_log2_1 + _log2_0 : -CAKE_INF); - } - #pragma unroll - for (int vs_1 = 0; vs_1 < 2; vs_1++) { - int o_base_epi = taddr + 256 + (unsigned int)(vs_1 * 128) + (unsigned int)corr_row; - #pragma unroll - for (int c_1 = 0; c_1 < 128; c_1 += 32) { - float _tmem_load_5[32]; - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(_tmem_load_5[0]), "=f"(_tmem_load_5[1]), "=f"(_tmem_load_5[2]), "=f"(_tmem_load_5[3]), "=f"(_tmem_load_5[4]), "=f"(_tmem_load_5[5]), "=f"(_tmem_load_5[6]), "=f"(_tmem_load_5[7]), "=f"(_tmem_load_5[8]), "=f"(_tmem_load_5[9]), "=f"(_tmem_load_5[10]), "=f"(_tmem_load_5[11]), "=f"(_tmem_load_5[12]), "=f"(_tmem_load_5[13]), "=f"(_tmem_load_5[14]), "=f"(_tmem_load_5[15]), "=f"(_tmem_load_5[16]), "=f"(_tmem_load_5[17]), "=f"(_tmem_load_5[18]), "=f"(_tmem_load_5[19]), "=f"(_tmem_load_5[20]), "=f"(_tmem_load_5[21]), "=f"(_tmem_load_5[22]), "=f"(_tmem_load_5[23]), "=f"(_tmem_load_5[24]), "=f"(_tmem_load_5[25]), "=f"(_tmem_load_5[26]), "=f"(_tmem_load_5[27]), "=f"(_tmem_load_5[28]), "=f"(_tmem_load_5[29]), "=f"(_tmem_load_5[30]), "=f"(_tmem_load_5[31]) - : "r"(o_base_epi + c_1)); - int gmem_base = o_offset + vs_1 * 128 + c_1; - if (head_idx < num_heads) { - #pragma unroll - for (int j = 0; j < 32; j += 16) { - { - const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; - #if __CUDA_ARCH__ >= 1000 - #pragma unroll - for (int _ps = 0; _ps < 8; _ps++) - mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_5[j])[_ps], _prescale2_1); - #else - #pragma unroll - for (int _ps = 0; _ps < 16; _ps++) - _tmem_load_5[j + _ps] *= inv_sum * output_scale; - #endif - { - __nv_bfloat162 _pk0 = __floats2bfloat162_rn(_tmem_load_5[j + 0], _tmem_load_5[j + 1]); - unsigned _pk_u0 = *reinterpret_cast(&_pk0); - __nv_bfloat162 _pk1 = __floats2bfloat162_rn(_tmem_load_5[j + 2], _tmem_load_5[j + 3]); - unsigned _pk_u1 = *reinterpret_cast(&_pk1); - __nv_bfloat162 _pk2 = __floats2bfloat162_rn(_tmem_load_5[j + 4], _tmem_load_5[j + 5]); - unsigned _pk_u2 = *reinterpret_cast(&_pk2); - __nv_bfloat162 _pk3 = __floats2bfloat162_rn(_tmem_load_5[j + 6], _tmem_load_5[j + 7]); - unsigned _pk_u3 = *reinterpret_cast(&_pk3); - __nv_bfloat162 _pk4 = __floats2bfloat162_rn(_tmem_load_5[j + 8], _tmem_load_5[j + 9]); - unsigned _pk_u4 = *reinterpret_cast(&_pk4); - __nv_bfloat162 _pk5 = __floats2bfloat162_rn(_tmem_load_5[j + 10], _tmem_load_5[j + 11]); - unsigned _pk_u5 = *reinterpret_cast(&_pk5); - __nv_bfloat162 _pk6 = __floats2bfloat162_rn(_tmem_load_5[j + 12], _tmem_load_5[j + 13]); - unsigned _pk_u6 = *reinterpret_cast(&_pk6); - __nv_bfloat162 _pk7 = __floats2bfloat162_rn(_tmem_load_5[j + 14], _tmem_load_5[j + 15]); - unsigned _pk_u7 = *reinterpret_cast(&_pk7); - asm volatile( - "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" - :: "l"((void*)(&((__nv_bfloat16*)(O + (gmem_base + j)))[0])), "r"(_pk_u0), "r"(_pk_u1), "r"(_pk_u2), "r"(_pk_u3), "r"(_pk_u4), "r"(_pk_u5), "r"(_pk_u6), "r"(_pk_u7) : "memory"); - } - } - } - } - } - } - correction_tile_cursor = correction_tile_cursor + num_kv_tiles_2; - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(tmem_dealloc_addr); - } - } - // ---- Role: mma_warp ---- - if (warp == 8) { - { // mma_warp_main - const int wg2_dummy_1 = 0; - int all_num_kv_tiles_3 = (sparse_topk + 128 - 1) / 128; - unsigned int mma_k_stage = 0; - unsigned int mma_v_stage = 0; - unsigned int mma_kv_stage = 0; - int mma_tile_cursor = 0; - unsigned int _phase_s_seeded_0 = 0; - unsigned int _phase_q_full_0_3 = 0; - unsigned int _phase_q_pair_ready_0 = 0; - unsigned int _phase_k_full = 0; - unsigned int _phase_kv_full = 0; - unsigned int _phase_kv_pair_ready = 0; - unsigned int _phase_v_full = 0; - #pragma unroll 1 - for (unsigned int work_idx_4 = bid; work_idx_4 < total_work_items; work_idx_4 += num_bids) { - int split_idx_3 = 0; - int tiles_per_split_3 = all_num_kv_tiles_3; - int first_tile_3 = split_idx_3 * tiles_per_split_3; - int num_kv_tiles_3 = tiles_per_split_3; - { - { - mbarrier_wait(s_seeded_addr, _phase_s_seeded_0); - _phase_s_seeded_0 ^= 1; - } - asm volatile("tcgen05.fence::after_thread_sync;"); - } - { - mbarrier_wait(q_full_addr, _phase_q_full_0_3); - _phase_q_full_0_3 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - } - int first_pv = 1; - #pragma unroll 1 - for (int tile_3 = 0; tile_3 < num_kv_tiles_3; tile_3++) { - int pipeline_tile_2 = mma_tile_cursor + tile_3; - int phase_2 = pipeline_tile_2 & 1; - int score_col = ((phase_2 != 0) ? 128 : 0); - { - #pragma unroll 1 - for (int n = 0; n < 4; n++) { - mbarrier_wait(kv_full_addr + (mma_kv_stage) * 8, _phase_kv_full); - asm volatile("tcgen05.fence::after_thread_sync;"); - { - int _mma_a_lo_4 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (n) * 1024); - int _mma_b_lo_4 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (mma_kv_stage) * 1024); - { - uint64_t _mma_ss_a_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_4); - uint64_t _mma_ss_b_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_4); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem_scratch + (score_col)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 136314896, ((n == 0) ? 0 : 1)); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem_scratch + (score_col)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 136314896, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem_scratch + (score_col)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 136314896, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem_scratch + (score_col)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 136314896, 1); - } - } - if (n == 3) { - elect_commit(s_full_addr + (phase_2) * 8); - } - elect_commit(kv_empty_addr + (mma_kv_stage) * 8); - } - mma_kv_stage += 1; - if (mma_kv_stage == 8) { mma_kv_stage = 0; _phase_kv_full ^= 1; _phase_kv_pair_ready ^= 1; } - } - } - if (tile_3 > 0) { - int prev_pipeline_tile = pipeline_tile_2 - 1; - int prev_phase = prev_pipeline_tile & 1; - int pv_wait_phase = prev_pipeline_tile >> 1 & 1; - mbarrier_wait(p_full_addr + (prev_phase) * 8, pv_wait_phase); - mbarrier_wait(corr_done_addr + (prev_phase) * 8, pv_wait_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - int p_col = ((prev_phase != 0) ? 224 : 96); - { - #pragma unroll - for (int vs_2 = 0; vs_2 < 2; vs_2++) { - mbarrier_wait(kv_full_addr + (mma_kv_stage) * 8, _phase_kv_full); - asm volatile("tcgen05.fence::after_thread_sync;"); - int output_col = 256 + vs_2 * 128; - { - int _mma_b_lo_8 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (mma_kv_stage) * 1024); - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 dhi, blo, ta, id;\n\t" - ".reg .b64 db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "" - "mov.b32 dhi, 0x40004040;\n\t" - "mov.b32 id, 136380432;\n\t" - "mov.b32 ta, %2;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p0;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" - "}\n" - :: "r"((tmem_tmem_scratch + (output_col))), "r"(_mma_b_lo_8), "r"(tmem_tmem_scratch + p_col), "r"(((first_pv) ? 0 : 1))); - elect_commit(kv_empty_addr + (mma_kv_stage) * 8); - } - mma_kv_stage += 1; - if (mma_kv_stage == 8) { mma_kv_stage = 0; _phase_kv_full ^= 1; _phase_kv_pair_ready ^= 1; } - } - } - first_pv = 0; - elect_commit(pv_done_addr); - } - } - int last_pipeline_tile = mma_tile_cursor + num_kv_tiles_3 - 1; - int last_phase = last_pipeline_tile & 1; - int drain_wait_phase = last_pipeline_tile >> 1 & 1; - mbarrier_wait(p_full_addr + (last_phase) * 8, drain_wait_phase); - mbarrier_wait(corr_done_addr + (last_phase) * 8, drain_wait_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - int p_col_last = ((last_phase != 0) ? 224 : 96); - { - #pragma unroll - for (int vs_3 = 0; vs_3 < 2; vs_3++) { - mbarrier_wait(kv_full_addr + (mma_kv_stage) * 8, _phase_kv_full); - asm volatile("tcgen05.fence::after_thread_sync;"); - int output_col_d = 256 + vs_3 * 128; - { - int _mma_b_lo_12 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (mma_kv_stage) * 1024); - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 dhi, blo, ta, id;\n\t" - ".reg .b64 db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "" - "mov.b32 dhi, 0x40004040;\n\t" - "mov.b32 id, 136380432;\n\t" - "mov.b32 ta, %2;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p0;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" - "}\n" - :: "r"((tmem_tmem_scratch + (output_col_d))), "r"(_mma_b_lo_12), "r"(tmem_tmem_scratch + p_col_last), "r"(((first_pv) ? 0 : 1))); - elect_commit(kv_empty_addr + (mma_kv_stage) * 8); - } - mma_kv_stage += 1; - if (mma_kv_stage == 8) { mma_kv_stage = 0; _phase_kv_full ^= 1; _phase_kv_pair_ready ^= 1; } - } - } - elect_commit(q_empty_addr); - elect_commit(o_done_addr); - mma_tile_cursor = mma_tile_cursor + num_kv_tiles_3; - } - mbarrier_arrive(tmem_dealloc_addr); - unsigned int _phase_tmem_dealloc_0 = 0; - mbarrier_wait(tmem_dealloc_addr, _phase_tmem_dealloc_0); - _phase_tmem_dealloc_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - { - int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); - asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); - } - } - } - // ---- Role: empty1 ---- - if (warp >= 14 && warp <= 15) { - { // empty1_main - const int wg2_dummy_2 = 0; - unsigned int _phase_q_full_0_4 = 0; - { - #pragma unroll 1 - for (unsigned int work_idx_5 = bid; work_idx_5 < total_work_items; work_idx_5 += num_bids) { - mbarrier_wait(q_full_addr, _phase_q_full_0_4); - _phase_q_full_0_4 ^= 1; - } - } - } - } - - // Cleanup -} - -} // extern "C" diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_c20607c57cc11784015e_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_c20607c57cc11784015e_binding.cu deleted file mode 100644 index a287b064d89..00000000000 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_c20607c57cc11784015e_binding.cu +++ /dev/null @@ -1,461 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -// Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_c20607c57cc11784015e'. -#include -#include -#include -#include - -#include "tvm_ffi_utils.h" - -#include -#include -#include -#include -#include -#include -#include - -extern "C" __global__ void kernel_cake_dsv4_c20607c57cc11784015e(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items); - - -namespace cake_host_shim_9514fa50724352e6 { - -using tvm::ffi::Optional; -using tvm::ffi::TensorView; - -class ScopedCudaDevice { - public: - explicit ScopedCudaDevice(int device_id) { - cudaError_t error = cudaGetDevice(&previous_device_); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed before host-shim launch: cudaError=" - << static_cast(error); - if (previous_device_ != device_id) { - error = cudaSetDevice(device_id); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaSetDevice failed before host-shim launch for cuda:" - << device_id << ": cudaError=" << static_cast(error); - restore_ = true; - } - } - - ScopedCudaDevice(const ScopedCudaDevice&) = delete; - ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; - - ~ScopedCudaDevice() noexcept { - if (restore_) { - (void)cudaSetDevice(previous_device_); - } - } - - private: - int previous_device_ = -1; - bool restore_ = false; -}; - -inline int64_t CakeDeviceMultiprocessorCount(int device_id) { - constexpr int kMaxCachedCudaDevices = 64; - TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) - << "physical-SM-count cache does not cover cuda:" << device_id; - static std::atomic count_by_device[kMaxCachedCudaDevices]{}; - int cached = count_by_device[device_id].load(std::memory_order_acquire); - if (cached > 0) return cached; - - int count = 0; - cudaError_t error = cudaDeviceGetAttribute( - &count, cudaDevAttrMultiProcessorCount, device_id); - TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) - << "querying multiProcessorCount failed for cuda:" << device_id - << ": cudaError=" << static_cast(error) << ", count=" << count; - // Concurrent first launches may repeat the immutable device query, but all - // publication is atomic and every later hot-path lookup is one acquire load. - count_by_device[device_id].store(count, std::memory_order_release); - return count; -} - -inline void CheckCudaTensor(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) - << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; -} - -inline void CheckSameCudaDevice( - const TensorView& t, - const TensorView& reference, - const char* name, - const char* reference_name) { - TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) - << name << " must be on the same CUDA device as " << reference_name - << ": got cuda:" << t.device().device_id - << " versus cuda:" << reference.device().device_id; -} - -inline void CheckCurrentCudaDevice( - const TensorView& reference, - const char* reference_name) { - int current_device = -1; - cudaError_t error = cudaGetDevice(¤t_device); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed while validating " << reference_name - << ": cudaError=" << static_cast(error); - TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) - << "current CUDA device must match " << reference_name - << ": current=cuda:" << current_device - << ", tensor=cuda:" << reference.device().device_id; -} - -inline void CheckContiguous(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; -} - -inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { - DLDataType d = t.dtype(); - TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) - << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits - << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits - << ", lanes=" << (int)d.lanes << ")"; -} - -// A logical axis.outer(trailing) folds every source dim above the trailing -// dimensions. Shape products are independent of physical strides, so verify -// the leading dimensions form one dense row-major chain instead of inventing -// a "folded stride". The descriptor reads its exact adjacent physical step -// separately through stride[-(trailing + 1)]. -inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { - TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) - << name << " cannot fold leading dimensions above " << trailing - << " trailing dims from ndim=" << t.ndim(); - int outer_last = t.ndim() - trailing - 1; - if (outer_last <= 0) { - return; - } - int64_t step = t.stride(outer_last); - TVM_FFI_CHECK(step > 0, ValueError) - << name << " physical strides must be positive"; - int64_t expected = step; - for (int axis = outer_last - 1; axis >= 0; --axis) { - expected *= t.size(axis + 1); - if (t.size(axis) > 1) { - TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) - << name << " leading dims are not physically foldable above " << trailing - << " trailing dims: stride(" << axis << ")=" << t.stride(axis) - << ", expected " << expected; - } - } -} - -// 4D TMA descriptor for buffer 'tmap_q' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 3, ValueError) - << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - int64_t d2 = t.size(t.ndim() - 2); - int64_t d3 = t.size(t.ndim() - 3); - TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) - << "TMA source 'tmap_q' trailing dims must be positive"; - TVM_FFI_CHECK(d1 % 128 == 0, ValueError) - << "TMA source 'tmap_q' extent " << d1 - << " must divide exactly by " << 128; - uint64_t global_dim[4] = {(uint64_t)(128), (uint64_t)(d2), (uint64_t)((d1 / 128)), (uint64_t)(d3)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; - TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) - << "TMA box (128, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; - int64_t carrier_stride_0 = d1; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; - int64_t carrier_stride_1 = 128; - TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; - TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; - TVM_FFI_CHECK((carrier_stride_1 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; - int64_t carrier_stride_2 = (d2 * d1); - TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; - TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; - TVM_FFI_CHECK((carrier_stride_2 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; - uint64_t global_strides[3] = { - (uint64_t)((carrier_stride_0 * 8) / 8), - (uint64_t)((carrier_stride_1 * 8) / 8), - (uint64_t)((carrier_stride_2 * 8) / 8), - }; - uint32_t box_dim[4] = {128u, 64u, 1u, 1u}; - uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_swa_kv' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_swa_kv' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; - TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (128, 1) exceeds resolved global dims for 'tmap_swa_kv'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 8) / 8), - }; - uint32_t box_dim[2] = {128u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_compressed_kv' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_compressed_kv(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_compressed_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_compressed_kv' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_compressed_kv' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_compressed_kv"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_compressed_kv' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved a non-positive global dim"; - TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (128, 1) exceeds resolved global dims for 'tmap_compressed_kv'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 8) / 8), - }; - uint32_t box_dim[2] = {128u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_kv') failed: CUresult=" << (int)r; - return tm; -} - -void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_compressed_kv, TensorView arg_O, TensorView arg_partial_lse, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_num_query_tokens, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, int64_t grid_x, int64_t grid_y, int64_t grid_z) { - DLDevice dev = arg_tmap_q.device(); - ScopedCudaDevice device_guard(dev.device_id); - CheckCudaTensor(arg_tmap_q, "tmap_q"); - CheckDtype(arg_tmap_q, "tmap_q", 1, 8, 1); - CheckContiguous(arg_tmap_q, "tmap_q"); - CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); - CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 1, 8, 1); - CheckCudaTensor(arg_tmap_compressed_kv, "tmap_compressed_kv"); - CheckDtype(arg_tmap_compressed_kv, "tmap_compressed_kv", 1, 8, 1); - CheckCudaTensor(arg_O, "O"); - CheckDtype(arg_O, "O", 4, 16, 1); - CheckContiguous(arg_O, "O"); - CheckCudaTensor(arg_partial_lse, "partial_lse"); - CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); - CheckContiguous(arg_partial_lse, "partial_lse"); - CheckCudaTensor(arg_swa_indices, "swa_indices"); - CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); - CheckContiguous(arg_swa_indices, "swa_indices"); - CheckCudaTensor(arg_compressed_indices, "compressed_indices"); - CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); - CheckContiguous(arg_compressed_indices, "compressed_indices"); - CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); - CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckCudaTensor(arg_sinks, "sinks"); - CheckDtype(arg_sinks, "sinks", 2, 32, 1); - CheckContiguous(arg_sinks, "sinks"); - CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); - CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); - CheckContiguous(arg_bmm1_scale, "bmm1_scale"); - CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); - CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); - CheckContiguous(arg_bmm2_scale, "bmm2_scale"); - TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) - << "scalar 'num_heads' value " << arg_num_heads - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) - << "scalar 'swa_index_stride' value " << arg_swa_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) - << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) - << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_num_query_tokens >= -2147483648LL && arg_num_query_tokens <= 2147483647LL, ValueError) - << "scalar 'num_query_tokens' value " << arg_num_query_tokens - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) - << "scalar 'sparse_topk' value " << arg_sparse_topk - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) - << "scalar 'has_sinks' value " << arg_has_sinks - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) - << "scalar 'total_work_items' value " << arg_total_work_items - << " is outside i32 range [-2147483648, 2147483647]"; - CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); - CheckSameCudaDevice(arg_tmap_compressed_kv, arg_tmap_q, "tmap_compressed_kv", "tmap_q"); - CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); - CheckSameCudaDevice(arg_partial_lse, arg_tmap_q, "partial_lse", "tmap_q"); - CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); - CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); - CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); - CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); - CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); - CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); - TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) - << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y - << ", " << grid_z << ")"; - TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) - << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z - << ") must be divisible by cluster dims (2, 1, 1)"; - - cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); - CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); - CUtensorMap p_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); - CUtensorMap p_tmap_compressed_kv = EncodeTma_tmap_compressed_kv(arg_tmap_compressed_kv); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); - float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); - int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); - int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); - int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); - float* p_sinks = static_cast(arg_sinks.data_ptr()); - float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); - float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); - int32_t v_num_heads = (int32_t)arg_num_heads; - int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; - int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; - int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; - int32_t v_num_query_tokens = (int32_t)arg_num_query_tokens; - int32_t v_sparse_topk = (int32_t)arg_sparse_topk; - int32_t v_has_sinks = (int32_t)arg_has_sinks; - int32_t v_total_work_items = (int32_t)arg_total_work_items; - void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_tmap_compressed_kv, &p_O, &p_partial_lse, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_num_query_tokens, &v_sparse_topk, &v_has_sinks, &v_total_work_items}; - - static std::mutex smem_mu; - static auto* smem_status_by_device = - new std::unordered_map(); - cudaError_t smem_status = cudaSuccess; - { - std::lock_guard lock(smem_mu); - auto it = smem_status_by_device->find(dev.device_id); - if (it == smem_status_by_device->end()) { - int current_device = -1; - smem_status = cudaGetDevice(¤t_device); - if (smem_status == cudaSuccess && current_device != dev.device_id) { - smem_status = cudaErrorInvalidDevice; - } - if (smem_status == cudaSuccess) { - smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_c20607c57cc11784015e), - cudaFuncAttributeMaxDynamicSharedMemorySize, - 190976); - } - smem_status_by_device->emplace(dev.device_id, smem_status); - } else { - smem_status = it->second; - } - } - TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_c20607c57cc11784015e failed: " - << cudaGetErrorString(smem_status); - dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); - dim3 block(512u, 1u, 1u); - - // Extended direct-source launch for a linked CUDA kernel symbol. - cudaLaunchAttribute attrs[2]{}; - int n = 0; - attrs[n].id = cudaLaunchAttributeClusterDimension; - attrs[n].val.clusterDim.x = 2u; - attrs[n].val.clusterDim.y = 1u; - attrs[n].val.clusterDim.z = 1u; - ++n; - attrs[n].id = cudaLaunchAttributeClusterSchedulingPolicyPreference; - attrs[n].val.clusterSchedulingPolicyPreference = cudaClusterSchedulingPolicySpread; - ++n; - cudaLaunchConfig_t config{}; - config.gridDim = grid; - config.blockDim = block; - config.dynamicSmemBytes = 190976u; - config.stream = stream; - config.attrs = attrs; - config.numAttrs = n; - cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_c20607c57cc11784015e), kargs); - TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_c20607c57cc11784015e failed: " - << cudaGetErrorString(launch_status); -} - -} // namespace cake_host_shim_9514fa50724352e6 - -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_9514fa50724352e6::Run); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_c20607c57cc11784015e_kernel.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_c20607c57cc11784015e_kernel.cu deleted file mode 100644 index e88df9fd3d5..00000000000 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_c20607c57cc11784015e_kernel.cu +++ /dev/null @@ -1,2821 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -typedef signed char int8_t; -typedef unsigned char uint8_t; -typedef unsigned short uint16_t; -typedef unsigned int uint32_t; -#if defined(__CUDACC_RTC__) -typedef unsigned long long uint64_t; -#else -typedef unsigned long uint64_t; -#endif -static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); -typedef signed int int32_t; -typedef short int int16_t; -struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; -struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; -static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); -static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); - -#if defined(__CUDACC_RTC__) -typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; -#else -#include -#endif - -static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); -static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); -#include -#include - -#define CAKE_INF CUDART_INF_F -#define TMEM_NCOLS 512 -#define TMEM_TMEM_SCRATCH_OFFSET 0 -#define NUM_K_PIPE_STAGES 2 -#define NUM_V_PIPE_STAGES 2 -#define NUM_KV_PIPE_STAGES 1 -#define NUM_INDEX_PIPE_STAGES 6 -#define SMEM_SMEM_Q_OFF 1024 -#define SMEM_SMEM_Q_STAGE_BYTES 8192 -#define SMEM_SMEM_Q_STRIDE 8192 -#define SMEM_SMEM_KV_OFF 33792 -#define SMEM_SMEM_KV_STAGE_BYTES 16384 -#define SMEM_SMEM_KV_STRIDE 16384 -#define SMEM_SMEM_V_OFF 33792 -#define SMEM_SMEM_V_STAGE_BYTES 16384 -#define SMEM_SMEM_V_STRIDE 16384 -#define SMEM_SMEM_Q_FULL_OFF 1024 -#define SMEM_SMEM_Q_FULL_STAGE_BYTES 32768 -#define SMEM_SMEM_Q_FULL_STRIDE 32768 -#define SMEM_SMEM_K_FULL_OFF 33792 -#define SMEM_SMEM_K_FULL_STAGE_BYTES 32768 -#define SMEM_SMEM_K_FULL_STRIDE 32768 -#define SMEM_SMEM_V_FULL_OFF 99328 -#define SMEM_SMEM_V_FULL_STAGE_BYTES 32768 -#define SMEM_SMEM_V_FULL_STRIDE 32768 -#define SMEM_SMEM_STATS_MAX_OFF 164864 -#define SMEM_SMEM_STATS_MAX_STAGE_BYTES 1024 -#define SMEM_SMEM_STATS_MAX_STRIDE 1024 -#define SMEM_SMEM_STATS_SUM_OFF 165888 -#define SMEM_SMEM_STATS_SUM_STAGE_BYTES 512 -#define SMEM_SMEM_STATS_SUM_STRIDE 512 -#define SMEM_SMEM_STATS_FINAL_MAX_OFF 166400 -#define SMEM_SMEM_STATS_FINAL_MAX_STAGE_BYTES 512 -#define SMEM_SMEM_STATS_FINAL_MAX_STRIDE 512 -#define SMEM_SMEM_SOFTMAX_WARP_PAIR_EXCHANGE_OFF 166912 -#define SMEM_SMEM_SOFTMAX_WARP_PAIR_EXCHANGE_STAGE_BYTES 1024 -#define SMEM_SMEM_SOFTMAX_WARP_PAIR_EXCHANGE_STRIDE 1024 -#define SMEM_SMEM_CORR_WARP_PAIR_EXCHANGE_OFF 167936 -#define SMEM_SMEM_CORR_WARP_PAIR_EXCHANGE_STAGE_BYTES 512 -#define SMEM_SMEM_CORR_WARP_PAIR_EXCHANGE_STRIDE 512 -#define SMEM_SMEM_SPARSE_INDICES_OFF 168448 -#define SMEM_SMEM_SPARSE_INDICES_STAGE_BYTES 1024 -#define SMEM_SMEM_SPARSE_INDICES_STRIDE 1024 -#define SMEM_SMEM_P_FP8_OFF 174592 -#define SMEM_SMEM_P_FP8_STAGE_BYTES 8192 -#define SMEM_SMEM_P_FP8_STRIDE 8192 -#define SMEM_TOTAL 190976 -#define THREADS 512 - -#include - -__device__ __forceinline__ uint32_t elect_sync() { - uint32_t pred = 0; - asm volatile( - "{\n\t" - ".reg .pred %%px;\n\t" - "elect.sync _|%%px, %1;\n\t" - "@%%px mov.s32 %0, 1;\n\t" - "}\n" - : "+r"(pred) - : "r"(0xFFFFFFFF)); - return pred; -} - - -__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { - asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" - :: "r"(mbar_addr), "r"(count) : "memory"); -} - -__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { - asm volatile("mbarrier.init.b64 [%0], %1;" - :: "l"(mbar_addr), "r"(count) : "memory"); -} - - -__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - - -// CTA-local pipelines have short, resident producer/consumer edges. Omitting -// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the -// explicit loop still makes this helper blocking until acquire succeeds. -__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE;\n\t" - "bra.uni LAB_WAIT;\n\t" - "DONE:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Source-faithful relaxed CTA wait used only by a typed protocol that does -// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. -__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, 10000000;\n\t" - "@P1 bra.uni DONE_RELAXED;\n\t" - "bra.uni LAB_WAIT_RELAXED;\n\t" - "DONE_RELAXED:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Exact source ports may request the PTX suspendTimeHint operand explicitly. -// The hint is expressed in nanoseconds and is kept separate from the canonical -// no-hint CTA helper so unrelated schedules retain their existing retry path. -__device__ __forceinline__ void mbarrier_wait_suspend( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_SUSPEND:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_SUSPEND;\n\t" - "bra.uni LAB_WAIT_SUSPEND;\n\t" - "DONE_SUSPEND:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE_CLUSTER;\n\t" - "bra.uni LAB_WAIT_CLUSTER;\n\t" - "DONE_CLUSTER:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - ".reg .u32 WAIT_ADDR;\n\t" - "mov.u32 WAIT_ADDR, %0;\n\t" - "LAB_WAIT_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [WAIT_ADDR], %1, %2;\n\t" - "@P1 bra.uni DONE_HINT;\n\t" - "bra.uni LAB_WAIT_HINT;\n\t" - "DONE_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -// Exact unqualified CTA wait used by source schedules whose PTX intentionally -// omits the acquire qualifier while retaining a typed suspendTimeHint operand. -__device__ __forceinline__ void mbarrier_wait_relaxed_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED_HINT:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra DONE_RELAXED_HINT;\n\t" - "bra LAB_WAIT_RELAXED_HINT;\n\t" - "DONE_RELAXED_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); -} - -__device__ __forceinline__ void mbarrier_wait_cluster_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" - "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" - "DONE_CLUSTER_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_suspend( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait_cluster(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); - } -} - - -__device__ __forceinline__ void tcgen05_mma_f8f6f4_cta2( - int taddr, uint64_t a_desc, uint64_t b_desc, - uint32_t i_desc, int enable_input_d) { - asm volatile( - "{\n\t" - ".reg .pred p;\n\t" - ".reg .b32 m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - "setp.ne.b32 p, %4, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\t" - "mov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "tcgen05.mma.cta_group::2.kind::f8f6f4 [%0], %1, %2, %3, {m0, m1, m2, m3, m4, m5, m6, m7}, p;\n\t" - "}\n" - :: "r"(taddr), "l"(a_desc), "l"(b_desc), - "r"(i_desc), "r"(enable_input_d) - : "memory"); -} - - -__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { - return (x & 0x3FFFFULL) >> 4ULL; -} - - -union MmaSmemDesc { - uint64_t u64; - uint32_t u32[2]; -}; - -__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { - MmaSmemDesc tmp; - tmp.u64 = smem_desc; - tmp.u32[0] += offset; - smem_desc = tmp.u64; -} - - -__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { - asm volatile( - "{\n\t" - ".reg .pred leader;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" - ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" - "}\n" - :: "r"(mbar_addr), "h"(cta_mask) : "memory"); -} - - -__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { - asm volatile( - "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" - :: "r"(mbar_addr), "r"(bytes) : "memory"); -} - - -__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7," - " %8, %9, %10, %11, %12, %13, %14, %15," - " %16, %17, %18, %19, %20, %21, %22, %23," - " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), - "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), - "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), - "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), - "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), - "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), - "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x16.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7," - " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), - "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), - "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ void tmem_ld_x4(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x4.b32" - " {%0, %1, %2, %3}, [%4];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ void tmem_st_x16_f32(int tmem_addr, const float* src) { - asm volatile( - "tcgen05.st.sync.aligned.32x32b.x16.b32" - " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," - " %9, %10, %11, %12, %13, %14, %15, %16};" - :: "r"(tmem_addr), - "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), - "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), - "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), - "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15])); -} - - -__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { - asm volatile( - "tcgen05.st.sync.aligned.32x32b.x4.b32" - " [%0], {%1, %2, %3, %4};" - :: "r"(tmem_addr), - "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); -} - - -__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { - uint32_t addr; - asm("{\n\t" - ".reg .u64 u64addr;\n\t" - "cvta.to.shared.u64 u64addr, %1;\n\t" - "cvt.u32.u64 %0, u64addr;\n\t" - "}\n" : "=r"(addr) : "l"(ptr)); - return addr; -} - - -__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { - uint32_t remote; - asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" - : "=r"(remote) : "r"(local_addr), "r"(rank)); - return remote; -} - - -__device__ __forceinline__ float approx_exp2(float x) { - float y; - asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float approx_rcp(float x) { - float y; - asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float max_noftz(float a, float b) { - float c; - asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); - return c; -} - - -__device__ __forceinline__ float warp_reduce_max(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); - return val; -} - - -__device__ __forceinline__ float warp_reduce_sum(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val += __shfl_xor_sync(0xFFFFFFFF, val, offset); - return val; -} - - -__device__ __forceinline__ float row_max_reduce(float2 acc) { - return max_noftz(acc.x, acc.y); -} - - -__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { - #pragma unroll - for (int j = 0; j < 16; j++) { - if (j % 2 == 0) - acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); - else - acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); - } -} - - -__device__ __forceinline__ float2 ex2_emulation_f32x2_value(float2 value) { - const float c0 = 1.0f, c1 = 0.695146143436431884765625f; - const float c2 = 0.227564394474029541015625f, c3 = 0.077119089663028717041015625f; - const float magic = 12582912.0f; - float x0 = max_noftz(value.x, -127.0f), x1 = max_noftz(value.y, -127.0f); - float2 xc2 = make_float2(x0, x1), magic2 = make_float2(magic, magic); - float2 xr2; - asm("add.rm.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xr2) - : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&magic2)); - float2 c3_2 = make_float2(c3, c3), c2_2 = make_float2(c2, c2); - float2 c1_2 = make_float2(c1, c1), c0_2 = make_float2(c0, c0); - float2 xrb2, xfrac2; - asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xrb2) - : "l"(*(unsigned long long*)&xr2), "l"(*(unsigned long long*)&magic2)); - asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xfrac2) - : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&xrb2)); - float2 poly2; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&c3_2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c2_2)); - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c1_2)); - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c0_2)); - int x0r_i, x1r_i, p0_i, p1_i; - asm("mov.b64 {%0, %1}, %2;" : "=r"(x0r_i), "=r"(x1r_i) : "l"(*(unsigned long long*)&xr2)); - asm("mov.b64 {%0, %1}, %2;" : "=r"(p0_i), "=r"(p1_i) : "l"(*(unsigned long long*)&poly2)); - float r0, r1; - asm("mov.b32 %0, %1;" : "=f"(r0) : "r"((x0r_i << 23) + p0_i)); - asm("mov.b32 %0, %1;" : "=f"(r1) : "r"((x1r_i << 23) + p1_i)); - return make_float2(r0, r1); -} - -__device__ __forceinline__ void ex2_emulation_f32x2(float* x0_ptr, float* x1_ptr) { - float2 result = ex2_emulation_f32x2_value(make_float2(*x0_ptr, *x1_ptr)); - *x0_ptr = result.x; *x1_ptr = result.y; -} - -__device__ __forceinline__ void softmax_frag_exp2_cast( - float* sv, uint32_t* pv, int use_emu) -{ - #pragma unroll - for (int j = 0; j < 16; j++) { - if (use_emu && j >= 12) - ex2_emulation_f32x2(&sv[j*2], &sv[j*2+1]); - else { - sv[j*2] = approx_exp2(sv[j*2]); - sv[j*2+1] = approx_exp2(sv[j*2+1]); - } - } - #pragma unroll - for (int j = 0; j < 16; j++) { - __nv_bfloat162 bf = __float22bfloat162_rn({sv[j*2], sv[j*2+1]}); - pv[j] = reinterpret_cast(bf); - } -} - - - -__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { - const float2* sv2 = reinterpret_cast(sv); - #pragma unroll - for (int j = 0; j < 16; j++) { - asm("add.f32x2 %0, %1, %2;" - : "+l"(reinterpret_cast(*acc)) - : "l"(reinterpret_cast(*acc)), - "l"(reinterpret_cast(sv2[j]))); - } -} - - -__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { - asm("mul.rn.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { - asm("mul.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { - asm("add.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { - asm("add.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { - asm("sub.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { - asm("sub.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("add.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { - float2 r; - asm("sub.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("sub.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ void fma_scale_x32( - float* sv, const float2* scale2, const float2* neg_max2) -{ - float2* sv_2 = reinterpret_cast(sv); - #pragma unroll - for (int j = 0; j < 16; j++) - fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); -} - -__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("mul.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - - -// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) - -__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("add.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("add.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("add.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("add.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("add.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("add.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("add.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - - -__device__ __forceinline__ void fence_async_shared() { - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); -} - - -__device__ __forceinline__ uint64_t make_smem_desc(int addr) { - const int SBO = 1024; - return desc_encode(addr) - | (desc_encode(SBO) << 32ULL) - | (1ULL << 46ULL) - | (2ULL << 61ULL); -} - - -__device__ __forceinline__ void tma_2d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3}], [%4];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_4d_gmem2smem_cta2( - int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.4d.shared::cluster.global" - ".mbarrier::complete_tx::bytes.cta_group::2" - " [%0], [%1, {%2, %3, %4, %5}], [%6];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr) { - // Canonical .shared::cta form for non-multicast gather4, matching - // trtllm-gen / cuda_ptx and the PTX ISA qualifier order - // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, - // .shared::cluster is reserved for the multicast variant (ctaMask). - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem_cta2( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr) { - // Canonical .shared::cta form; see tma_gather4_gmem2smem above. - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" - ".mbarrier::complete_tx::bytes.cta_group::2" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem_mc( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr, unsigned short cta_mask) { - // Multicast variant: the PTX grammar ties the .shared::cluster - // destination to .multicast::cluster + ctaMask (cf. cuda_ptx / - // SM100_TMA_LOAD_MULTICAST_2D_GATHER4 in CUTLASS). - asm volatile( - "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" - ".mbarrier::complete_tx::bytes.multicast::cluster" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr), "h"(cta_mask) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem_mc_cta2( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr, unsigned short cta_mask) { - // Multicast + cta_group::2 variant; see tma_gather4_gmem2smem_mc. - asm volatile( - "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" - ".mbarrier::complete_tx::bytes.multicast::cluster.cta_group::2" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7], %8;" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr), "h"(cta_mask) : "memory"); -} - - -__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { - asm volatile( - "{\n\t" - ".reg .b16 lo, hi;\n\t" - "mov.b32 {lo, hi}, %1;\n\t" - "tcgen05.commit.cta_group::2.mbarrier::arrive::one" - ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" - "}\n" - :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); -} - - -__device__ __forceinline__ void tmem_st_x8_u32(int addr, const uint32_t* src) { - asm volatile( - "tcgen05.st.sync.aligned.32x32b.x8.b32" - " [%0], {%1,%2,%3,%4,%5,%6,%7,%8};" - :: "r"(addr), - "r"(src[0]), "r"(src[1]), "r"(src[2]), "r"(src[3]), - "r"(src[4]), "r"(src[5]), "r"(src[6]), "r"(src[7])); -} - - -__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { - uint32_t result; - asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" - : "=r"(result) : "r"(val)); - return result; -} - -extern "C" { - -__global__ __launch_bounds__(512, 1) __cluster_dims__(2,1,1) void -kernel_cake_dsv4_c20607c57cc11784015e(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int num_query_tokens, int sparse_topk, int has_sinks, int total_work_items) -{ - const int tid = threadIdx.x; - const int warp = make_warp_uniform(tid / 32); - const int lane = tid % 32; - - extern __shared__ __align__(1024) char smem_raw[]; - int smem; - smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); - - const int mbar_base = smem; - #define q_full_addr (mbar_base + 0) - #define q_empty_addr (mbar_base + 8) - #define k_full_addr (mbar_base + 16) - #define k_empty_addr (mbar_base + 32) - #define v_full_addr (mbar_base + 48) - #define v_empty_addr (mbar_base + 64) - #define index_full_addr (mbar_base + 80) - #define index_empty_addr (mbar_base + 128) - #define s_full_addr (mbar_base + 176) - #define p_full_addr (mbar_base + 192) - #define corr_done_addr (mbar_base + 208) - #define stats_addr (mbar_base + 224) - #define sum_ready_addr (mbar_base + 240) - #define o_done_addr (mbar_base + 248) - #define pv_done_addr (mbar_base + 256) - #define s_seeded_addr (mbar_base + 264) - #define q_pair_ready_addr (mbar_base + 272) - #define kv_pair_ready_addr (mbar_base + 280) - #define pv_pair_ready_addr (mbar_base + 288) - #define tmem_dealloc_addr (mbar_base + 304) - #define tmem_dealloc_peer_addr (mbar_base + 312) - - const int bid = blockIdx.x; - const int num_bids = gridDim.x; - const unsigned int clusters_x = gridDim.x / 2; - const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; - const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; - - int cta_rank; - asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); - - // Kernel setup ops - uint8_t* smem_q = reinterpret_cast(smem_raw + 1024); - const int smem_q_addr = smem + 1024; - uint8_t* smem_kv = reinterpret_cast(smem_raw + 33792); - const int smem_kv_addr = smem + 33792; - uint8_t* smem_v = reinterpret_cast(smem_raw + 33792); - const int smem_v_addr = smem + 33792; - uint8_t* smem_q_full = reinterpret_cast(smem_raw + 1024); - const int smem_q_full_addr = smem + 1024; - uint8_t* smem_k_full = reinterpret_cast(smem_raw + 33792); - const int smem_k_full_addr = smem + 33792; - uint8_t* smem_v_full = reinterpret_cast(smem_raw + 99328); - const int smem_v_full_addr = smem + 99328; - float* smem_stats_max = reinterpret_cast(smem_raw + 164864); - const int smem_stats_max_addr = smem + 164864; - float* smem_stats_sum = reinterpret_cast(smem_raw + 165888); - const int smem_stats_sum_addr = smem + 165888; - float* smem_stats_final_max = reinterpret_cast(smem_raw + 166400); - const int smem_stats_final_max_addr = smem + 166400; - float* smem_softmax_warp_pair_exchange = reinterpret_cast(smem_raw + 166912); - const int smem_softmax_warp_pair_exchange_addr = smem + 166912; - float* smem_corr_warp_pair_exchange = reinterpret_cast(smem_raw + 167936); - const int smem_corr_warp_pair_exchange_addr = smem + 167936; - int* smem_sparse_indices = reinterpret_cast(smem_raw + 168448); - const int smem_sparse_indices_addr = smem + 168448; - uint8_t* smem_p_fp8 = reinterpret_cast(smem_raw + 174592); - const int smem_p_fp8_addr = smem + 174592; - { - asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); - asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); - } - - // Mbarrier init (21 pipeline groups, 0 ordered-sequence groups, 40 barriers) - // Mbarriers at smem_raw[0..320) - - if (warp == 12) { - uint32_t leader = elect_sync(); - if (leader) { - // q_full: 1 barriers, init_count=2 - mbarrier_init(smem + 0, 2); - // q_empty: 1 barriers, init_count=1 - mbarrier_init(smem + 8, 1); - // --- pipeline 'k_pipe' --- - // k_full: 2 barriers, init_count=2 - mbarrier_init(smem + 16, 2); - mbarrier_init(smem + 24, 2); - // k_empty: 2 barriers, init_count=1 - mbarrier_init(smem + 32, 1); - mbarrier_init(smem + 40, 1); - // --- pipeline 'v_pipe' --- - // v_full: 2 barriers, init_count=2 - mbarrier_init(smem + 48, 2); - mbarrier_init(smem + 56, 2); - // v_empty: 2 barriers, init_count=1 - mbarrier_init(smem + 64, 1); - mbarrier_init(smem + 72, 1); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - if (warp == 9) { - uint32_t leader = elect_sync(); - if (leader) { - // --- pipeline 'index_pipe' --- - // index_full: 6 barriers, init_count=32 - mbarrier_init(smem + 80, 32); - mbarrier_init(smem + 88, 32); - mbarrier_init(smem + 96, 32); - mbarrier_init(smem + 104, 32); - mbarrier_init(smem + 112, 32); - mbarrier_init(smem + 120, 32); - // index_empty: 6 barriers, init_count=192 - mbarrier_init(smem + 128, 192); - mbarrier_init(smem + 136, 192); - mbarrier_init(smem + 144, 192); - mbarrier_init(smem + 152, 192); - mbarrier_init(smem + 160, 192); - mbarrier_init(smem + 168, 192); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - if (warp == 0) { - uint32_t leader = elect_sync(); - if (leader) { - // s_full: stages (0,), init_count=1 - mbarrier_init(smem + 176, 1); - // p_full: stages (0,), init_count=64 - mbarrier_init(smem + 192, 64); - // o_done: 1 barriers, init_count=1 - mbarrier_init(smem + 248, 1); - // pv_done: 1 barriers, init_count=1 - mbarrier_init(smem + 256, 1); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - if (warp == 1) { - uint32_t leader = elect_sync(); - if (leader) { - // s_full: stages (1,), init_count=1 - mbarrier_init(smem + 184, 1); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - if (warp == 2) { - uint32_t leader = elect_sync(); - if (leader) { - // p_full: stages (1,), init_count=64 - mbarrier_init(smem + 200, 64); - // s_seeded: 1 barriers, init_count=256 - mbarrier_init(smem + 264, 256); - // q_pair_ready: 1 barriers, init_count=64 - mbarrier_init(smem + 272, 64); - // --- pipeline 'kv_pipe' --- - // kv_pair_ready: 1 barriers, init_count=64 - mbarrier_init(smem + 280, 64); - // pv_pair_ready: stages (0,), init_count=64 - mbarrier_init(smem + 288, 64); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - if (warp == 4) { - uint32_t leader = elect_sync(); - if (leader) { - // corr_done: stages (0,), init_count=64 - mbarrier_init(smem + 208, 64); - // stats: stages (0,), init_count=64 - mbarrier_init(smem + 224, 64); - // sum_ready: 1 barriers, init_count=64 - mbarrier_init(smem + 240, 64); - // pv_pair_ready: stages (1,), init_count=64 - mbarrier_init(smem + 296, 64); - // tmem_dealloc: 1 barriers, init_count=288 - mbarrier_init(smem + 304, 288); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - if (warp == 6) { - uint32_t leader = elect_sync(); - if (leader) { - // corr_done: stages (1,), init_count=64 - mbarrier_init(smem + 216, 64); - // stats: stages (1,), init_count=64 - mbarrier_init(smem + 232, 64); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - if (warp == 5) { - uint32_t leader = elect_sync(); - if (leader) { - // tmem_dealloc_peer: 1 barriers, init_count=32 - mbarrier_init(smem + 312, 32); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - - __syncwarp(); - - asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); - asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); - - // TMEM alloc (512 columns, 512 used) - volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 320); - if (warp == 0) { - int _tmem_hold = smem + 320; - asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); - __syncwarp(); - asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); - } - - __syncthreads(); - asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); - asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); - - const int taddr = tmem_addr_storage[0]; - - // Kernel post-init ops - const int tmem_tmem_scratch = taddr; - - // ---- Ordered hardware-WG register redistribution ---- - // Inc phase consumes the registers released above. - if (warp >= 0 && warp <= 11) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 136;"); - } - - // ---- Role: index_warp ---- - if (warp == 9) { - { // index_warp_main - const int index_dummy = 0; - unsigned int _phase_index_empty = 1; - unsigned int _phase_q_full_0 = 0; - { - int all_num_kv_tiles = (sparse_topk + 128 - 1) / 128; - unsigned int index_prod_stage = 0; - #pragma unroll - for (unsigned int static_work = 0; static_work < 1; static_work++) { - int split_idx = 0; - int query_idx = bid / 2; - int tiles_per_split = all_num_kv_tiles + 1 - 1; - int first_tile = split_idx * tiles_per_split; - int num_index_passes = (tiles_per_split + 1) / 2; - #pragma unroll 1 - for (int index_pass = 0; index_pass < num_index_passes; index_pass++) { - mbarrier_wait(index_empty_addr + (index_prod_stage) * 8, _phase_index_empty); - int index_stage_base = smem_sparse_indices_addr + index_prod_stage * 1024; - int stage_row_base = index_pass * 256; - #pragma unroll - for (int index_half = 0; index_half < 2; index_half++) { - int index_offset = index_half * 128 + lane * 4; - int tile_col_base = first_tile * 128 + stage_row_base + index_half * 128; - int* row_tile_ptr = ((tile_col_base < 128) ? swa_indices + (query_idx * swa_index_stride) : compressed_indices + (query_idx * compressed_index_stride + (tile_col_base - 128))); - int lane_col = tile_col_base + lane * 4; - if (lane_col + 4 <= sparse_topk) { - asm volatile("cp.async.cg.shared::cta.global.L2::128B [%0], [%1], 16;" - :: "r"(index_stage_base + index_offset * 4), "l"(row_tile_ptr + (lane * 4))); - } else { - int tail_rows[4]; - #pragma unroll - for (int row_i = 0; row_i < 4; row_i++) { - tail_rows[row_i] = -1; - } - #pragma unroll - for (int row_i_1 = 0; row_i_1 < 4; row_i_1++) { - if (lane_col + row_i_1 < sparse_topk) { - tail_rows[row_i_1] = row_tile_ptr[lane * 4 + row_i_1]; - } - } - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(smem_sparse_indices_addr + ((unsigned int)index_stage_base - smem_sparse_indices_addr + (unsigned int)(index_offset * 4))), "r"(tail_rows[0]), "r"(tail_rows[1]), "r"(tail_rows[2]), "r"(tail_rows[3]) : "memory"); - } - } - asm volatile("cp.async.commit_group;"); - asm volatile("cp.async.wait_group 0;"); - mbarrier_arrive(index_full_addr + (index_prod_stage) * 8); - index_prod_stage += 1; - if (index_prod_stage == 6) { index_prod_stage = 0; _phase_index_empty ^= 1; } - } - } - } - } - } - // ---- Role: load_warp ---- - if (warp >= 12 && warp <= 15) { - asm volatile("setmaxnreg.dec.sync.aligned.u32 96;"); - { // load_warp_main - const int wg2_dummy = 0; - const int load_warp_rank = warp - ((1) ? 12 : 9); - int all_num_kv_tiles_1 = (sparse_topk + 128 - 1) / 128; - unsigned int load_k_stage = 0; - unsigned int load_v_stage = 0; - unsigned int load_kv_stage = 0; - unsigned int load_k_index_stage = 0; - unsigned int load_v_index_stage = 0; - int k_cta_offset = bid % 2 * 64; - unsigned int _phase_q_empty_0 = 1; - unsigned int _phase_index_full = 0; - unsigned int _phase_k_empty = 1; - unsigned int _phase_k_empty_1 = 1; - unsigned int _phase_v_empty = 1; - #pragma unroll - for (unsigned int work_idx = 0; work_idx < 1; work_idx++) { - int split_idx_1 = 0; - int query_idx_1 = bid / 2 >> 1; - int v_chunk = bid / 2 & 1; - { - split_idx_1 = 0; - query_idx_1 = bid / 2; - v_chunk = bid % 2; - } - int tiles_per_split_1 = all_num_kv_tiles_1 + 1 - 1; - int first_tile_1 = split_idx_1 * tiles_per_split_1; - int num_kv_tiles = tiles_per_split_1; - int sparse_extent = num_kv_tiles * 128; - mbarrier_wait(q_empty_addr, _phase_q_empty_0); - _phase_q_empty_0 ^= 1; - if (load_warp_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); - #pragma unroll - for (int q_stage = 0; q_stage < 4; q_stage++) { - asm volatile( - "cp.async.bulk.tensor.4d.shared::cluster.global.mbarrier::complete_tx::bytes.multicast::cluster.cta_group::2" - " [%0], [%1, {%2, %3, %4, %5}], [%6], %7;" - :: "r"(smem_q_full_addr + (unsigned int)(q_stage * 8192)), "l"((&tmap_q)), "r"(0), "r"(bid % 2 * 64), "r"(q_stage), "r"(query_idx_1), - "r"(((q_full_addr) & 0xFEFFFFFF)), "h"((uint16_t)(1 << bid % 2)) : "memory"); - } - } - } - { - mbarrier_wait(index_full_addr + (load_k_index_stage) * 8, _phase_index_full); - asm volatile("tcgen05.fence::after_thread_sync;"); - } - int k_index_stage_base = smem_sparse_indices_addr + load_k_index_stage * 1024; - { - mbarrier_wait(k_empty_addr + (load_k_stage) * 8, _phase_k_empty); - if (load_warp_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); - } - } - int k_dst = smem_k_full_addr + load_k_stage * 32768 + (unsigned int)(load_warp_rank * 8192); - int group = lane; - if (group < 16) { - int group_offset = k_cta_offset + group * 4; - int raw_rows[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows[0])), "=r"(*reinterpret_cast(&raw_rows[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows[(0) + 3])) - : "r"(k_index_stage_base + group_offset * 4)); - int row0 = ((raw_rows[0] >= 0) ? raw_rows[0] : 0); - int row1 = ((raw_rows[1] >= 0) ? raw_rows[1] : 0); - int row2 = ((raw_rows[2] >= 0) ? raw_rows[2] : 0); - int row3 = ((raw_rows[3] >= 0) ? raw_rows[3] : 0); - if (first_tile_1 == 0) { - tma_gather4_gmem2smem_mc_cta2(k_dst + group * 512, (&tmap_swa_kv), load_warp_rank * 128, row0, row1, row2, row3, ((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } else { - tma_gather4_gmem2smem_mc_cta2(k_dst + group * 512, (&tmap_compressed_kv), load_warp_rank * 128, row0, row1, row2, row3, ((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } - } - load_k_stage += 1; - if (load_k_stage == 2) { load_k_stage = 0; _phase_k_empty ^= 1; } - } - #pragma unroll 1 - for (int tile = 1; tile < num_kv_tiles; tile++) { - if ((1 & (int)((tile & 1) == 0)) != 0) { - mbarrier_wait(index_full_addr + (load_k_index_stage) * 8, _phase_index_full); - asm volatile("tcgen05.fence::after_thread_sync;"); - } - k_index_stage_base = smem_sparse_indices_addr + load_k_index_stage * 1024; - { - mbarrier_wait(k_empty_addr + (load_k_stage) * 8, _phase_k_empty); - if (load_warp_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); - } - } - int k_dst_1 = smem_k_full_addr + load_k_stage * 32768 + (unsigned int)(load_warp_rank * 8192); - int group_1 = lane; - if (group_1 < 16) { - int group_offset_1 = (tile & 1) * 128 + k_cta_offset + group_1 * 4; - int raw_rows_1[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_1[0])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 3])) - : "r"(k_index_stage_base + group_offset_1 * 4)); - int row0_1 = ((raw_rows_1[0] >= 0) ? raw_rows_1[0] : 0); - int row1_1 = ((raw_rows_1[1] >= 0) ? raw_rows_1[1] : 0); - int row2_1 = ((raw_rows_1[2] >= 0) ? raw_rows_1[2] : 0); - int row3_1 = ((raw_rows_1[3] >= 0) ? raw_rows_1[3] : 0); - tma_gather4_gmem2smem_mc_cta2(k_dst_1 + group_1 * 512, (&tmap_compressed_kv), load_warp_rank * 128, row0_1, row1_1, row2_1, row3_1, ((k_full_addr + (load_k_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } - load_k_stage += 1; - if (load_k_stage == 2) { load_k_stage = 0; _phase_k_empty ^= 1; } - if ((tile & 1) != 0) { - load_k_index_stage += 1; - if (load_k_index_stage == 6) { load_k_index_stage = 0; _phase_index_full ^= 1; } - } - } - int prev_tile = tile - 1; - int global_prev_tile = first_tile_1 + prev_tile; - int v_index_stage_base = smem_sparse_indices_addr + load_v_index_stage * 1024; - { - mbarrier_wait(v_empty_addr + (load_v_stage) * 8, _phase_v_empty); - if (load_warp_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); - } - } - if (lane < 16) { - int v_dst = smem_v_full_addr + load_v_stage * 32768 + (unsigned int)(load_warp_rank / 2 * 16384); - int v_col = v_chunk * 256 + load_warp_rank / 2 * 128; - int group_2 = (load_warp_rank & 1) * 16 + lane; - int group_offset_2 = (prev_tile & 1) * 128 + group_2 * 4; - int raw_rows_2[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_2[0])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_2[(0) + 3])) - : "r"(v_index_stage_base + group_offset_2 * 4)); - int row0_2 = ((raw_rows_2[0] >= 0) ? raw_rows_2[0] : 0); - int row1_2 = ((raw_rows_2[1] >= 0) ? raw_rows_2[1] : 0); - int row2_2 = ((raw_rows_2[2] >= 0) ? raw_rows_2[2] : 0); - int row3_2 = ((raw_rows_2[3] >= 0) ? raw_rows_2[3] : 0); - if (global_prev_tile == 0) { - tma_gather4_gmem2smem_mc_cta2(v_dst + group_2 * 512, (&tmap_swa_kv), v_col, row0_2, row1_2, row2_2, row3_2, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } else { - tma_gather4_gmem2smem_mc_cta2(v_dst + group_2 * 512, (&tmap_compressed_kv), v_col, row0_2, row1_2, row2_2, row3_2, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } - } - load_v_stage += 1; - if (load_v_stage == 2) { load_v_stage = 0; _phase_v_empty ^= 1; } - if ((prev_tile & 1) != 0) { - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(index_empty_addr + (load_v_index_stage) * 8); - load_v_index_stage = load_v_index_stage + 1; - if (load_v_index_stage == 6) { - load_v_index_stage = 0; - } - } - } - } - int last_tile = num_kv_tiles - 1; - int global_last_tile = first_tile_1 + last_tile; - { - mbarrier_wait(v_empty_addr + (load_v_stage) * 8, _phase_v_empty); - if (load_warp_rank == 0) { - if (elect_sync()) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" - :: "r"((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); - } - } - if (lane < 16) { - int v_dst_1 = smem_v_full_addr + load_v_stage * 32768 + (unsigned int)(load_warp_rank / 2 * 16384); - int v_col_1 = v_chunk * 256 + load_warp_rank / 2 * 128; - int group_3 = (load_warp_rank & 1) * 16 + lane; - int group_offset_3 = (last_tile & 1) * 128 + group_3 * 4; - int raw_rows_3[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_3[0])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_3[(0) + 3])) - : "r"(smem_sparse_indices_addr + load_v_index_stage * 1024 + (unsigned int)(group_offset_3 * 4))); - int row0_3 = ((raw_rows_3[0] >= 0) ? raw_rows_3[0] : 0); - int row1_3 = ((raw_rows_3[1] >= 0) ? raw_rows_3[1] : 0); - int row2_3 = ((raw_rows_3[2] >= 0) ? raw_rows_3[2] : 0); - int row3_3 = ((raw_rows_3[3] >= 0) ? raw_rows_3[3] : 0); - if (global_last_tile == 0) { - tma_gather4_gmem2smem_mc_cta2(v_dst_1 + group_3 * 512, (&tmap_swa_kv), v_col_1, row0_3, row1_3, row2_3, row3_3, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } else { - tma_gather4_gmem2smem_mc_cta2(v_dst_1 + group_3 * 512, (&tmap_compressed_kv), v_col_1, row0_3, row1_3, row2_3, row3_3, ((v_full_addr + (load_v_stage) * 8) & 0xFEFFFFFF), 1 << bid % 2); - } - } - load_v_stage += 1; - if (load_v_stage == 2) { load_v_stage = 0; _phase_v_empty ^= 1; } - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(index_empty_addr + (load_v_index_stage) * 8); - load_v_index_stage = load_v_index_stage + 1; - if (load_v_index_stage == 6) { - load_v_index_stage = 0; - } - } - } - mbarrier_arrive(tmem_dealloc_addr); - } - } - // ---- Role: softmax_wg ---- - if (warp == 0 || warp == 2) { - { // softmax_wg_main - float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; - const int wg_dummy_inc = 0; - int all_num_kv_tiles_2 = (sparse_topk + 128 - 1) / 128; - const int tmem_row_base = ((1) ? warp % 2 * 32 : warp % 4 * 32); - const int tmem_score_row_base = ((1) ? (int)(warp % 4 * 32) : tmem_row_base); - const int n_half = ((1) ? (int)(warp % 4 / 2) : 0); - const int my_row = tmem_row_base + lane; - const int stats_row = n_half * 64 + my_row; - int softmax_tile_cursor = 0; - unsigned int softmax_index_stage = 0; - unsigned int _phase_q_full_0_1 = 0; - #pragma unroll - for (unsigned int work_idx_1 = 0; work_idx_1 < 1; work_idx_1++) { - int split_idx_2 = 0; - int query_idx_2 = bid / 2 >> 1; - { - split_idx_2 = 0; - query_idx_2 = bid / 2; - } - int tiles_per_split_2 = all_num_kv_tiles_2 + 1 - 1; - int first_tile_2 = split_idx_2 * tiles_per_split_2; - int num_kv_tiles_1 = tiles_per_split_2; - int _max_0 = ((sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) : (0)); - int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); - int active_topk = _min_0; - float row_max_val = -CAKE_INF; - float row_sum_val = 0.0f; - int sink_head = ((1) ? bid % 2 * 64 + my_row : my_row); - if (has_sinks != 0 && sink_head < num_heads && split_idx_2 == 0) { - row_max_val = sinks[sink_head] * 1.4426950408889634f / softmax_scale_log2; - row_sum_val = 1.0f; - } - #pragma unroll 1 - for (int tile_1 = 0; tile_1 < num_kv_tiles_1; tile_1++) { - int pipeline_tile = softmax_tile_cursor + tile_1; - int phase = pipeline_tile & 1; - int s_wait_phase = pipeline_tile >> 1 & 1; - mbarrier_wait(s_full_addr + (phase) * 8, s_wait_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - int s_off = ((phase != 0) ? 128 : 0); - int s_base = taddr + (unsigned int)s_off + (unsigned int)(tmem_score_row_base << 16); - float new_max = row_max_val; - int valid_sparse_cols = ((active_topk < sparse_topk) ? active_topk : sparse_topk); - valid_sparse_cols = valid_sparse_cols - (first_tile_2 + tile_1) * 128 - n_half * 64; - if (valid_sparse_cols < 0) { - valid_sparse_cols = 0; - } - if (valid_sparse_cols > 64) { - valid_sparse_cols = 64; - } - int index_tile_addr = smem_sparse_indices_addr + softmax_index_stage * 1024 + (unsigned int)((tile_1 & 1) * 512); - int index_lane_addr = index_tile_addr + (n_half * 64 + lane) * 4; - int staged_index[2]; - unsigned int invalid_cols[2]; - unsigned int any_invalid = 0; - #pragma unroll - for (int w = 0; w < 2; w++) { - asm volatile("ld.shared.b32 %0, [%1];" : "=r"(*reinterpret_cast(&staged_index[w])) : "r"(index_lane_addr + w * 128)); - unsigned int _vote_0 = __ballot_sync(0xFFFFFFFF, staged_index[w] < 0); - invalid_cols[w] = _vote_0; - any_invalid = any_invalid | invalid_cols[w]; - } - if ((tile_1 & 1) != 0 || tile_1 + 1 == num_kv_tiles_1) { - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(index_empty_addr + (softmax_index_stage) * 8); - softmax_index_stage = softmax_index_stage + 1; - if (softmax_index_stage == 6) { - softmax_index_stage = 0; - } - } - float _tmem_load_0[4]; - tmem_ld_x4(&_tmem_load_0[0], s_base); - float sv_split[64]; - float half_max = -CAKE_INF; - if (valid_sparse_cols == 64 && any_invalid == 0) { - #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 1000) - #error "TmemLoadRed requires tcgen05.ld.red support (sm_103/sm_101-sm_110 family), not sm_100" - #endif - asm volatile( - "tcgen05.ld.red.sync.aligned.32x32b.x64.max.f32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, %64, [%65];" - : "=f"(sv_split[0]), "=f"(sv_split[1]), "=f"(sv_split[2]), "=f"(sv_split[3]), "=f"(sv_split[4]), "=f"(sv_split[5]), "=f"(sv_split[6]), "=f"(sv_split[7]), "=f"(sv_split[8]), "=f"(sv_split[9]), "=f"(sv_split[10]), "=f"(sv_split[11]), "=f"(sv_split[12]), "=f"(sv_split[13]), "=f"(sv_split[14]), "=f"(sv_split[15]), "=f"(sv_split[16]), "=f"(sv_split[17]), "=f"(sv_split[18]), "=f"(sv_split[19]), "=f"(sv_split[20]), "=f"(sv_split[21]), "=f"(sv_split[22]), "=f"(sv_split[23]), "=f"(sv_split[24]), "=f"(sv_split[25]), "=f"(sv_split[26]), "=f"(sv_split[27]), "=f"(sv_split[28]), "=f"(sv_split[29]), "=f"(sv_split[30]), "=f"(sv_split[31]), "=f"(sv_split[32]), "=f"(sv_split[33]), "=f"(sv_split[34]), "=f"(sv_split[35]), "=f"(sv_split[36]), "=f"(sv_split[37]), "=f"(sv_split[38]), "=f"(sv_split[39]), "=f"(sv_split[40]), "=f"(sv_split[41]), "=f"(sv_split[42]), "=f"(sv_split[43]), "=f"(sv_split[44]), "=f"(sv_split[45]), "=f"(sv_split[46]), "=f"(sv_split[47]), "=f"(sv_split[48]), "=f"(sv_split[49]), "=f"(sv_split[50]), "=f"(sv_split[51]), "=f"(sv_split[52]), "=f"(sv_split[53]), "=f"(sv_split[54]), "=f"(sv_split[55]), "=f"(sv_split[56]), "=f"(sv_split[57]), "=f"(sv_split[58]), "=f"(sv_split[59]), "=f"(sv_split[60]), "=f"(sv_split[61]), "=f"(sv_split[62]), "=f"(sv_split[63]), "=f"(half_max) - : "r"(s_base)); - asm volatile("tcgen05.wait::ld.sync.aligned;"); - } else { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(sv_split[0]), "=f"(sv_split[1]), "=f"(sv_split[2]), "=f"(sv_split[3]), "=f"(sv_split[4]), "=f"(sv_split[5]), "=f"(sv_split[6]), "=f"(sv_split[7]), "=f"(sv_split[8]), "=f"(sv_split[9]), "=f"(sv_split[10]), "=f"(sv_split[11]), "=f"(sv_split[12]), "=f"(sv_split[13]), "=f"(sv_split[14]), "=f"(sv_split[15]), "=f"(sv_split[16]), "=f"(sv_split[17]), "=f"(sv_split[18]), "=f"(sv_split[19]), "=f"(sv_split[20]), "=f"(sv_split[21]), "=f"(sv_split[22]), "=f"(sv_split[23]), "=f"(sv_split[24]), "=f"(sv_split[25]), "=f"(sv_split[26]), "=f"(sv_split[27]), "=f"(sv_split[28]), "=f"(sv_split[29]), "=f"(sv_split[30]), "=f"(sv_split[31]) - : "r"(s_base)); - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(sv_split[32]), "=f"(sv_split[33]), "=f"(sv_split[34]), "=f"(sv_split[35]), "=f"(sv_split[36]), "=f"(sv_split[37]), "=f"(sv_split[38]), "=f"(sv_split[39]), "=f"(sv_split[40]), "=f"(sv_split[41]), "=f"(sv_split[42]), "=f"(sv_split[43]), "=f"(sv_split[44]), "=f"(sv_split[45]), "=f"(sv_split[46]), "=f"(sv_split[47]), "=f"(sv_split[48]), "=f"(sv_split[49]), "=f"(sv_split[50]), "=f"(sv_split[51]), "=f"(sv_split[52]), "=f"(sv_split[53]), "=f"(sv_split[54]), "=f"(sv_split[55]), "=f"(sv_split[56]), "=f"(sv_split[57]), "=f"(sv_split[58]), "=f"(sv_split[59]), "=f"(sv_split[60]), "=f"(sv_split[61]), "=f"(sv_split[62]), "=f"(sv_split[63]) - : "r"(s_base + 32)); - uint32_t _slice_lo_mask_0; - { - int _lim_0 = valid_sparse_cols; - if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } - else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); - } - } - if (!(_slice_lo_mask_0 & (1u << 0))) sv_split[0] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 1))) sv_split[1] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 2))) sv_split[2] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 3))) sv_split[3] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 4))) sv_split[4] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 5))) sv_split[5] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 6))) sv_split[6] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 7))) sv_split[7] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 8))) sv_split[8] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 9))) sv_split[9] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 10))) sv_split[10] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 11))) sv_split[11] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 12))) sv_split[12] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 13))) sv_split[13] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 14))) sv_split[14] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 15))) sv_split[15] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 16))) sv_split[16] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 17))) sv_split[17] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 18))) sv_split[18] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 19))) sv_split[19] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 20))) sv_split[20] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 21))) sv_split[21] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 22))) sv_split[22] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 23))) sv_split[23] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 24))) sv_split[24] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 25))) sv_split[25] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 26))) sv_split[26] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 27))) sv_split[27] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 28))) sv_split[28] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 29))) sv_split[29] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 30))) sv_split[30] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 31))) sv_split[31] = -CAKE_INF; - uint32_t _slice_lo_mask_1; - { - int _lim_1 = valid_sparse_cols - 32; - if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } - else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); - } - } - if (!(_slice_lo_mask_1 & (1u << 0))) sv_split[32] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 1))) sv_split[33] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 2))) sv_split[34] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 3))) sv_split[35] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 4))) sv_split[36] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 5))) sv_split[37] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 6))) sv_split[38] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 7))) sv_split[39] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 8))) sv_split[40] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 9))) sv_split[41] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 10))) sv_split[42] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 11))) sv_split[43] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 12))) sv_split[44] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 13))) sv_split[45] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 14))) sv_split[46] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 15))) sv_split[47] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 16))) sv_split[48] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 17))) sv_split[49] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 18))) sv_split[50] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 19))) sv_split[51] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 20))) sv_split[52] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 21))) sv_split[53] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 22))) sv_split[54] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 23))) sv_split[55] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 24))) sv_split[56] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 25))) sv_split[57] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 26))) sv_split[58] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 27))) sv_split[59] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 28))) sv_split[60] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 29))) sv_split[61] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 30))) sv_split[62] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 31))) sv_split[63] = -CAKE_INF; - if (any_invalid != 0) { - #pragma unroll - for (int w_1 = 0; w_1 < 2; w_1++) { - unsigned int mask_word = invalid_cols[w_1]; - sv_split[w_1 * 32] = (((mask_word & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32]); - sv_split[w_1 * 32 + 1] = (((mask_word >> 1 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 1]); - sv_split[w_1 * 32 + 2] = (((mask_word >> 2 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 2]); - sv_split[w_1 * 32 + 3] = (((mask_word >> 3 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 3]); - sv_split[w_1 * 32 + 4] = (((mask_word >> 4 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 4]); - sv_split[w_1 * 32 + 5] = (((mask_word >> 5 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 5]); - sv_split[w_1 * 32 + 6] = (((mask_word >> 6 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 6]); - sv_split[w_1 * 32 + 7] = (((mask_word >> 7 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 7]); - sv_split[w_1 * 32 + 8] = (((mask_word >> 8 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 8]); - sv_split[w_1 * 32 + 9] = (((mask_word >> 9 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 9]); - sv_split[w_1 * 32 + 10] = (((mask_word >> 10 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 10]); - sv_split[w_1 * 32 + 11] = (((mask_word >> 11 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 11]); - sv_split[w_1 * 32 + 12] = (((mask_word >> 12 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 12]); - sv_split[w_1 * 32 + 13] = (((mask_word >> 13 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 13]); - sv_split[w_1 * 32 + 14] = (((mask_word >> 14 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 14]); - sv_split[w_1 * 32 + 15] = (((mask_word >> 15 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 15]); - sv_split[w_1 * 32 + 16] = (((mask_word >> 16 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 16]); - sv_split[w_1 * 32 + 17] = (((mask_word >> 17 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 17]); - sv_split[w_1 * 32 + 18] = (((mask_word >> 18 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 18]); - sv_split[w_1 * 32 + 19] = (((mask_word >> 19 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 19]); - sv_split[w_1 * 32 + 20] = (((mask_word >> 20 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 20]); - sv_split[w_1 * 32 + 21] = (((mask_word >> 21 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 21]); - sv_split[w_1 * 32 + 22] = (((mask_word >> 22 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 22]); - sv_split[w_1 * 32 + 23] = (((mask_word >> 23 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 23]); - sv_split[w_1 * 32 + 24] = (((mask_word >> 24 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 24]); - sv_split[w_1 * 32 + 25] = (((mask_word >> 25 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 25]); - sv_split[w_1 * 32 + 26] = (((mask_word >> 26 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 26]); - sv_split[w_1 * 32 + 27] = (((mask_word >> 27 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 27]); - sv_split[w_1 * 32 + 28] = (((mask_word >> 28 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 28]); - sv_split[w_1 * 32 + 29] = (((mask_word >> 29 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 29]); - sv_split[w_1 * 32 + 30] = (((mask_word >> 30 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 30]); - sv_split[w_1 * 32 + 31] = (((mask_word >> 31 & 1) != 0) ? -CAKE_INF : sv_split[w_1 * 32 + 31]); - } - } - float2 _reg_reduce_max2_2 = {-CAKE_INF, -CAKE_INF}; - row_max_x32_accum(&sv_split[0], _reg_reduce_max2_2); - row_max_x32_accum(&sv_split[32], _reg_reduce_max2_2); - float sv_split_max = row_max_reduce(_reg_reduce_max2_2); - half_max = sv_split_max; - } - { - float _max_1 = max_noftz(new_max, half_max); - new_max = _max_1; - smem_softmax_warp_pair_exchange[phase * 128 + stats_row] = new_max; - asm volatile("barrier.sync %0, 64;" :: "r"(2 + warp % 2) : "memory"); - float _max_2 = max_noftz(new_max, smem_softmax_warp_pair_exchange[phase * 128 + (stats_row ^ 64)]); - new_max = _max_2; - } - float _fma_0 = __fmaf_rn(row_max_val, softmax_scale_log2, (-new_max) * softmax_scale_log2); - float delta = _fma_0; - float _exp2_0 = approx_exp2(delta); - float exp_delta = _exp2_0; - float acc_scale = ((row_max_val > -CAKE_INF) ? exp_delta : 1.0f); - smem_stats_max[phase * 128 + stats_row] = acc_scale; - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(stats_addr + (phase) * 8); - row_max_val = new_max; - float safe_max = ((new_max == -CAKE_INF) ? 0.0f : new_max); - float max_scaled = safe_max * softmax_scale_log2; - float block_sum = 0.0f; - { - { - const float2 _fma_b2_3 = {softmax_scale_log2, softmax_scale_log2}; - const float2 _fma_c2_4 = {-max_scaled, -max_scaled}; - #pragma unroll - for (int _lf = 0; _lf < 32; _lf++) - fma_f32x2_inplace(&reinterpret_cast(sv_split)[_lf], _fma_b2_3, _fma_c2_4); - #pragma unroll - for (int _le = 0; _le < 64; _le++) { - sv_split[_le] = approx_exp2(sv_split[_le]); - } - { - const float2* _reg_reduce_src2_5 = reinterpret_cast(&sv_split[0]); - float2 _reg_reduce_sum2_5_0 = make_float2(0.0f, 0.0f); - float2 _reg_reduce_sum2_5_1 = make_float2(0.0f, 0.0f); - float2 _reg_reduce_sum2_5_2 = make_float2(0.0f, 0.0f); - float2 _reg_reduce_sum2_5_3 = make_float2(0.0f, 0.0f); - #pragma unroll - for (int _rr = 0; _rr < 8; _rr++) { - add_f32x2_inplace(&_reg_reduce_sum2_5_0, _reg_reduce_src2_5[_rr * 4 + 0]); - add_f32x2_inplace(&_reg_reduce_sum2_5_1, _reg_reduce_src2_5[_rr * 4 + 1]); - add_f32x2_inplace(&_reg_reduce_sum2_5_2, _reg_reduce_src2_5[_rr * 4 + 2]); - add_f32x2_inplace(&_reg_reduce_sum2_5_3, _reg_reduce_src2_5[_rr * 4 + 3]); - } - add_f32x2_inplace(&_reg_reduce_sum2_5_0, _reg_reduce_sum2_5_1); - add_f32x2_inplace(&_reg_reduce_sum2_5_2, _reg_reduce_sum2_5_3); - add_f32x2_inplace(&_reg_reduce_sum2_5_0, _reg_reduce_sum2_5_2); - float sv_split_sum = _reg_reduce_sum2_5_0.x + _reg_reduce_sum2_5_0.y; - block_sum = sv_split_sum; - } - uint32_t _fp8_0[16]; - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(sv_split[0]), "f"(sv_split[1]), - "f"(sv_split[2]), "f"(sv_split[3])); - _fp8_0[0] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(sv_split[4]), "f"(sv_split[5]), - "f"(sv_split[6]), "f"(sv_split[7])); - _fp8_0[1] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(sv_split[8]), "f"(sv_split[9]), - "f"(sv_split[10]), "f"(sv_split[11])); - _fp8_0[2] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(sv_split[12]), "f"(sv_split[13]), - "f"(sv_split[14]), "f"(sv_split[15])); - _fp8_0[3] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(sv_split[16]), "f"(sv_split[17]), - "f"(sv_split[18]), "f"(sv_split[19])); - _fp8_0[4] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(sv_split[20]), "f"(sv_split[21]), - "f"(sv_split[22]), "f"(sv_split[23])); - _fp8_0[5] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(sv_split[24]), "f"(sv_split[25]), - "f"(sv_split[26]), "f"(sv_split[27])); - _fp8_0[6] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(sv_split[28]), "f"(sv_split[29]), - "f"(sv_split[30]), "f"(sv_split[31])); - _fp8_0[7] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(sv_split[32]), "f"(sv_split[33]), - "f"(sv_split[34]), "f"(sv_split[35])); - _fp8_0[8] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(sv_split[36]), "f"(sv_split[37]), - "f"(sv_split[38]), "f"(sv_split[39])); - _fp8_0[9] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(sv_split[40]), "f"(sv_split[41]), - "f"(sv_split[42]), "f"(sv_split[43])); - _fp8_0[10] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(sv_split[44]), "f"(sv_split[45]), - "f"(sv_split[46]), "f"(sv_split[47])); - _fp8_0[11] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(sv_split[48]), "f"(sv_split[49]), - "f"(sv_split[50]), "f"(sv_split[51])); - _fp8_0[12] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(sv_split[52]), "f"(sv_split[53]), - "f"(sv_split[54]), "f"(sv_split[55])); - _fp8_0[13] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(sv_split[56]), "f"(sv_split[57]), - "f"(sv_split[58]), "f"(sv_split[59])); - _fp8_0[14] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(sv_split[60]), "f"(sv_split[61]), - "f"(sv_split[62]), "f"(sv_split[63])); - _fp8_0[15] = _packed; - } - int p_stage_base = smem_p_fp8_addr + (unsigned int)(phase * 8192); - #pragma unroll - for (int p_vec = 0; p_vec < 4; p_vec++) { - asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((p_stage_base + (my_row * 128 + (n_half * 4 + p_vec) * 16 ^ (my_row * 128 + (n_half * 4 + p_vec) * 16 >> 7 & 7) << 4))), "r"(_fp8_0[p_vec * 4]), "r"(_fp8_0[p_vec * 4 + 1]), "r"(_fp8_0[p_vec * 4 + 2]), "r"(_fp8_0[p_vec * 4 + 3]) : "memory"); - } - } - } - { - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - } - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - { - asm volatile("" ::: "memory"); - } - mbarrier_arrive(p_full_addr + (phase) * 8); - { - asm volatile("" ::: "memory"); - float _fma_3 = __fmaf_rn(row_sum_val, acc_scale, block_sum); - row_sum_val = _fma_3; - } - } - smem_stats_sum[stats_row] = row_sum_val; - smem_stats_final_max[stats_row] = row_max_val; - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(sum_ready_addr); - softmax_tile_cursor = softmax_tile_cursor + num_kv_tiles_1; - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(tmem_dealloc_addr); - } - } - // ---- Role: softmax_padding ---- - if (warp == 1 || warp == 3) { - { // softmax_padding_main - const int lowhead_softmax_padding = 0; - } - } - // ---- Role: correction_wg ---- - if (warp == 4 || warp == 6) { - { // correction_wg_main - float softmax_scale_log2_1 = bmm1_scale[0] * 1.4426950408889634f; - float output_scale = bmm2_scale[0]; - const int wg_dummy_inc_1 = 0; - int all_num_kv_tiles_3 = (sparse_topk + 128 - 1) / 128; - const int tmem_row_base_1 = ((1) ? warp % 2 * 32 : warp % 4 * 32); - const int n_half_1 = ((1) ? (int)(warp % 4 / 2) : 0); - const int my_row_1 = tmem_row_base_1 + lane; - const int stats_row_1 = n_half_1 * 64 + my_row_1; - const int corr_row = tmem_row_base_1 << 16; - int correction_tile_cursor = 0; - unsigned int _phase_q_full_0_2 = 0; - unsigned int _phase_pv_done_0 = 0; - unsigned int _phase_o_done_0 = 0; - unsigned int _phase_sum_ready_0 = 0; - #pragma unroll - for (unsigned int work_idx_2 = 0; work_idx_2 < 1; work_idx_2++) { - int split_idx_3 = 0; - int query_idx_3 = bid / 2 >> 1; - int v_chunk_1 = bid / 2 & 1; - { - split_idx_3 = 0; - query_idx_3 = bid / 2; - v_chunk_1 = bid % 2; - } - int tiles_per_split_3 = all_num_kv_tiles_3 + 1 - 1; - int first_tile_3 = split_idx_3 * tiles_per_split_3; - int num_kv_tiles_2 = tiles_per_split_3; - #pragma unroll 1 - for (int tile_2 = 0; tile_2 < num_kv_tiles_2; tile_2++) { - int pipeline_tile_1 = correction_tile_cursor + tile_2; - int phase_1 = pipeline_tile_1 & 1; - int stats_wait_phase = pipeline_tile_1 >> 1 & 1; - mbarrier_wait(stats_addr + (phase_1) * 8, stats_wait_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - float acc_scale_1 = smem_stats_max[phase_1 * 128 + stats_row_1]; - if (tile_2 > 0) { - mbarrier_wait(pv_done_addr, _phase_pv_done_0); - _phase_pv_done_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - } - if (tile_2 > 0) { - int _vote_1 = __any_sync(0xFFFFFFFF, acc_scale_1 < 1.0f); - int any_rescale = _vote_1; - if (any_rescale != 0) { - #pragma unroll - for (int vs = 0; vs < 2; vs++) { - int o_base = taddr + 256 + (unsigned int)(vs * 128) + (unsigned int)corr_row; - #pragma unroll - for (int c = 0; c < 128; c += 16) { - float _tmem_load_3[16]; - tmem_ld_x16(&_tmem_load_3[0], o_base + c); - const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; - #pragma unroll - for (int _ls = 0; _ls < 8; _ls++) - mul_f32x2_inplace(&reinterpret_cast(_tmem_load_3)[_ls], _scale2_0); - tmem_st_x16_f32(o_base + c, _tmem_load_3); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - } - } - } - } - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(corr_done_addr + (phase_1) * 8); - } - mbarrier_wait(o_done_addr, _phase_o_done_0); - _phase_o_done_0 ^= 1; - mbarrier_wait(sum_ready_addr, _phase_sum_ready_0); - _phase_sum_ready_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - float total_sum = smem_stats_sum[stats_row_1]; - float final_max = smem_stats_final_max[stats_row_1]; - { - smem_corr_warp_pair_exchange[stats_row_1] = total_sum; - asm volatile("barrier.sync %0, 64;" :: "r"(4 + warp % 2) : "memory"); - total_sum = total_sum + smem_corr_warp_pair_exchange[stats_row_1 ^ 64]; - } - float _rcp_0 = approx_rcp(total_sum); - float inv_sum = ((total_sum > 0.0f) ? _rcp_0 : 0.0f); - int head_idx = my_row_1; - { - head_idx = bid % 2 * 64 + my_row_1; - } - int direct_o_offset = (query_idx_3 * num_heads + head_idx) * 512 + v_chunk_1 * 256; - int partial_o_offset = (query_idx_3 * num_heads + head_idx + split_idx_3) * 512 + ((0) ? v_chunk_1 : n_half_1) * 256; - int o_offset = partial_o_offset; - #pragma unroll - for (int vs_1 = 0; vs_1 < 2; vs_1++) { - int o_base_epi = taddr + 256 + (unsigned int)(vs_1 * 128) + (unsigned int)corr_row; - #pragma unroll - for (int c_1 = 0; c_1 < 128; c_1 += 32) { - float _tmem_load_4[32]; - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(_tmem_load_4[0]), "=f"(_tmem_load_4[1]), "=f"(_tmem_load_4[2]), "=f"(_tmem_load_4[3]), "=f"(_tmem_load_4[4]), "=f"(_tmem_load_4[5]), "=f"(_tmem_load_4[6]), "=f"(_tmem_load_4[7]), "=f"(_tmem_load_4[8]), "=f"(_tmem_load_4[9]), "=f"(_tmem_load_4[10]), "=f"(_tmem_load_4[11]), "=f"(_tmem_load_4[12]), "=f"(_tmem_load_4[13]), "=f"(_tmem_load_4[14]), "=f"(_tmem_load_4[15]), "=f"(_tmem_load_4[16]), "=f"(_tmem_load_4[17]), "=f"(_tmem_load_4[18]), "=f"(_tmem_load_4[19]), "=f"(_tmem_load_4[20]), "=f"(_tmem_load_4[21]), "=f"(_tmem_load_4[22]), "=f"(_tmem_load_4[23]), "=f"(_tmem_load_4[24]), "=f"(_tmem_load_4[25]), "=f"(_tmem_load_4[26]), "=f"(_tmem_load_4[27]), "=f"(_tmem_load_4[28]), "=f"(_tmem_load_4[29]), "=f"(_tmem_load_4[30]), "=f"(_tmem_load_4[31]) - : "r"(o_base_epi + c_1)); - int gmem_base = o_offset + vs_1 * 128 + c_1; - if (head_idx < num_heads) { - #pragma unroll - for (int j = 0; j < 32; j += 8) { - { - const float2 _prescale2_1 = {inv_sum * output_scale, inv_sum * output_scale}; - #if __CUDA_ARCH__ >= 1000 - #pragma unroll - for (int _ps = 0; _ps < 4; _ps++) - mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_4[j])[_ps], _prescale2_1); - #else - #pragma unroll - for (int _ps = 0; _ps < 8; _ps++) - _tmem_load_4[j + _ps] *= inv_sum * output_scale; - #endif - __nv_bfloat162 _pk[4]; - _pk[0] = __floats2bfloat162_rn(_tmem_load_4[j + 0], _tmem_load_4[j + 1]); - _pk[1] = __floats2bfloat162_rn(_tmem_load_4[j + 2], _tmem_load_4[j + 3]); - _pk[2] = __floats2bfloat162_rn(_tmem_load_4[j + 4], _tmem_load_4[j + 5]); - _pk[3] = __floats2bfloat162_rn(_tmem_load_4[j + 6], _tmem_load_4[j + 7]); - *reinterpret_cast(&((__nv_bfloat16*)(O + (gmem_base + j)))[0]) = *reinterpret_cast(&_pk[0]); - } - } - } - } - } - correction_tile_cursor = correction_tile_cursor + num_kv_tiles_2; - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(tmem_dealloc_addr); - } - } - // ---- Role: correction_padding ---- - if (warp == 5 || warp == 7) { - { // correction_padding_main - const int lowhead_correction_padding = 0; - } - } - // ---- Role: mma_warp ---- - if (warp == 8) { - { // mma_warp_main - const int wg2_dummy_1 = 0; - int all_num_kv_tiles_4 = (sparse_topk + 128 - 1) / 128; - unsigned int mma_k_stage = 0; - unsigned int mma_v_stage = 0; - unsigned int mma_kv_stage = 0; - int mma_tile_cursor = 0; - unsigned int _phase_s_seeded_0 = 0; - unsigned int _phase_q_full_0_3 = 0; - unsigned int _phase_q_pair_ready_0 = 0; - unsigned int _phase_k_full = 0; - unsigned int _phase_k_full_1 = 0; - unsigned int _phase_kv_pair_ready = 0; - unsigned int _phase_v_full = 0; - #pragma unroll - for (unsigned int work_idx_3 = 0; work_idx_3 < 1; work_idx_3++) { - int split_idx_4 = 0; - { - split_idx_4 = 0; - } - int tiles_per_split_4 = all_num_kv_tiles_4 + 1 - 1; - int first_tile_4 = split_idx_4 * tiles_per_split_4; - int num_kv_tiles_3 = tiles_per_split_4; - { - if (cta_rank == 0) { - mbarrier_wait(q_full_addr, _phase_q_full_0_3); - _phase_q_full_0_3 ^= 1; - } - } - int first_pv = 1; - #pragma unroll 1 - for (int tile_3 = 0; tile_3 < num_kv_tiles_3; tile_3++) { - int pipeline_tile_2 = mma_tile_cursor + tile_3; - int phase_2 = pipeline_tile_2 & 1; - int score_col = ((phase_2 != 0) ? 128 : 0); - { - if (cta_rank == 0) { - mbarrier_wait(k_full_addr + (mma_k_stage) * 8, _phase_k_full); - int _mma_a_lo_0 = (((smem_q_full_addr) >> 4) & 0x3FFF) + (0) * 2048; - int _mma_b_lo_0 = (((smem_k_full_addr) >> 4) & 0x3FFF) + (mma_k_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 136314896;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_tmem_scratch + (score_col))), "r"(0)); - int _mma_a_lo_1 = (((smem_q_full_addr + 8192) >> 4) & 0x3FFF) + (0) * 2048; - int _mma_b_lo_1 = (((smem_k_full_addr + 8192) >> 4) & 0x3FFF) + (mma_k_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 136314896;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); - int _mma_a_lo_2 = (((smem_q_full_addr + 16384) >> 4) & 0x3FFF) + (0) * 2048; - int _mma_b_lo_2 = (((smem_k_full_addr + 16384) >> 4) & 0x3FFF) + (mma_k_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 136314896;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); - int _mma_a_lo_3 = (((smem_q_full_addr + 24576) >> 4) & 0x3FFF) + (0) * 2048; - int _mma_b_lo_3 = (((smem_k_full_addr + 24576) >> 4) & 0x3FFF) + (mma_k_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 136314896;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 2;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"((tmem_tmem_scratch + (score_col))), "r"(1)); - if (cta_rank == 0) { - elect_commit_cg2_multicast(s_full_addr + (phase_2) * 8, (uint16_t)(3)); - } - if (cta_rank == 0) { - elect_commit_cg2_multicast(k_empty_addr + (mma_k_stage) * 8, (uint16_t)(3)); - } - } - mma_k_stage += 1; - if (mma_k_stage == 2) { mma_k_stage = 0; _phase_k_full ^= 1; } - } - if (tile_3 > 0) { - int prev_pipeline_tile = pipeline_tile_2 - 1; - int prev_phase = prev_pipeline_tile & 1; - int pv_wait_phase = prev_pipeline_tile >> 1 & 1; - mbarrier_wait(p_full_addr + (prev_phase) * 8, pv_wait_phase); - mbarrier_wait(corr_done_addr + (prev_phase) * 8, pv_wait_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - { - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cluster.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(pv_pair_ready_addr + (unsigned int)(prev_phase * 8)), "r"(0) : "memory"); - if (cta_rank == 0) { - mbarrier_wait_cluster_hint(pv_pair_ready_addr + (prev_phase) * 8, pv_wait_phase, 10000000); - } - } - int p_col = ((prev_phase != 0) ? 224 : 96); - { - if (cta_rank == 0) { - mbarrier_wait(v_full_addr + (mma_v_stage) * 8, _phase_v_full); - int _mma_a_lo_6 = (((smem_p_fp8_addr) >> 4) & 0x3FFF) + (prev_phase) * 512; - int _mma_b_lo_6 = ((((smem_v_full_addr) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 138477584;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_6), "r"(_mma_b_lo_6), "r"((tmem_tmem_scratch + (256))), "r"(((first_pv) ? 0 : 1))); - int _mma_a_lo_7 = (((smem_p_fp8_addr) >> 4) & 0x3FFF) + (prev_phase) * 512; - int _mma_b_lo_7 = ((((smem_v_full_addr + 16384) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 138477584;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_7), "r"(_mma_b_lo_7), "r"((tmem_tmem_scratch + (384))), "r"(((first_pv) ? 0 : 1))); - if (cta_rank == 0) { - elect_commit_cg2_multicast(v_empty_addr + (mma_v_stage) * 8, (uint16_t)(3)); - } - } - mma_v_stage += 1; - if (mma_v_stage == 2) { mma_v_stage = 0; _phase_v_full ^= 1; } - } - first_pv = 0; - if (cta_rank == 0) { - elect_commit_cg2_multicast(pv_done_addr, (uint16_t)(3)); - } - } - } - int last_pipeline_tile = mma_tile_cursor + num_kv_tiles_3 - 1; - int last_phase = last_pipeline_tile & 1; - int drain_wait_phase = last_pipeline_tile >> 1 & 1; - mbarrier_wait(p_full_addr + (last_phase) * 8, drain_wait_phase); - mbarrier_wait(corr_done_addr + (last_phase) * 8, drain_wait_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - { - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cluster.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(pv_pair_ready_addr + (unsigned int)(last_phase * 8)), "r"(0) : "memory"); - if (cta_rank == 0) { - mbarrier_wait_cluster_hint(pv_pair_ready_addr + (last_phase) * 8, drain_wait_phase, 10000000); - } - } - int p_col_last = ((last_phase != 0) ? 224 : 96); - { - if (cta_rank == 0) { - mbarrier_wait(v_full_addr + (mma_v_stage) * 8, _phase_v_full); - int _mma_a_lo_10 = (((smem_p_fp8_addr) >> 4) & 0x3FFF) + (last_phase) * 512; - int _mma_b_lo_10 = ((((smem_v_full_addr) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 138477584;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_10), "r"(_mma_b_lo_10), "r"((tmem_tmem_scratch + (256))), "r"(((first_pv) ? 0 : 1))); - int _mma_a_lo_11 = (((smem_p_fp8_addr) >> 4) & 0x3FFF) + (last_phase) * 512; - int _mma_b_lo_11 = ((((smem_v_full_addr + 16384) >> 4) & 0x3FFF) | 0x4000000) + (mma_v_stage) * 2048; - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" - ".reg .b64 da, db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" - "mov.b32 adhi, 0x40004040;\n\t" - "mov.b32 bdhi, 0x40004040;\n\t" - "mov.b32 id, 138477584;\n\t" - "mov.b32 alo, %0;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "add.u32 alo, alo, 2;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 da, {alo, adhi};\n\t" - "mov.b64 db, {blo, bdhi};\n\t" - "@leader tcgen05.mma.cta_group::2.kind::f8f6f4 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" - "}\n" - :: "r"(_mma_a_lo_11), "r"(_mma_b_lo_11), "r"((tmem_tmem_scratch + (384))), "r"(((first_pv) ? 0 : 1))); - if (cta_rank == 0) { - elect_commit_cg2_multicast(v_empty_addr + (mma_v_stage) * 8, (uint16_t)(3)); - } - } - mma_v_stage += 1; - if (mma_v_stage == 2) { mma_v_stage = 0; _phase_v_full ^= 1; } - } - if (cta_rank == 0) { - elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); - } - if (cta_rank == 0) { - elect_commit_cg2_multicast(o_done_addr, (uint16_t)(3)); - } - mma_tile_cursor = mma_tile_cursor + num_kv_tiles_3; - } - mbarrier_arrive(tmem_dealloc_addr); - unsigned int _phase_tmem_dealloc_0 = 0; - mbarrier_wait(tmem_dealloc_addr, _phase_tmem_dealloc_0); - _phase_tmem_dealloc_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - { - int tmem_dealloc_peer_rank = bid % 2 ^ 1; - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cluster.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(tmem_dealloc_peer_addr), "r"(tmem_dealloc_peer_rank) : "memory"); - mbarrier_wait(tmem_dealloc_peer_addr, 0); - int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); - asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); - } - } - } - // ---- Role: empty1 ---- - if (warp >= 10 && warp <= 11) { - { // empty1_main - const int wg2_dummy_2 = 0; - unsigned int _phase_q_full_0_4 = 0; - } - } - - // Cleanup -} - -} // extern "C" diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_c967e7efe220160e28d3_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_c967e7efe220160e28d3_binding.cu deleted file mode 100644 index eeec0e7ee24..00000000000 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_c967e7efe220160e28d3_binding.cu +++ /dev/null @@ -1,534 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -// Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_c967e7efe220160e28d3'. -#include -#include -#include -#include - -#include "tvm_ffi_utils.h" - -#include -#include -#include -#include -#include -#include -#include - -struct CakeTensorMap; -extern "C" __global__ void kernel_cake_dsv4_c967e7efe220160e28d3(CakeTensorMap const* tmap_q, CakeTensorMap const* tmap_swa_kv, CakeTensorMap const* tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int batch_size, int max_q_len, int ragged_query, int has_sinks); - - -namespace cake_host_shim_2f26cde03513c4ba { - -using tvm::ffi::Optional; -using tvm::ffi::TensorView; - -class ScopedCudaDevice { - public: - explicit ScopedCudaDevice(int device_id) { - cudaError_t error = cudaGetDevice(&previous_device_); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed before host-shim launch: cudaError=" - << static_cast(error); - if (previous_device_ != device_id) { - error = cudaSetDevice(device_id); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaSetDevice failed before host-shim launch for cuda:" - << device_id << ": cudaError=" << static_cast(error); - restore_ = true; - } - } - - ScopedCudaDevice(const ScopedCudaDevice&) = delete; - ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; - - ~ScopedCudaDevice() noexcept { - if (restore_) { - (void)cudaSetDevice(previous_device_); - } - } - - private: - int previous_device_ = -1; - bool restore_ = false; -}; - -inline int64_t CakeDeviceMultiprocessorCount(int device_id) { - constexpr int kMaxCachedCudaDevices = 64; - TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) - << "physical-SM-count cache does not cover cuda:" << device_id; - static std::atomic count_by_device[kMaxCachedCudaDevices]{}; - int cached = count_by_device[device_id].load(std::memory_order_acquire); - if (cached > 0) return cached; - - int count = 0; - cudaError_t error = cudaDeviceGetAttribute( - &count, cudaDevAttrMultiProcessorCount, device_id); - TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) - << "querying multiProcessorCount failed for cuda:" << device_id - << ": cudaError=" << static_cast(error) << ", count=" << count; - // Concurrent first launches may repeat the immutable device query, but all - // publication is atomic and every later hot-path lookup is one acquire load. - count_by_device[device_id].store(count, std::memory_order_release); - return count; -} - -inline void CheckCudaTensor(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) - << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; -} - -inline void CheckSameCudaDevice( - const TensorView& t, - const TensorView& reference, - const char* name, - const char* reference_name) { - TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) - << name << " must be on the same CUDA device as " << reference_name - << ": got cuda:" << t.device().device_id - << " versus cuda:" << reference.device().device_id; -} - -inline void CheckCurrentCudaDevice( - const TensorView& reference, - const char* reference_name) { - int current_device = -1; - cudaError_t error = cudaGetDevice(¤t_device); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed while validating " << reference_name - << ": cudaError=" << static_cast(error); - TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) - << "current CUDA device must match " << reference_name - << ": current=cuda:" << current_device - << ", tensor=cuda:" << reference.device().device_id; -} - -inline void CheckContiguous(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; -} - -inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { - DLDataType d = t.dtype(); - TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) - << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits - << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits - << ", lanes=" << (int)d.lanes << ")"; -} - -// A logical axis.outer(trailing) folds every source dim above the trailing -// dimensions. Shape products are independent of physical strides, so verify -// the leading dimensions form one dense row-major chain instead of inventing -// a "folded stride". The descriptor reads its exact adjacent physical step -// separately through stride[-(trailing + 1)]. -inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { - TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) - << name << " cannot fold leading dimensions above " << trailing - << " trailing dims from ndim=" << t.ndim(); - int outer_last = t.ndim() - trailing - 1; - if (outer_last <= 0) { - return; - } - int64_t step = t.stride(outer_last); - TVM_FFI_CHECK(step > 0, ValueError) - << name << " physical strides must be positive"; - int64_t expected = step; - for (int axis = outer_last - 1; axis >= 0; --axis) { - expected *= t.size(axis + 1); - if (t.size(axis) > 1) { - TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) - << name << " leading dims are not physically foldable above " << trailing - << " trailing dims: stride(" << axis << ")=" << t.stride(axis) - << ", expected " << expected; - } - } -} - -// Explicit immutable caller storage: initialize once, outside graph capture, -// exactly as the production TmaDeviceSlot path. The owner must retain the -// allocation and must never expose it as mutable scratch or rewrite its bytes. -template -static inline void PrepareImmutableCallerTmaWorkspace( - const tvm::ffi::Tensor& retained_workspace, - const CUtensorMap (&maps)[N], cudaStream_t stream) { - void* workspace = retained_workspace.data_ptr(); - CUcontext context = nullptr; - CUresult result = cuCtxGetCurrent(&context); - TVM_FFI_CHECK(result == CUDA_SUCCESS && context != nullptr, RuntimeError) - << "immutable TMA workspace requires an active CUDA context"; - CUcontext allocation_context = nullptr; - result = cuPointerGetAttribute( - &allocation_context, CU_POINTER_ATTRIBUTE_CONTEXT, - reinterpret_cast(workspace)); - TVM_FFI_CHECK(result == CUDA_SUCCESS && allocation_context == context, RuntimeError) - << "immutable TMA workspace must belong to the active CUDA context"; - unsigned long long buffer_id = 0; - result = cuPointerGetAttribute( - &buffer_id, CU_POINTER_ATTRIBUTE_BUFFER_ID, - reinterpret_cast(workspace)); - TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) - << "querying immutable TMA allocation identity failed"; - std::string key = std::to_string(reinterpret_cast(context)); - key += ":" + std::to_string(reinterpret_cast(workspace)); - key += ":" + std::to_string(buffer_id); - std::string descriptor_bytes(reinterpret_cast(maps), sizeof(maps)); - struct InitializedWorkspace { - tvm::ffi::Tensor owner; - std::string descriptor_bytes; - }; - static std::mutex mu; - // Keep the actual Tensor allocation alive: a caching allocator may reuse a - // freed suballocation without changing its enclosing CUDA buffer identity. - static auto* initialized = new std::unordered_map(); - std::lock_guard lock(mu); - auto it = initialized->find(key); - if (it != initialized->end()) { - TVM_FFI_CHECK(it->second.descriptor_bytes == descriptor_bytes, ValueError) - << "immutable TMA descriptors changed; provide new private descriptor storage"; - return; - } - CUstreamCaptureStatus capture_status = CU_STREAM_CAPTURE_STATUS_NONE; - result = cuStreamIsCapturing(reinterpret_cast(stream), &capture_status); - TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) - << "querying immutable TMA initialization capture state failed"; - TVM_FFI_CHECK(capture_status == CU_STREAM_CAPTURE_STATUS_NONE, RuntimeError) - << "immutable TMA storage must be initialized before CUDA Graph capture"; - // Synchronous host-to-device initialization completes before publishing the - // slot to any thread/stream, matching TmaDeviceSlot. The consumer retains its - // generated tensor-map acquire; the mutable upload's release fence is unchanged. - result = cuMemcpyHtoD(reinterpret_cast(workspace), maps, sizeof(maps)); - TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) - << "initializing immutable caller TMA descriptors failed"; - initialized->emplace( - std::move(key), InitializedWorkspace{retained_workspace, std::move(descriptor_bytes)}); -} - -static inline void* CallerTmaWorkspaceSlot( - const tvm::ffi::TensorView& workspace, size_t slot) { - return static_cast(workspace.data_ptr()) + slot * sizeof(CUtensorMap); -} - -// 4D TMA descriptor for buffer 'tmap_q' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 3, ValueError) - << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - int64_t d2 = t.size(t.ndim() - 2); - int64_t d3 = t.size(t.ndim() - 3); - TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) - << "TMA source 'tmap_q' trailing dims must be positive"; - TVM_FFI_CHECK(d1 % 64 == 0, ValueError) - << "TMA source 'tmap_q' extent " << d1 - << " must divide exactly by " << 64; - uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(d3)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 8u <= global_dim[2] && 1u <= global_dim[3], ValueError) - << "TMA box (64, 64, 8, 1) exceeds resolved global dims for 'tmap_q'"; - int64_t carrier_stride_0 = d1; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; - int64_t carrier_stride_1 = 64; - TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; - TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; - TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; - int64_t carrier_stride_2 = (d2 * d1); - TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; - TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; - TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; - uint64_t global_strides[3] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - (uint64_t)((carrier_stride_1 * 16) / 8), - (uint64_t)((carrier_stride_2 * 16) / 8), - }; - uint32_t box_dim[4] = {64u, 64u, 8u, 1u}; - uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_swa_kv' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_swa_kv' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_kv'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_compressed_kv' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_compressed_kv(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_compressed_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_compressed_kv' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_compressed_kv' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_compressed_kv"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_compressed_kv' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_kv'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_kv') failed: CUresult=" << (int)r; - return tm; -} - -void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_compressed_kv, TensorView arg_O, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_seq_lens, TensorView arg_cum_seq_lens_q, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_batch_size, int64_t arg_max_q_len, int64_t arg_ragged_query, int64_t arg_has_sinks, tvm::ffi::Tensor arg_tma_descriptor_workspace, int64_t grid_x, int64_t grid_y, int64_t grid_z) { - DLDevice dev = arg_tmap_q.device(); - ScopedCudaDevice device_guard(dev.device_id); - CheckCudaTensor(arg_tmap_q, "tmap_q"); - CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); - CheckContiguous(arg_tmap_q, "tmap_q"); - CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); - CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 4, 16, 1); - CheckCudaTensor(arg_tmap_compressed_kv, "tmap_compressed_kv"); - CheckDtype(arg_tmap_compressed_kv, "tmap_compressed_kv", 4, 16, 1); - CheckCudaTensor(arg_O, "O"); - CheckDtype(arg_O, "O", 4, 16, 1); - CheckContiguous(arg_O, "O"); - CheckCudaTensor(arg_swa_indices, "swa_indices"); - CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); - CheckContiguous(arg_swa_indices, "swa_indices"); - CheckCudaTensor(arg_compressed_indices, "compressed_indices"); - CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); - CheckContiguous(arg_compressed_indices, "compressed_indices"); - CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); - CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckCudaTensor(arg_seq_lens, "seq_lens"); - CheckDtype(arg_seq_lens, "seq_lens", 0, 32, 1); - CheckContiguous(arg_seq_lens, "seq_lens"); - CheckCudaTensor(arg_cum_seq_lens_q, "cum_seq_lens_q"); - CheckDtype(arg_cum_seq_lens_q, "cum_seq_lens_q", 0, 32, 1); - CheckContiguous(arg_cum_seq_lens_q, "cum_seq_lens_q"); - CheckCudaTensor(arg_sinks, "sinks"); - CheckDtype(arg_sinks, "sinks", 2, 32, 1); - CheckContiguous(arg_sinks, "sinks"); - CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); - CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); - CheckContiguous(arg_bmm1_scale, "bmm1_scale"); - CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); - CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); - CheckContiguous(arg_bmm2_scale, "bmm2_scale"); - TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) - << "scalar 'num_heads' value " << arg_num_heads - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) - << "scalar 'swa_index_stride' value " << arg_swa_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) - << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) - << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) - << "scalar 'sparse_topk' value " << arg_sparse_topk - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_batch_size >= -2147483648LL && arg_batch_size <= 2147483647LL, ValueError) - << "scalar 'batch_size' value " << arg_batch_size - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_max_q_len >= -2147483648LL && arg_max_q_len <= 2147483647LL, ValueError) - << "scalar 'max_q_len' value " << arg_max_q_len - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_ragged_query >= -2147483648LL && arg_ragged_query <= 2147483647LL, ValueError) - << "scalar 'ragged_query' value " << arg_ragged_query - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) - << "scalar 'has_sinks' value " << arg_has_sinks - << " is outside i32 range [-2147483648, 2147483647]"; - CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); - CheckSameCudaDevice(arg_tmap_compressed_kv, arg_tmap_q, "tmap_compressed_kv", "tmap_q"); - CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); - CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); - CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); - CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); - CheckSameCudaDevice(arg_seq_lens, arg_tmap_q, "seq_lens", "tmap_q"); - CheckSameCudaDevice(arg_cum_seq_lens_q, arg_tmap_q, "cum_seq_lens_q", "tmap_q"); - CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); - CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); - CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); - CheckCudaTensor(arg_tma_descriptor_workspace, "tma_descriptor_workspace"); - CheckSameCudaDevice(arg_tma_descriptor_workspace, arg_tmap_q, "tma_descriptor_workspace", "tmap_q"); - CheckContiguous(arg_tma_descriptor_workspace, "tma_descriptor_workspace"); - CheckDtype(arg_tma_descriptor_workspace, "tma_descriptor_workspace", 1, 8, 1); - TVM_FFI_CHECK(arg_tma_descriptor_workspace.numel() >= 384, ValueError) - << "tma_descriptor_workspace requires at least 384 bytes"; - TVM_FFI_CHECK( - reinterpret_cast(arg_tma_descriptor_workspace.data_ptr()) % 128 == 0, - ValueError) - << "tma_descriptor_workspace must be 128-byte aligned"; - TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) - << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y - << ", " << grid_z << ")"; - - cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); - CUtensorMap h_tmap_q = EncodeTma_tmap_q(arg_tmap_q); - CakeTensorMap const* p_tmap_q = static_cast(CallerTmaWorkspaceSlot(arg_tma_descriptor_workspace, 0u)); - CUtensorMap h_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); - CakeTensorMap const* p_tmap_swa_kv = static_cast(CallerTmaWorkspaceSlot(arg_tma_descriptor_workspace, 1u)); - CUtensorMap h_tmap_compressed_kv = EncodeTma_tmap_compressed_kv(arg_tmap_compressed_kv); - CakeTensorMap const* p_tmap_compressed_kv = static_cast(CallerTmaWorkspaceSlot(arg_tma_descriptor_workspace, 2u)); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); - int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); - int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); - int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); - int* p_seq_lens = static_cast(arg_seq_lens.data_ptr()); - int* p_cum_seq_lens_q = static_cast(arg_cum_seq_lens_q.data_ptr()); - float* p_sinks = static_cast(arg_sinks.data_ptr()); - float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); - float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); - int32_t v_num_heads = (int32_t)arg_num_heads; - int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; - int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; - int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; - int32_t v_sparse_topk = (int32_t)arg_sparse_topk; - int32_t v_batch_size = (int32_t)arg_batch_size; - int32_t v_max_q_len = (int32_t)arg_max_q_len; - int32_t v_ragged_query = (int32_t)arg_ragged_query; - int32_t v_has_sinks = (int32_t)arg_has_sinks; - void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_tmap_compressed_kv, &p_O, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_seq_lens, &p_cum_seq_lens_q, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_batch_size, &v_max_q_len, &v_ragged_query, &v_has_sinks}; - - static std::mutex smem_mu; - static auto* smem_status_by_device = - new std::unordered_map(); - cudaError_t smem_status = cudaSuccess; - { - std::lock_guard lock(smem_mu); - auto it = smem_status_by_device->find(dev.device_id); - if (it == smem_status_by_device->end()) { - int current_device = -1; - smem_status = cudaGetDevice(¤t_device); - if (smem_status == cudaSuccess && current_device != dev.device_id) { - smem_status = cudaErrorInvalidDevice; - } - if (smem_status == cudaSuccess) { - smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_c967e7efe220160e28d3), - cudaFuncAttributeMaxDynamicSharedMemorySize, - 198144); - } - smem_status_by_device->emplace(dev.device_id, smem_status); - } else { - smem_status = it->second; - } - } - TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_c967e7efe220160e28d3 failed: " - << cudaGetErrorString(smem_status); - dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); - dim3 block(256u, 1u, 1u); - const CUtensorMap workspace_maps[] = {h_tmap_q, h_tmap_swa_kv, h_tmap_compressed_kv}; - PrepareImmutableCallerTmaWorkspace(arg_tma_descriptor_workspace, workspace_maps, stream); - - cudaLaunchConfig_t config{}; - config.gridDim = grid; - config.blockDim = block; - config.dynamicSmemBytes = 198144u; - config.stream = stream; - cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_c967e7efe220160e28d3), kargs); - TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_c967e7efe220160e28d3 failed: " - << cudaGetErrorString(launch_status); -} - -} // namespace cake_host_shim_2f26cde03513c4ba - -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_2f26cde03513c4ba::Run); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_c967e7efe220160e28d3_kernel.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_c967e7efe220160e28d3_kernel.cu deleted file mode 100644 index b4dce1a94ee..00000000000 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_c967e7efe220160e28d3_kernel.cu +++ /dev/null @@ -1,1686 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -typedef signed char int8_t; -typedef unsigned char uint8_t; -typedef unsigned short uint16_t; -typedef unsigned int uint32_t; -#if defined(__CUDACC_RTC__) -typedef unsigned long long uint64_t; -#else -typedef unsigned long uint64_t; -#endif -static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); -typedef signed int int32_t; -typedef short int int16_t; -struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; -struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; -static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); -static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); -template -struct __align__(128) CakeTensorMapPack { CakeTensorMap maps[N]; }; - -#if defined(__CUDACC_RTC__) -typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; -#else -#include -#endif - -static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); -static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); -#include -#include - -#define CAKE_INF CUDART_INF_F -#define TMEM_NCOLS 512 -#define TMEM_TMEM_OFFSET 0 -#define NUM_KV_PIPE_STAGES 4 -#define SMEM_SMEM_Q_OFF 1024 -#define SMEM_SMEM_Q_STAGE_BYTES 8192 -#define SMEM_SMEM_Q_STRIDE 8192 -#define SMEM_SMEM_KV_OFF 66560 -#define SMEM_SMEM_KV_STAGE_BYTES 32768 -#define SMEM_SMEM_KV_STRIDE 32768 -#define SMEM_SMEM_V_OFF 66560 -#define SMEM_SMEM_V_STAGE_BYTES 32768 -#define SMEM_SMEM_V_STRIDE 32768 -#define SMEM_TOTAL 198144 -#define THREADS 256 - -#include - -__device__ __forceinline__ uint32_t elect_sync() { - uint32_t pred = 0; - asm volatile( - "{\n\t" - ".reg .pred %%px;\n\t" - "elect.sync _|%%px, %1;\n\t" - "@%%px mov.s32 %0, 1;\n\t" - "}\n" - : "+r"(pred) - : "r"(0xFFFFFFFF)); - return pred; -} - - -__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { - asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" - :: "r"(mbar_addr), "r"(count) : "memory"); -} - -__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { - asm volatile("mbarrier.init.b64 [%0], %1;" - :: "l"(mbar_addr), "r"(count) : "memory"); -} - - -__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - - -// CTA-local pipelines have short, resident producer/consumer edges. Omitting -// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the -// explicit loop still makes this helper blocking until acquire succeeds. -__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE;\n\t" - "bra.uni LAB_WAIT;\n\t" - "DONE:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Source-faithful relaxed CTA wait used only by a typed protocol that does -// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. -__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, 10000000;\n\t" - "@P1 bra.uni DONE_RELAXED;\n\t" - "bra.uni LAB_WAIT_RELAXED;\n\t" - "DONE_RELAXED:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Exact source ports may request the PTX suspendTimeHint operand explicitly. -// The hint is expressed in nanoseconds and is kept separate from the canonical -// no-hint CTA helper so unrelated schedules retain their existing retry path. -__device__ __forceinline__ void mbarrier_wait_suspend( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_SUSPEND:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_SUSPEND;\n\t" - "bra.uni LAB_WAIT_SUSPEND;\n\t" - "DONE_SUSPEND:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE_CLUSTER;\n\t" - "bra.uni LAB_WAIT_CLUSTER;\n\t" - "DONE_CLUSTER:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - ".reg .u32 WAIT_ADDR;\n\t" - "mov.u32 WAIT_ADDR, %0;\n\t" - "LAB_WAIT_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [WAIT_ADDR], %1, %2;\n\t" - "@P1 bra.uni DONE_HINT;\n\t" - "bra.uni LAB_WAIT_HINT;\n\t" - "DONE_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -// Exact unqualified CTA wait used by source schedules whose PTX intentionally -// omits the acquire qualifier while retaining a typed suspendTimeHint operand. -__device__ __forceinline__ void mbarrier_wait_relaxed_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED_HINT:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra DONE_RELAXED_HINT;\n\t" - "bra LAB_WAIT_RELAXED_HINT;\n\t" - "DONE_RELAXED_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); -} - -__device__ __forceinline__ void mbarrier_wait_cluster_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" - "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" - "DONE_CLUSTER_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_suspend( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait_cluster(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); - } -} - - -__device__ __forceinline__ void tcgen05_mma_f16( - int taddr, uint64_t a_desc, uint64_t b_desc, - uint32_t i_desc, int enable_input_d) { - asm volatile( - "{\n\t" - ".reg .pred p;\n\t" - "setp.ne.b32 p, %4, 0;\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" - "}\n" - :: "r"(taddr), "l"(a_desc), "l"(b_desc), - "r"(i_desc), "r"(enable_input_d) - : "memory"); -} - - -__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { - return (x & 0x3FFFFULL) >> 4ULL; -} - - -union MmaSmemDesc { - uint64_t u64; - uint32_t u32[2]; -}; - -__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { - MmaSmemDesc tmp; - tmp.u64 = smem_desc; - tmp.u32[0] += offset; - smem_desc = tmp.u64; -} - - -__device__ __forceinline__ void elect_commit(int mbar_addr) { - asm volatile( - "{\n\t" - ".reg .pred leader;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" - ".shared::cluster.b64 [%0];\n\t" - "}\n" - :: "r"(mbar_addr)); -} - - -__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { - asm volatile( - "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" - :: "r"(mbar_addr), "r"(bytes) : "memory"); -} - - -__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7," - " %8, %9, %10, %11, %12, %13, %14, %15," - " %16, %17, %18, %19, %20, %21, %22, %23," - " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), - "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), - "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), - "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), - "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), - "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), - "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ float approx_exp2(float x) { - float y; - asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float approx_rcp(float x) { - float y; - asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float max_noftz(float a, float b) { - float c; - asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); - return c; -} - - -__device__ __forceinline__ float warp_reduce_max(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); - return val; -} - - -__device__ __forceinline__ float warp_reduce_sum(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val += __shfl_xor_sync(0xFFFFFFFF, val, offset); - return val; -} - - -__device__ __forceinline__ float row_max_reduce(float2 acc) { - return max_noftz(acc.x, acc.y); -} - - -__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { - #pragma unroll - for (int j = 0; j < 16; j++) { - if (j % 2 == 0) - acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); - else - acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); - } -} - - -__device__ __forceinline__ float2 ex2_emulation_f32x2_value(float2 value) { - const float c0 = 1.0f, c1 = 0.695146143436431884765625f; - const float c2 = 0.227564394474029541015625f, c3 = 0.077119089663028717041015625f; - const float magic = 12582912.0f; - float x0 = max_noftz(value.x, -127.0f), x1 = max_noftz(value.y, -127.0f); - float2 xc2 = make_float2(x0, x1), magic2 = make_float2(magic, magic); - float2 xr2; - asm("add.rm.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xr2) - : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&magic2)); - float2 c3_2 = make_float2(c3, c3), c2_2 = make_float2(c2, c2); - float2 c1_2 = make_float2(c1, c1), c0_2 = make_float2(c0, c0); - float2 xrb2, xfrac2; - asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xrb2) - : "l"(*(unsigned long long*)&xr2), "l"(*(unsigned long long*)&magic2)); - asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xfrac2) - : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&xrb2)); - float2 poly2; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&c3_2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c2_2)); - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c1_2)); - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c0_2)); - int x0r_i, x1r_i, p0_i, p1_i; - asm("mov.b64 {%0, %1}, %2;" : "=r"(x0r_i), "=r"(x1r_i) : "l"(*(unsigned long long*)&xr2)); - asm("mov.b64 {%0, %1}, %2;" : "=r"(p0_i), "=r"(p1_i) : "l"(*(unsigned long long*)&poly2)); - float r0, r1; - asm("mov.b32 %0, %1;" : "=f"(r0) : "r"((x0r_i << 23) + p0_i)); - asm("mov.b32 %0, %1;" : "=f"(r1) : "r"((x1r_i << 23) + p1_i)); - return make_float2(r0, r1); -} - -__device__ __forceinline__ void ex2_emulation_f32x2(float* x0_ptr, float* x1_ptr) { - float2 result = ex2_emulation_f32x2_value(make_float2(*x0_ptr, *x1_ptr)); - *x0_ptr = result.x; *x1_ptr = result.y; -} - -__device__ __forceinline__ void softmax_frag_exp2_cast( - float* sv, uint32_t* pv, int use_emu) -{ - #pragma unroll - for (int j = 0; j < 16; j++) { - if (use_emu && j >= 12) - ex2_emulation_f32x2(&sv[j*2], &sv[j*2+1]); - else { - sv[j*2] = approx_exp2(sv[j*2]); - sv[j*2+1] = approx_exp2(sv[j*2+1]); - } - } - #pragma unroll - for (int j = 0; j < 16; j++) { - __nv_bfloat162 bf = __float22bfloat162_rn({sv[j*2], sv[j*2+1]}); - pv[j] = reinterpret_cast(bf); - } -} - - - -__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { - const float2* sv2 = reinterpret_cast(sv); - #pragma unroll - for (int j = 0; j < 16; j++) { - asm("add.f32x2 %0, %1, %2;" - : "+l"(reinterpret_cast(*acc)) - : "l"(reinterpret_cast(*acc)), - "l"(reinterpret_cast(sv2[j]))); - } -} - - -__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { - asm("mul.rn.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { - asm("mul.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { - asm("add.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { - asm("add.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { - asm("sub.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { - asm("sub.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("add.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { - float2 r; - asm("sub.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("sub.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ void fma_scale_x32( - float* sv, const float2* scale2, const float2* neg_max2) -{ - float2* sv_2 = reinterpret_cast(sv); - #pragma unroll - for (int j = 0; j < 16; j++) - fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); -} - -__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("mul.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - - -// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) - -__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("add.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("add.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("add.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("add.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("add.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("add.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("add.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - - -__device__ __forceinline__ void fence_async_shared() { - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); -} - - -__device__ __forceinline__ uint64_t make_smem_desc(int addr) { - const int SBO = 1024; - return desc_encode(addr) - | (desc_encode(SBO) << 32ULL) - | (1ULL << 46ULL) - | (2ULL << 61ULL); -} - - -__device__ __forceinline__ void tma_2d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3}], [%4];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_4d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.4d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5}], [%6];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr) { - // Canonical .shared::cta form for non-multicast gather4, matching - // trtllm-gen / cuda_ptx and the PTX ISA qualifier order - // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, - // .shared::cluster is reserved for the multicast variant (ctaMask). - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { - asm volatile( - "tcgen05.commit.cta_group::1.mbarrier::arrive::one" - ".shared::cluster.b64 [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { - uint32_t result; - asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" - : "=r"(result) : "r"(val)); - return result; -} - -extern "C" { - -__global__ __launch_bounds__(256, 1) void -kernel_cake_dsv4_c967e7efe220160e28d3(CakeTensorMap const* tmap_q, CakeTensorMap const* tmap_swa_kv, CakeTensorMap const* tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int batch_size, int max_q_len, int ragged_query, int has_sinks) -{ - const int tid = threadIdx.x; - const int warp = make_warp_uniform(tid / 32); - const int lane = tid % 32; - - extern __shared__ __align__(1024) char smem_raw[]; - int smem; - smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); - - const int mbar_base = smem; - #define q_full_addr (mbar_base + 0) - #define kv_full_addr (mbar_base + 8) - #define s_full_addr (mbar_base + 40) - #define p_full_addr (mbar_base + 56) - #define pv_done_addr (mbar_base + 72) - #define tmem_dealloc_addr (mbar_base + 80) - - const int bid = blockIdx.x; - const int num_bids = gridDim.x; - - const int cta_rank = 0; - if (tid == 0) { - asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_q)) : "memory"); - asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_swa_kv)) : "memory"); - asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_compressed_kv)) : "memory"); - } - __syncthreads(); - - - // Kernel setup ops - __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); - const int smem_q_addr = smem + 1024; - __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); - const int smem_kv_addr = smem + 66560; - __nv_bfloat16* smem_v = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); - const int smem_v_addr = smem + 66560; - - // Mbarrier init (6 pipeline groups, 0 ordered-sequence groups, 11 barriers) - // Mbarriers at smem_raw[0..88) - - if (warp == 0) { - uint32_t leader = elect_sync(); - if (leader) { - // q_full: 1 barriers, init_count=1 - mbarrier_init(smem + 0, 1); - // --- pipeline 'kv_pipe' --- - // kv_full: 4 barriers, init_count=3 - mbarrier_init(smem + 8, 3); - mbarrier_init(smem + 16, 3); - mbarrier_init(smem + 24, 3); - mbarrier_init(smem + 32, 3); - // s_full: 2 barriers, init_count=1 - mbarrier_init(smem + 40, 1); - mbarrier_init(smem + 48, 1); - // p_full: 2 barriers, init_count=4 - mbarrier_init(smem + 56, 4); - mbarrier_init(smem + 64, 4); - // pv_done: 1 barriers, init_count=1 - mbarrier_init(smem + 72, 1); - // tmem_dealloc: 1 barriers, init_count=7 - mbarrier_init(smem + 80, 7); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - - __syncwarp(); - - // TMEM alloc (512 columns, 512 used) - volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 88); - if (warp == 0) { - int _tmem_hold = smem + 88; - asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); - __syncwarp(); - asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); - } - - __syncthreads(); - asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); - - const int taddr = tmem_addr_storage[0]; - - // Kernel post-init ops - const int tmem_tmem = taddr; - - // ---- Ordered hardware-WG register redistribution ---- - // Dec phase frees registers before any WG attempts inc. - if (warp >= 4 && warp <= 7) { - asm volatile("setmaxnreg.dec.sync.aligned.u32 64;"); - } - - // ---- Role: softmax ---- - if (warp <= 3) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 176;"); - { // softmax_main - float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; - int query_idx = blockIdx.x; - int _max_2 = ((sparse_topk_lens[query_idx] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx] + sparse_topk_lens_offset) : (0)); - int _min_2 = ((_max_2) < (sparse_topk) ? (_max_2) : (sparse_topk)); - int active_topk = _min_2; - int tile_count = (active_topk + 128 - 1) / 128; - int query_batch = query_idx / max_q_len; - int query_offset = query_idx - query_batch * max_q_len; - int query_length = max_q_len; - if (ragged_query != 0) { - query_batch = 0; - #pragma unroll 1 - for (int batch = 0; batch < batch_size; batch++) { - if (query_idx >= cum_seq_lens_q[batch + 1]) { - query_batch = batch + 1; - } - } - int query_begin = cum_seq_lens_q[query_batch]; - query_length = cum_seq_lens_q[query_batch + 1] - query_begin; - query_offset = query_idx - query_begin; - } - int swa_valid_columns = seq_lens[query_batch] - query_length + query_offset + 1; - if (swa_valid_columns < 0) { - swa_valid_columns = 0; - } - if (swa_valid_columns > 128) { - swa_valid_columns = 128; - } - const int warp_in_compute = warp; - const int tmem_row_origin = warp_in_compute * 32; - const int logical_row_origin = warp_in_compute * 16; - const int my_row = logical_row_origin + lane % 16; - const int col_half = lane / 16; - float row_max = -CAKE_INF; - float row_sum = 0.0f; - if (has_sinks != 0 && my_row < num_heads) { - row_max = sinks[my_row] * 1.4426950408889634f / softmax_scale_log2; - row_sum = 1.0f; - } - #pragma unroll - for (int tile = 0; tile < 1; tile++) { - if (tile_count > tile) { - int phase = tile & 1; - int wait_phase = tile >> 1 & 1; - mbarrier_wait(s_full_addr + (phase) * 8, wait_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - int score_col = ((phase != 0) ? 128 : 0); - int score_addr = taddr + (unsigned int)score_col + (unsigned int)(tmem_row_origin << 16); - float _tmem_load_0[64]; - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_0[0])), "=r"(*reinterpret_cast(&_tmem_load_0[1])), "=r"(*reinterpret_cast(&_tmem_load_0[2])), "=r"(*reinterpret_cast(&_tmem_load_0[3])), "=r"(*reinterpret_cast(&_tmem_load_0[4])), "=r"(*reinterpret_cast(&_tmem_load_0[5])), "=r"(*reinterpret_cast(&_tmem_load_0[6])), "=r"(*reinterpret_cast(&_tmem_load_0[7])), "=r"(*reinterpret_cast(&_tmem_load_0[8])), "=r"(*reinterpret_cast(&_tmem_load_0[9])), "=r"(*reinterpret_cast(&_tmem_load_0[10])), "=r"(*reinterpret_cast(&_tmem_load_0[11])), "=r"(*reinterpret_cast(&_tmem_load_0[12])), "=r"(*reinterpret_cast(&_tmem_load_0[13])), "=r"(*reinterpret_cast(&_tmem_load_0[14])), "=r"(*reinterpret_cast(&_tmem_load_0[15])), "=r"(*reinterpret_cast(&_tmem_load_0[16])), "=r"(*reinterpret_cast(&_tmem_load_0[17])), "=r"(*reinterpret_cast(&_tmem_load_0[18])), "=r"(*reinterpret_cast(&_tmem_load_0[19])), "=r"(*reinterpret_cast(&_tmem_load_0[20])), "=r"(*reinterpret_cast(&_tmem_load_0[21])), "=r"(*reinterpret_cast(&_tmem_load_0[22])), "=r"(*reinterpret_cast(&_tmem_load_0[23])), "=r"(*reinterpret_cast(&_tmem_load_0[24])), "=r"(*reinterpret_cast(&_tmem_load_0[25])), "=r"(*reinterpret_cast(&_tmem_load_0[26])), "=r"(*reinterpret_cast(&_tmem_load_0[27])), "=r"(*reinterpret_cast(&_tmem_load_0[28])), "=r"(*reinterpret_cast(&_tmem_load_0[29])), "=r"(*reinterpret_cast(&_tmem_load_0[30])), "=r"(*reinterpret_cast(&_tmem_load_0[31])) - : "r"(score_addr)); - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_0[32])), "=r"(*reinterpret_cast(&_tmem_load_0[33])), "=r"(*reinterpret_cast(&_tmem_load_0[34])), "=r"(*reinterpret_cast(&_tmem_load_0[35])), "=r"(*reinterpret_cast(&_tmem_load_0[36])), "=r"(*reinterpret_cast(&_tmem_load_0[37])), "=r"(*reinterpret_cast(&_tmem_load_0[38])), "=r"(*reinterpret_cast(&_tmem_load_0[39])), "=r"(*reinterpret_cast(&_tmem_load_0[40])), "=r"(*reinterpret_cast(&_tmem_load_0[41])), "=r"(*reinterpret_cast(&_tmem_load_0[42])), "=r"(*reinterpret_cast(&_tmem_load_0[43])), "=r"(*reinterpret_cast(&_tmem_load_0[44])), "=r"(*reinterpret_cast(&_tmem_load_0[45])), "=r"(*reinterpret_cast(&_tmem_load_0[46])), "=r"(*reinterpret_cast(&_tmem_load_0[47])), "=r"(*reinterpret_cast(&_tmem_load_0[48])), "=r"(*reinterpret_cast(&_tmem_load_0[49])), "=r"(*reinterpret_cast(&_tmem_load_0[50])), "=r"(*reinterpret_cast(&_tmem_load_0[51])), "=r"(*reinterpret_cast(&_tmem_load_0[52])), "=r"(*reinterpret_cast(&_tmem_load_0[53])), "=r"(*reinterpret_cast(&_tmem_load_0[54])), "=r"(*reinterpret_cast(&_tmem_load_0[55])), "=r"(*reinterpret_cast(&_tmem_load_0[56])), "=r"(*reinterpret_cast(&_tmem_load_0[57])), "=r"(*reinterpret_cast(&_tmem_load_0[58])), "=r"(*reinterpret_cast(&_tmem_load_0[59])), "=r"(*reinterpret_cast(&_tmem_load_0[60])), "=r"(*reinterpret_cast(&_tmem_load_0[61])), "=r"(*reinterpret_cast(&_tmem_load_0[62])), "=r"(*reinterpret_cast(&_tmem_load_0[63])) - : "r"(score_addr + 32)); - int valid_cols = active_topk - tile * 128 - col_half * 64; - if (valid_cols < 0) { - valid_cols = 0; - } - if (valid_cols > 64) { - valid_cols = 64; - } - if (tile == 0) { - int valid_swa_cols = swa_valid_columns - col_half * 64; - if (valid_swa_cols < 0) { - valid_swa_cols = 0; - } - if (valid_swa_cols < valid_cols) { - valid_cols = valid_swa_cols; - } - } - uint32_t _slice_lo_mask_0; - { - int _lim_0 = valid_cols; - if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } - else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); - } - } - if (!(_slice_lo_mask_0 & (1u << 0))) _tmem_load_0[0] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 1))) _tmem_load_0[1] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 2))) _tmem_load_0[2] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 3))) _tmem_load_0[3] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 4))) _tmem_load_0[4] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 5))) _tmem_load_0[5] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 6))) _tmem_load_0[6] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 7))) _tmem_load_0[7] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 8))) _tmem_load_0[8] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 9))) _tmem_load_0[9] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 10))) _tmem_load_0[10] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 11))) _tmem_load_0[11] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 12))) _tmem_load_0[12] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 13))) _tmem_load_0[13] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 14))) _tmem_load_0[14] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 15))) _tmem_load_0[15] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 16))) _tmem_load_0[16] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 17))) _tmem_load_0[17] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 18))) _tmem_load_0[18] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 19))) _tmem_load_0[19] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 20))) _tmem_load_0[20] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 21))) _tmem_load_0[21] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 22))) _tmem_load_0[22] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 23))) _tmem_load_0[23] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 24))) _tmem_load_0[24] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 25))) _tmem_load_0[25] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 26))) _tmem_load_0[26] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 27))) _tmem_load_0[27] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 28))) _tmem_load_0[28] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 29))) _tmem_load_0[29] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 30))) _tmem_load_0[30] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 31))) _tmem_load_0[31] = -CAKE_INF; - uint32_t _slice_lo_mask_1; - { - int _lim_1 = valid_cols - 32; - if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } - else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); - } - } - if (!(_slice_lo_mask_1 & (1u << 0))) _tmem_load_0[32] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 1))) _tmem_load_0[33] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 2))) _tmem_load_0[34] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 3))) _tmem_load_0[35] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 4))) _tmem_load_0[36] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 5))) _tmem_load_0[37] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 6))) _tmem_load_0[38] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 7))) _tmem_load_0[39] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 8))) _tmem_load_0[40] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 9))) _tmem_load_0[41] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 10))) _tmem_load_0[42] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 11))) _tmem_load_0[43] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 12))) _tmem_load_0[44] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 13))) _tmem_load_0[45] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 14))) _tmem_load_0[46] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 15))) _tmem_load_0[47] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 16))) _tmem_load_0[48] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 17))) _tmem_load_0[49] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 18))) _tmem_load_0[50] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 19))) _tmem_load_0[51] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 20))) _tmem_load_0[52] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 21))) _tmem_load_0[53] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 22))) _tmem_load_0[54] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 23))) _tmem_load_0[55] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 24))) _tmem_load_0[56] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 25))) _tmem_load_0[57] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 26))) _tmem_load_0[58] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 27))) _tmem_load_0[59] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 28))) _tmem_load_0[60] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 29))) _tmem_load_0[61] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 30))) _tmem_load_0[62] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 31))) _tmem_load_0[63] = -CAKE_INF; - int* index_row = swa_indices + (query_idx * swa_index_stride + tile * 128 + col_half * 64); - #pragma unroll - for (int chunk = 0; chunk < 64; chunk += 4) { - if (valid_cols > chunk) { - int _vec_load_1[4]; - { - const int4* _ivptr_2 = reinterpret_cast(index_row + chunk + 0); - int4 _ivld_2; - _ivld_2 = *_ivptr_2; - _vec_load_1[0 + 0] = _ivld_2.x; - _vec_load_1[0 + 1] = _ivld_2.y; - _vec_load_1[0 + 2] = _ivld_2.z; - _vec_load_1[0 + 3] = _ivld_2.w; - } - #pragma unroll - for (int i = 0; i < 4; i++) { - if (_vec_load_1[i] < 0) { - _tmem_load_0[chunk + i] = -CAKE_INF; - } - } - } - } - float2 _reg_reduce_max2_3 = {-CAKE_INF, -CAKE_INF}; - row_max_x32_accum(&_tmem_load_0[0], _reg_reduce_max2_3); - row_max_x32_accum(&_tmem_load_0[32], _reg_reduce_max2_3); - float _tmem_load_0_max = row_max_reduce(_reg_reduce_max2_3); - float tile_max = _tmem_load_0_max; - float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, tile_max, 16); - float _max_3 = max_noftz(tile_max, _shfl_xor_0); - tile_max = _max_3; - float _max_4 = max_noftz(row_max, tile_max); - float new_max = _max_4; - float _fma_0 = __fmaf_rn(row_max, softmax_scale_log2, (-new_max) * softmax_scale_log2); - float delta = _fma_0; - float _exp2_0 = approx_exp2(delta); - float acc_scale = ((row_max > -CAKE_INF) ? _exp2_0 : 1.0f); - float safe_max = ((new_max == -CAKE_INF) ? 0.0f : new_max); - float max_scaled = safe_max * softmax_scale_log2; - const float2 _fma_b2_4 = {softmax_scale_log2, softmax_scale_log2}; - const float2 _fma_c2_5 = {-max_scaled, -max_scaled}; - #pragma unroll - for (int _lf = 0; _lf < 32; _lf++) - fma_f32x2_inplace(&reinterpret_cast(_tmem_load_0)[_lf], _fma_b2_4, _fma_c2_5); - #pragma unroll - for (int _le = 0; _le < 64; _le++) { - _tmem_load_0[_le] = approx_exp2(_tmem_load_0[_le]); - } - unsigned int packed_p[32]; - #pragma unroll - for (int _lp = 0; _lp < 32; _lp++) { - __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_0[_lp*2 + 0], _tmem_load_0[_lp*2+1 + 0])); - packed_p[_lp] = *(uint32_t*)&_bf2; - } - int p_addr = taddr + (unsigned int)score_col + 64 + (unsigned int)(tmem_row_origin << 16); - asm volatile( - "tcgen05.st.sync.aligned.16x32bx2.x16.b32" - " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" - :: "r"(p_addr), "r"(*reinterpret_cast(&packed_p[0])), "r"(*reinterpret_cast(&packed_p[1])), "r"(*reinterpret_cast(&packed_p[2])), "r"(*reinterpret_cast(&packed_p[3])), "r"(*reinterpret_cast(&packed_p[4])), "r"(*reinterpret_cast(&packed_p[5])), "r"(*reinterpret_cast(&packed_p[6])), "r"(*reinterpret_cast(&packed_p[7])), "r"(*reinterpret_cast(&packed_p[8])), "r"(*reinterpret_cast(&packed_p[9])), "r"(*reinterpret_cast(&packed_p[10])), "r"(*reinterpret_cast(&packed_p[11])), "r"(*reinterpret_cast(&packed_p[12])), "r"(*reinterpret_cast(&packed_p[13])), "r"(*reinterpret_cast(&packed_p[14])), "r"(*reinterpret_cast(&packed_p[15]))); - asm volatile( - "tcgen05.st.sync.aligned.16x32bx2.x16.b32" - " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" - :: "r"(p_addr + 16), "r"(*reinterpret_cast(&(packed_p + 16)[0])), "r"(*reinterpret_cast(&(packed_p + 16)[1])), "r"(*reinterpret_cast(&(packed_p + 16)[2])), "r"(*reinterpret_cast(&(packed_p + 16)[3])), "r"(*reinterpret_cast(&(packed_p + 16)[4])), "r"(*reinterpret_cast(&(packed_p + 16)[5])), "r"(*reinterpret_cast(&(packed_p + 16)[6])), "r"(*reinterpret_cast(&(packed_p + 16)[7])), "r"(*reinterpret_cast(&(packed_p + 16)[8])), "r"(*reinterpret_cast(&(packed_p + 16)[9])), "r"(*reinterpret_cast(&(packed_p + 16)[10])), "r"(*reinterpret_cast(&(packed_p + 16)[11])), "r"(*reinterpret_cast(&(packed_p + 16)[12])), "r"(*reinterpret_cast(&(packed_p + 16)[13])), "r"(*reinterpret_cast(&(packed_p + 16)[14])), "r"(*reinterpret_cast(&(packed_p + 16)[15]))); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - if (elect_sync()) { - mbarrier_arrive(p_full_addr + (phase) * 8); - } - float2 _reg_reduce_sum2_6 = make_float2(0.0f, 0.0f); - softmax_block_sum(&_tmem_load_0[0], &_reg_reduce_sum2_6); - softmax_block_sum(&_tmem_load_0[32], &_reg_reduce_sum2_6); - float _tmem_load_0_sum = _reg_reduce_sum2_6.x + _reg_reduce_sum2_6.y; - float block_sum = _tmem_load_0_sum; - float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, block_sum, 16); - block_sum = block_sum + _shfl_xor_1; - float _fma_1 = __fmaf_rn(row_sum, acc_scale, block_sum); - row_sum = _fma_1; - row_max = new_max; - } - } - if (sparse_topk == 128) { - mbarrier_wait(pv_done_addr, 0); - asm volatile("tcgen05.fence::after_thread_sync;"); - float output_scale = bmm2_scale[0]; - float _rcp_0 = approx_rcp(row_sum); - float inv_sum = ((row_sum > 0.0f) ? _rcp_0 : 0.0f); - int output_base = (query_idx * num_heads + my_row) * 512 + blockIdx.y * 256; - #pragma unroll - for (int v_stage = 0; v_stage < 2; v_stage++) { - int output_col = 256 + v_stage * 128; - int output_addr = taddr + (unsigned int)output_col + (unsigned int)(tmem_row_origin << 16); - float _tmem_load_1[64]; - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_1[0])), "=r"(*reinterpret_cast(&_tmem_load_1[1])), "=r"(*reinterpret_cast(&_tmem_load_1[2])), "=r"(*reinterpret_cast(&_tmem_load_1[3])), "=r"(*reinterpret_cast(&_tmem_load_1[4])), "=r"(*reinterpret_cast(&_tmem_load_1[5])), "=r"(*reinterpret_cast(&_tmem_load_1[6])), "=r"(*reinterpret_cast(&_tmem_load_1[7])), "=r"(*reinterpret_cast(&_tmem_load_1[8])), "=r"(*reinterpret_cast(&_tmem_load_1[9])), "=r"(*reinterpret_cast(&_tmem_load_1[10])), "=r"(*reinterpret_cast(&_tmem_load_1[11])), "=r"(*reinterpret_cast(&_tmem_load_1[12])), "=r"(*reinterpret_cast(&_tmem_load_1[13])), "=r"(*reinterpret_cast(&_tmem_load_1[14])), "=r"(*reinterpret_cast(&_tmem_load_1[15])), "=r"(*reinterpret_cast(&_tmem_load_1[16])), "=r"(*reinterpret_cast(&_tmem_load_1[17])), "=r"(*reinterpret_cast(&_tmem_load_1[18])), "=r"(*reinterpret_cast(&_tmem_load_1[19])), "=r"(*reinterpret_cast(&_tmem_load_1[20])), "=r"(*reinterpret_cast(&_tmem_load_1[21])), "=r"(*reinterpret_cast(&_tmem_load_1[22])), "=r"(*reinterpret_cast(&_tmem_load_1[23])), "=r"(*reinterpret_cast(&_tmem_load_1[24])), "=r"(*reinterpret_cast(&_tmem_load_1[25])), "=r"(*reinterpret_cast(&_tmem_load_1[26])), "=r"(*reinterpret_cast(&_tmem_load_1[27])), "=r"(*reinterpret_cast(&_tmem_load_1[28])), "=r"(*reinterpret_cast(&_tmem_load_1[29])), "=r"(*reinterpret_cast(&_tmem_load_1[30])), "=r"(*reinterpret_cast(&_tmem_load_1[31])) - : "r"(output_addr)); - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_1[32])), "=r"(*reinterpret_cast(&_tmem_load_1[33])), "=r"(*reinterpret_cast(&_tmem_load_1[34])), "=r"(*reinterpret_cast(&_tmem_load_1[35])), "=r"(*reinterpret_cast(&_tmem_load_1[36])), "=r"(*reinterpret_cast(&_tmem_load_1[37])), "=r"(*reinterpret_cast(&_tmem_load_1[38])), "=r"(*reinterpret_cast(&_tmem_load_1[39])), "=r"(*reinterpret_cast(&_tmem_load_1[40])), "=r"(*reinterpret_cast(&_tmem_load_1[41])), "=r"(*reinterpret_cast(&_tmem_load_1[42])), "=r"(*reinterpret_cast(&_tmem_load_1[43])), "=r"(*reinterpret_cast(&_tmem_load_1[44])), "=r"(*reinterpret_cast(&_tmem_load_1[45])), "=r"(*reinterpret_cast(&_tmem_load_1[46])), "=r"(*reinterpret_cast(&_tmem_load_1[47])), "=r"(*reinterpret_cast(&_tmem_load_1[48])), "=r"(*reinterpret_cast(&_tmem_load_1[49])), "=r"(*reinterpret_cast(&_tmem_load_1[50])), "=r"(*reinterpret_cast(&_tmem_load_1[51])), "=r"(*reinterpret_cast(&_tmem_load_1[52])), "=r"(*reinterpret_cast(&_tmem_load_1[53])), "=r"(*reinterpret_cast(&_tmem_load_1[54])), "=r"(*reinterpret_cast(&_tmem_load_1[55])), "=r"(*reinterpret_cast(&_tmem_load_1[56])), "=r"(*reinterpret_cast(&_tmem_load_1[57])), "=r"(*reinterpret_cast(&_tmem_load_1[58])), "=r"(*reinterpret_cast(&_tmem_load_1[59])), "=r"(*reinterpret_cast(&_tmem_load_1[60])), "=r"(*reinterpret_cast(&_tmem_load_1[61])), "=r"(*reinterpret_cast(&_tmem_load_1[62])), "=r"(*reinterpret_cast(&_tmem_load_1[63])) - : "r"(output_addr + 32)); - int col_base = v_stage * 128 + col_half * 64; - if (my_row < num_heads) { - #pragma unroll - for (int offset = 0; offset < 64; offset += 8) { - { - const float2 _prescale2_7 = {inv_sum * output_scale, inv_sum * output_scale}; - #if __CUDA_ARCH__ >= 1000 - #pragma unroll - for (int _ps = 0; _ps < 4; _ps++) - mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_1[offset])[_ps], _prescale2_7); - #else - #pragma unroll - for (int _ps = 0; _ps < 8; _ps++) - _tmem_load_1[offset + _ps] *= inv_sum * output_scale; - #endif - __nv_bfloat162 _pk[4]; - _pk[0] = __floats2bfloat162_rn(_tmem_load_1[offset + 0], _tmem_load_1[offset + 1]); - _pk[1] = __floats2bfloat162_rn(_tmem_load_1[offset + 2], _tmem_load_1[offset + 3]); - _pk[2] = __floats2bfloat162_rn(_tmem_load_1[offset + 4], _tmem_load_1[offset + 5]); - _pk[3] = __floats2bfloat162_rn(_tmem_load_1[offset + 6], _tmem_load_1[offset + 7]); - *reinterpret_cast(&((__nv_bfloat16*)(O + (output_base + col_base + offset)))[0]) = *reinterpret_cast(&_pk[0]); - } - } - } - } - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - if (elect_sync()) { - mbarrier_arrive(tmem_dealloc_addr); - } - } - } - // ---- Role: mma_warp ---- - if (warp == 4) { - { // mma_warp_main - int query_idx_1 = blockIdx.x; - int _max_1 = ((sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) : (0)); - int _min_1 = ((_max_1) < (sparse_topk) ? (_max_1) : (sparse_topk)); - int active_topk_1 = _min_1; - int tile_count_1 = (active_topk_1 + 128 - 1) / 128; - unsigned int _phase_q_full_0 = 0; - mbarrier_wait(q_full_addr, _phase_q_full_0); - _phase_q_full_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - #pragma unroll - for (int tile_1 = 0; tile_1 < 1; tile_1++) { - if (tile_count_1 > tile_1) { - int phase_1 = tile_1 & 1; - int score_col_1 = ((phase_1 != 0) ? 128 : 0); - #pragma unroll - for (int k_stage = 0; k_stage < 4; k_stage++) { - mbarrier_wait(kv_full_addr + (k_stage) * 8, phase_1); - asm volatile("tcgen05.fence::after_thread_sync;"); - int _mma_a_lo_0 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (k_stage * 2) * 512); - int _mma_b_lo_0 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (k_stage) * 2048); - { - uint64_t _mma_ss_a_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_0); - uint64_t _mma_ss_b_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_0); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, ((k_stage == 0) ? 0 : 1)); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); - } - } - int _mma_a_lo_1 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (k_stage * 2 + 1) * 512); - int _mma_b_lo_1 = make_warp_uniform((((smem_kv_addr + 16384) >> 4) & 0x3FFF) + (k_stage) * 2048); - { - uint64_t _mma_ss_a_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_1); - uint64_t _mma_ss_b_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_1); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); - } - } - if (k_stage == 3) { - elect_commit(s_full_addr + (phase_1) * 8); - } - } - int wait_phase_1 = tile_1 >> 1 & 1; - mbarrier_wait(p_full_addr + (phase_1) * 8, wait_phase_1); - asm volatile("tcgen05.fence::after_thread_sync;"); - #pragma unroll - for (int v_stage_1 = 0; v_stage_1 < 2; v_stage_1++) { - int _mma_b_lo_2 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (blockIdx.y * 2 + v_stage_1) * 2048); - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 dhi, blo, ta, id;\n\t" - ".reg .b64 db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "" - "mov.b32 dhi, 0x40004040;\n\t" - "mov.b32 id, 69272720;\n\t" - "mov.b32 ta, %2;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p0;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "}\n" - :: "r"((tmem_tmem + (256 + v_stage_1 * 128))), "r"(_mma_b_lo_2), "r"(tmem_tmem + (score_col_1 + 64)), "r"(0)); - } - elect_commit(pv_done_addr); - } - } - unsigned int _phase_tmem_dealloc_0 = 0; - mbarrier_wait(tmem_dealloc_addr, _phase_tmem_dealloc_0); - _phase_tmem_dealloc_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); - asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); - } - } - // ---- Role: load_warp ---- - if (warp >= 5 && warp <= 7) { - { // load_warp_main - const int load_warp_rank = warp - 5; - int query_idx_2 = blockIdx.x; - int _max_0 = ((sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) : (0)); - int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); - int active_topk_2 = _min_0; - int tile_count_2 = (active_topk_2 + 128 - 1) / 128; - unsigned int load_kv_stage = 0; - unsigned int load_kv_phase = 1; - if (load_warp_rank == 0) { - if (elect_sync()) { - mbarrier_arrive_expect_tx(q_full_addr, 65536); - #pragma unroll - for (int q_stage = 0; q_stage < 8; q_stage += 8) { - tma_4d_gmem2smem(smem_q_addr + (unsigned int)(q_stage * 8192), tmap_q, 0, 0, q_stage, query_idx_2, q_full_addr); - } - } - } - int group = load_warp_rank + lane * 3; - int row0 = 0; - int row1 = 0; - int row2 = 0; - int row3 = 0; - if (group < 32) { - int index_base = group * 4; - int* index_row_1 = swa_indices + (query_idx_2 * swa_index_stride); - int raw_rows[4]; - if (index_base + 3 < sparse_topk) { - int _vec_load_0[4]; - { - const int4* _ivptr_0 = reinterpret_cast(index_row_1 + index_base + 0); - int4 _ivld_0; - _ivld_0 = *_ivptr_0; - _vec_load_0[0 + 0] = _ivld_0.x; - _vec_load_0[0 + 1] = _ivld_0.y; - _vec_load_0[0 + 2] = _ivld_0.z; - _vec_load_0[0 + 3] = _ivld_0.w; - } - #pragma unroll - for (int i_1 = 0; i_1 < 4; i_1++) { - raw_rows[i_1] = _vec_load_0[i_1]; - } - } else { - #pragma unroll - for (int i_2 = 0; i_2 < 4; i_2++) { - raw_rows[i_2] = -1; - if (index_base + i_2 < sparse_topk) { - raw_rows[i_2] = index_row_1[index_base + i_2]; - } - } - } - row0 = ((raw_rows[0] >= 0) ? raw_rows[0] : 0); - row1 = ((raw_rows[1] >= 0) ? raw_rows[1] : 0); - row2 = ((raw_rows[2] >= 0) ? raw_rows[2] : 0); - row3 = ((raw_rows[3] >= 0) ? raw_rows[3] : 0); - } - #pragma unroll - for (int tile_2 = 0; tile_2 < 1; tile_2++) { - if (tile_count_2 > tile_2) { - #pragma unroll - for (int k_stage_1 = 0; k_stage_1 < 4; k_stage_1++) { - if (elect_sync()) { - int producer_groups = ((load_warp_rank < 2) ? 11 : 10); - mbarrier_arrive_expect_tx(kv_full_addr + (load_kv_stage) * 8, producer_groups * 512 * 2); - } - if (group < 32) { - int dst_k = smem_kv_addr + load_kv_stage * 32768; - if (tile_2 == 0) { - tma_gather4_gmem2smem(dst_k + group * 512, tmap_swa_kv, k_stage_1 * 128, row0, row1, row2, row3, kv_full_addr + (load_kv_stage) * 8); - tma_gather4_gmem2smem(dst_k + 16384 + group * 512, tmap_swa_kv, k_stage_1 * 128 + 64, row0, row1, row2, row3, kv_full_addr + (load_kv_stage) * 8); - } else { - tma_gather4_gmem2smem(dst_k + group * 512, tmap_compressed_kv, k_stage_1 * 128, row0, row1, row2, row3, kv_full_addr + (load_kv_stage) * 8); - tma_gather4_gmem2smem(dst_k + 16384 + group * 512, tmap_compressed_kv, k_stage_1 * 128 + 64, row0, row1, row2, row3, kv_full_addr + (load_kv_stage) * 8); - } - } - load_kv_stage += 1; - if (load_kv_stage == 4) { load_kv_stage = 0; load_kv_phase ^= 1; } - } - } - } - if (elect_sync()) { - mbarrier_arrive(tmem_dealloc_addr); - } - } - } - - // Cleanup -} - -} // extern "C" diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_d23652aae3370c30d978_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_d23652aae3370c30d978_binding.cu deleted file mode 100644 index 51c14700a2a..00000000000 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_d23652aae3370c30d978_binding.cu +++ /dev/null @@ -1,526 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -// Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_d23652aae3370c30d978'. -#include -#include -#include -#include - -#include "tvm_ffi_utils.h" - -#include -#include -#include -#include -#include -#include -#include - -struct CakeTensorMap; -extern "C" __global__ void kernel_cake_dsv4_d23652aae3370c30d978(CakeTensorMap const* tmap_q, CakeTensorMap const* tmap_swa_kv, CakeTensorMap const* tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int max_q_len, int batch_size, int has_sinks); - - -namespace cake_host_shim_9b1b8d760815189c { - -using tvm::ffi::Optional; -using tvm::ffi::TensorView; - -class ScopedCudaDevice { - public: - explicit ScopedCudaDevice(int device_id) { - cudaError_t error = cudaGetDevice(&previous_device_); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed before host-shim launch: cudaError=" - << static_cast(error); - if (previous_device_ != device_id) { - error = cudaSetDevice(device_id); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaSetDevice failed before host-shim launch for cuda:" - << device_id << ": cudaError=" << static_cast(error); - restore_ = true; - } - } - - ScopedCudaDevice(const ScopedCudaDevice&) = delete; - ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; - - ~ScopedCudaDevice() noexcept { - if (restore_) { - (void)cudaSetDevice(previous_device_); - } - } - - private: - int previous_device_ = -1; - bool restore_ = false; -}; - -inline int64_t CakeDeviceMultiprocessorCount(int device_id) { - constexpr int kMaxCachedCudaDevices = 64; - TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) - << "physical-SM-count cache does not cover cuda:" << device_id; - static std::atomic count_by_device[kMaxCachedCudaDevices]{}; - int cached = count_by_device[device_id].load(std::memory_order_acquire); - if (cached > 0) return cached; - - int count = 0; - cudaError_t error = cudaDeviceGetAttribute( - &count, cudaDevAttrMultiProcessorCount, device_id); - TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) - << "querying multiProcessorCount failed for cuda:" << device_id - << ": cudaError=" << static_cast(error) << ", count=" << count; - // Concurrent first launches may repeat the immutable device query, but all - // publication is atomic and every later hot-path lookup is one acquire load. - count_by_device[device_id].store(count, std::memory_order_release); - return count; -} - -inline void CheckCudaTensor(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) - << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; -} - -inline void CheckSameCudaDevice( - const TensorView& t, - const TensorView& reference, - const char* name, - const char* reference_name) { - TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) - << name << " must be on the same CUDA device as " << reference_name - << ": got cuda:" << t.device().device_id - << " versus cuda:" << reference.device().device_id; -} - -inline void CheckCurrentCudaDevice( - const TensorView& reference, - const char* reference_name) { - int current_device = -1; - cudaError_t error = cudaGetDevice(¤t_device); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed while validating " << reference_name - << ": cudaError=" << static_cast(error); - TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) - << "current CUDA device must match " << reference_name - << ": current=cuda:" << current_device - << ", tensor=cuda:" << reference.device().device_id; -} - -inline void CheckContiguous(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; -} - -inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { - DLDataType d = t.dtype(); - TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) - << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits - << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits - << ", lanes=" << (int)d.lanes << ")"; -} - -// A logical axis.outer(trailing) folds every source dim above the trailing -// dimensions. Shape products are independent of physical strides, so verify -// the leading dimensions form one dense row-major chain instead of inventing -// a "folded stride". The descriptor reads its exact adjacent physical step -// separately through stride[-(trailing + 1)]. -inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { - TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) - << name << " cannot fold leading dimensions above " << trailing - << " trailing dims from ndim=" << t.ndim(); - int outer_last = t.ndim() - trailing - 1; - if (outer_last <= 0) { - return; - } - int64_t step = t.stride(outer_last); - TVM_FFI_CHECK(step > 0, ValueError) - << name << " physical strides must be positive"; - int64_t expected = step; - for (int axis = outer_last - 1; axis >= 0; --axis) { - expected *= t.size(axis + 1); - if (t.size(axis) > 1) { - TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) - << name << " leading dims are not physically foldable above " << trailing - << " trailing dims: stride(" << axis << ")=" << t.stride(axis) - << ", expected " << expected; - } - } -} - -// Explicit immutable caller storage: initialize once, outside graph capture, -// exactly as the production TmaDeviceSlot path. The owner must retain the -// allocation and must never expose it as mutable scratch or rewrite its bytes. -template -static inline void PrepareImmutableCallerTmaWorkspace( - const tvm::ffi::Tensor& retained_workspace, - const CUtensorMap (&maps)[N], cudaStream_t stream) { - void* workspace = retained_workspace.data_ptr(); - CUcontext context = nullptr; - CUresult result = cuCtxGetCurrent(&context); - TVM_FFI_CHECK(result == CUDA_SUCCESS && context != nullptr, RuntimeError) - << "immutable TMA workspace requires an active CUDA context"; - CUcontext allocation_context = nullptr; - result = cuPointerGetAttribute( - &allocation_context, CU_POINTER_ATTRIBUTE_CONTEXT, - reinterpret_cast(workspace)); - TVM_FFI_CHECK(result == CUDA_SUCCESS && allocation_context == context, RuntimeError) - << "immutable TMA workspace must belong to the active CUDA context"; - unsigned long long buffer_id = 0; - result = cuPointerGetAttribute( - &buffer_id, CU_POINTER_ATTRIBUTE_BUFFER_ID, - reinterpret_cast(workspace)); - TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) - << "querying immutable TMA allocation identity failed"; - std::string key = std::to_string(reinterpret_cast(context)); - key += ":" + std::to_string(reinterpret_cast(workspace)); - key += ":" + std::to_string(buffer_id); - std::string descriptor_bytes(reinterpret_cast(maps), sizeof(maps)); - struct InitializedWorkspace { - tvm::ffi::Tensor owner; - std::string descriptor_bytes; - }; - static std::mutex mu; - // Keep the actual Tensor allocation alive: a caching allocator may reuse a - // freed suballocation without changing its enclosing CUDA buffer identity. - static auto* initialized = new std::unordered_map(); - std::lock_guard lock(mu); - auto it = initialized->find(key); - if (it != initialized->end()) { - TVM_FFI_CHECK(it->second.descriptor_bytes == descriptor_bytes, ValueError) - << "immutable TMA descriptors changed; provide new private descriptor storage"; - return; - } - CUstreamCaptureStatus capture_status = CU_STREAM_CAPTURE_STATUS_NONE; - result = cuStreamIsCapturing(reinterpret_cast(stream), &capture_status); - TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) - << "querying immutable TMA initialization capture state failed"; - TVM_FFI_CHECK(capture_status == CU_STREAM_CAPTURE_STATUS_NONE, RuntimeError) - << "immutable TMA storage must be initialized before CUDA Graph capture"; - // Synchronous host-to-device initialization completes before publishing the - // slot to any thread/stream, matching TmaDeviceSlot. The consumer retains its - // generated tensor-map acquire; the mutable upload's release fence is unchanged. - result = cuMemcpyHtoD(reinterpret_cast(workspace), maps, sizeof(maps)); - TVM_FFI_CHECK(result == CUDA_SUCCESS, RuntimeError) - << "initializing immutable caller TMA descriptors failed"; - initialized->emplace( - std::move(key), InitializedWorkspace{retained_workspace, std::move(descriptor_bytes)}); -} - -static inline void* CallerTmaWorkspaceSlot( - const tvm::ffi::TensorView& workspace, size_t slot) { - return static_cast(workspace.data_ptr()) + slot * sizeof(CUtensorMap); -} - -// 4D TMA descriptor for buffer 'tmap_q' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 3, ValueError) - << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - int64_t d2 = t.size(t.ndim() - 2); - int64_t d3 = t.size(t.ndim() - 3); - TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) - << "TMA source 'tmap_q' trailing dims must be positive"; - TVM_FFI_CHECK(d1 % 64 == 0, ValueError) - << "TMA source 'tmap_q' extent " << d1 - << " must divide exactly by " << 64; - uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(d3)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 2u <= global_dim[2] && 1u <= global_dim[3], ValueError) - << "TMA box (64, 8, 2, 1) exceeds resolved global dims for 'tmap_q'"; - int64_t carrier_stride_0 = d1; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; - int64_t carrier_stride_1 = 64; - TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; - TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; - TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; - int64_t carrier_stride_2 = (d2 * d1); - TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; - TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; - TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; - uint64_t global_strides[3] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - (uint64_t)((carrier_stride_1 * 16) / 8), - (uint64_t)((carrier_stride_2 * 16) / 8), - }; - uint32_t box_dim[4] = {64u, 8u, 2u, 1u}; - uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_swa_kv' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_swa_kv' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_kv'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_compressed_kv' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_compressed_kv(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_compressed_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_compressed_kv' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_compressed_kv' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_compressed_kv"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_compressed_kv' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_kv'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_kv') failed: CUresult=" << (int)r; - return tm; -} - -void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_compressed_kv, TensorView arg_O, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_seq_lens, TensorView arg_cum_seq_lens_q, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_max_q_len, int64_t arg_batch_size, int64_t arg_has_sinks, tvm::ffi::Tensor arg_tma_descriptor_workspace, int64_t grid_x, int64_t grid_y, int64_t grid_z) { - DLDevice dev = arg_tmap_q.device(); - ScopedCudaDevice device_guard(dev.device_id); - CheckCudaTensor(arg_tmap_q, "tmap_q"); - CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); - CheckContiguous(arg_tmap_q, "tmap_q"); - CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); - CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 4, 16, 1); - CheckCudaTensor(arg_tmap_compressed_kv, "tmap_compressed_kv"); - CheckDtype(arg_tmap_compressed_kv, "tmap_compressed_kv", 4, 16, 1); - CheckCudaTensor(arg_O, "O"); - CheckDtype(arg_O, "O", 4, 16, 1); - CheckContiguous(arg_O, "O"); - CheckCudaTensor(arg_swa_indices, "swa_indices"); - CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); - CheckContiguous(arg_swa_indices, "swa_indices"); - CheckCudaTensor(arg_compressed_indices, "compressed_indices"); - CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); - CheckContiguous(arg_compressed_indices, "compressed_indices"); - CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); - CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckCudaTensor(arg_seq_lens, "seq_lens"); - CheckDtype(arg_seq_lens, "seq_lens", 0, 32, 1); - CheckContiguous(arg_seq_lens, "seq_lens"); - CheckCudaTensor(arg_cum_seq_lens_q, "cum_seq_lens_q"); - CheckDtype(arg_cum_seq_lens_q, "cum_seq_lens_q", 0, 32, 1); - CheckContiguous(arg_cum_seq_lens_q, "cum_seq_lens_q"); - CheckCudaTensor(arg_sinks, "sinks"); - CheckDtype(arg_sinks, "sinks", 2, 32, 1); - CheckContiguous(arg_sinks, "sinks"); - CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); - CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); - CheckContiguous(arg_bmm1_scale, "bmm1_scale"); - CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); - CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); - CheckContiguous(arg_bmm2_scale, "bmm2_scale"); - TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) - << "scalar 'swa_index_stride' value " << arg_swa_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) - << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) - << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) - << "scalar 'sparse_topk' value " << arg_sparse_topk - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_max_q_len >= -2147483648LL && arg_max_q_len <= 2147483647LL, ValueError) - << "scalar 'max_q_len' value " << arg_max_q_len - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_batch_size >= -2147483648LL && arg_batch_size <= 2147483647LL, ValueError) - << "scalar 'batch_size' value " << arg_batch_size - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) - << "scalar 'has_sinks' value " << arg_has_sinks - << " is outside i32 range [-2147483648, 2147483647]"; - CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); - CheckSameCudaDevice(arg_tmap_compressed_kv, arg_tmap_q, "tmap_compressed_kv", "tmap_q"); - CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); - CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); - CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); - CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); - CheckSameCudaDevice(arg_seq_lens, arg_tmap_q, "seq_lens", "tmap_q"); - CheckSameCudaDevice(arg_cum_seq_lens_q, arg_tmap_q, "cum_seq_lens_q", "tmap_q"); - CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); - CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); - CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); - CheckCudaTensor(arg_tma_descriptor_workspace, "tma_descriptor_workspace"); - CheckSameCudaDevice(arg_tma_descriptor_workspace, arg_tmap_q, "tma_descriptor_workspace", "tmap_q"); - CheckContiguous(arg_tma_descriptor_workspace, "tma_descriptor_workspace"); - CheckDtype(arg_tma_descriptor_workspace, "tma_descriptor_workspace", 1, 8, 1); - TVM_FFI_CHECK(arg_tma_descriptor_workspace.numel() >= 384, ValueError) - << "tma_descriptor_workspace requires at least 384 bytes"; - TVM_FFI_CHECK( - reinterpret_cast(arg_tma_descriptor_workspace.data_ptr()) % 128 == 0, - ValueError) - << "tma_descriptor_workspace must be 128-byte aligned"; - TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) - << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y - << ", " << grid_z << ")"; - - cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); - CUtensorMap h_tmap_q = EncodeTma_tmap_q(arg_tmap_q); - CakeTensorMap const* p_tmap_q = static_cast(CallerTmaWorkspaceSlot(arg_tma_descriptor_workspace, 0u)); - CUtensorMap h_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); - CakeTensorMap const* p_tmap_swa_kv = static_cast(CallerTmaWorkspaceSlot(arg_tma_descriptor_workspace, 1u)); - CUtensorMap h_tmap_compressed_kv = EncodeTma_tmap_compressed_kv(arg_tmap_compressed_kv); - CakeTensorMap const* p_tmap_compressed_kv = static_cast(CallerTmaWorkspaceSlot(arg_tma_descriptor_workspace, 2u)); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); - int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); - int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); - int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); - int* p_seq_lens = static_cast(arg_seq_lens.data_ptr()); - int* p_cum_seq_lens_q = static_cast(arg_cum_seq_lens_q.data_ptr()); - float* p_sinks = static_cast(arg_sinks.data_ptr()); - float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); - float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); - int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; - int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; - int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; - int32_t v_sparse_topk = (int32_t)arg_sparse_topk; - int32_t v_max_q_len = (int32_t)arg_max_q_len; - int32_t v_batch_size = (int32_t)arg_batch_size; - int32_t v_has_sinks = (int32_t)arg_has_sinks; - void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_tmap_compressed_kv, &p_O, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_seq_lens, &p_cum_seq_lens_q, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_max_q_len, &v_batch_size, &v_has_sinks}; - - static std::mutex smem_mu; - static auto* smem_status_by_device = - new std::unordered_map(); - cudaError_t smem_status = cudaSuccess; - { - std::lock_guard lock(smem_mu); - auto it = smem_status_by_device->find(dev.device_id); - if (it == smem_status_by_device->end()) { - int current_device = -1; - smem_status = cudaGetDevice(¤t_device); - if (smem_status == cudaSuccess && current_device != dev.device_id) { - smem_status = cudaErrorInvalidDevice; - } - if (smem_status == cudaSuccess) { - smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_d23652aae3370c30d978), - cudaFuncAttributeMaxDynamicSharedMemorySize, - 218752); - } - smem_status_by_device->emplace(dev.device_id, smem_status); - } else { - smem_status = it->second; - } - } - TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_d23652aae3370c30d978 failed: " - << cudaGetErrorString(smem_status); - dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); - dim3 block(640u, 1u, 1u); - const CUtensorMap workspace_maps[] = {h_tmap_q, h_tmap_swa_kv, h_tmap_compressed_kv}; - PrepareImmutableCallerTmaWorkspace(arg_tma_descriptor_workspace, workspace_maps, stream); - - cudaLaunchConfig_t config{}; - config.gridDim = grid; - config.blockDim = block; - config.dynamicSmemBytes = 218752u; - config.stream = stream; - cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_d23652aae3370c30d978), kargs); - TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_d23652aae3370c30d978 failed: " - << cudaGetErrorString(launch_status); -} - -} // namespace cake_host_shim_9b1b8d760815189c - -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_9b1b8d760815189c::Run); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_d23652aae3370c30d978_kernel.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_d23652aae3370c30d978_kernel.cu deleted file mode 100644 index f3f50174875..00000000000 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_d23652aae3370c30d978_kernel.cu +++ /dev/null @@ -1,2704 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -typedef signed char int8_t; -typedef unsigned char uint8_t; -typedef unsigned short uint16_t; -typedef unsigned int uint32_t; -#if defined(__CUDACC_RTC__) -typedef unsigned long long uint64_t; -#else -typedef unsigned long uint64_t; -#endif -static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); -typedef signed int int32_t; -typedef short int int16_t; -struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; -struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; -static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); -static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); -template -struct __align__(128) CakeTensorMapPack { CakeTensorMap maps[N]; }; - -#if defined(__CUDACC_RTC__) -typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; -#else -#include -#endif - -static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); -static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); -#include -#include - -#define CAKE_INF CUDART_INF_F -#define TMEM_NCOLS 512 -#define TMEM_SOURCE_EXACT_SCRATCH_OFFSET 0 -#define NUM_Q_PIPE_STAGES 1 -#define NUM_KV_PIPE_STAGES 6 -#define NUM_PAGE_PIPE_STAGES 6 -#define NUM_SCORE_PIPE_STAGES 2 -#define NUM_STATS_PIPE_STAGES 2 -#define NUM_P_PIPE_STAGES 2 -#define NUM_O_PIPE_STAGES 1 -#define NUM_WORK_PIPE_STAGES 2 -#define NUM_THROTTLE_PIPE_STAGES 2 -#define SMEM_SMEM_Q_OFF 1024 -#define SMEM_SMEM_Q_STAGE_BYTES 2048 -#define SMEM_SMEM_Q_STRIDE 2048 -#define SMEM_SMEM_KV_OFF 9216 -#define SMEM_SMEM_KV_STAGE_BYTES 32768 -#define SMEM_SMEM_KV_STRIDE 32768 -#define SMEM_SMEM_V_OFF 9216 -#define SMEM_SMEM_V_STAGE_BYTES 32768 -#define SMEM_SMEM_V_STRIDE 32768 -#define SMEM_SMEM_P_OFF 205824 -#define SMEM_SMEM_P_STAGE_BYTES 2048 -#define SMEM_SMEM_P_STRIDE 2048 -#define SMEM_SMEM_PAGE_OFFSETS_OFF 209920 -#define SMEM_SMEM_PAGE_OFFSETS_STAGE_BYTES 1024 -#define SMEM_SMEM_PAGE_OFFSETS_STRIDE 1024 -#define SMEM_SMEM_O_OFF 216064 -#define SMEM_SMEM_O_STAGE_BYTES 2048 -#define SMEM_SMEM_O_STRIDE 2048 -#define SMEM_SMEM_SOFTMAX_OFF 218112 -#define SMEM_SMEM_SOFTMAX_STAGE_BYTES 128 -#define SMEM_SMEM_SOFTMAX_STRIDE 128 -#define SMEM_SMEM_SOFTMAX_U32_OFF 218112 -#define SMEM_SMEM_SOFTMAX_U32_STAGE_BYTES 128 -#define SMEM_SMEM_SOFTMAX_U32_STRIDE 128 -#define SMEM_SMEM_SOFTMAX_UNUSED_OFF 218240 -#define SMEM_SMEM_SOFTMAX_UNUSED_STAGE_BYTES 128 -#define SMEM_SMEM_SOFTMAX_UNUSED_STRIDE 128 -#define SMEM_SMEM_CORR_OFF 218368 -#define SMEM_SMEM_CORR_STAGE_BYTES 128 -#define SMEM_SMEM_CORR_STRIDE 128 -#define SMEM_SMEM_STATS_OFF 218496 -#define SMEM_SMEM_STATS_STAGE_BYTES 96 -#define SMEM_SMEM_STATS_STRIDE 96 -#define SMEM_WORK_RESPONSE_OFF 218688 -#define SMEM_WORK_RESPONSE_STAGE_BYTES 16 -#define SMEM_WORK_RESPONSE_STRIDE 16 -#define SMEM_TOTAL 218752 -#define THREADS 640 - -#include - -__device__ __forceinline__ uint32_t elect_sync() { - uint32_t pred = 0; - asm volatile( - "{\n\t" - ".reg .pred %%px;\n\t" - "elect.sync _|%%px, %1;\n\t" - "@%%px mov.s32 %0, 1;\n\t" - "}\n" - : "+r"(pred) - : "r"(0xFFFFFFFF)); - return pred; -} - - -__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { - asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" - :: "r"(mbar_addr), "r"(count) : "memory"); -} - -__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { - asm volatile("mbarrier.init.b64 [%0], %1;" - :: "l"(mbar_addr), "r"(count) : "memory"); -} - - -__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - - -// CTA-local pipelines have short, resident producer/consumer edges. Omitting -// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the -// explicit loop still makes this helper blocking until acquire succeeds. -__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE;\n\t" - "bra.uni LAB_WAIT;\n\t" - "DONE:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Source-faithful relaxed CTA wait used only by a typed protocol that does -// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. -__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, 10000000;\n\t" - "@P1 bra.uni DONE_RELAXED;\n\t" - "bra.uni LAB_WAIT_RELAXED;\n\t" - "DONE_RELAXED:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Exact source ports may request the PTX suspendTimeHint operand explicitly. -// The hint is expressed in nanoseconds and is kept separate from the canonical -// no-hint CTA helper so unrelated schedules retain their existing retry path. -__device__ __forceinline__ void mbarrier_wait_suspend( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_SUSPEND:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_SUSPEND;\n\t" - "bra.uni LAB_WAIT_SUSPEND;\n\t" - "DONE_SUSPEND:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE_CLUSTER;\n\t" - "bra.uni LAB_WAIT_CLUSTER;\n\t" - "DONE_CLUSTER:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - ".reg .u32 WAIT_ADDR;\n\t" - "mov.u32 WAIT_ADDR, %0;\n\t" - "LAB_WAIT_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [WAIT_ADDR], %1, %2;\n\t" - "@P1 bra.uni DONE_HINT;\n\t" - "bra.uni LAB_WAIT_HINT;\n\t" - "DONE_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -// Exact unqualified CTA wait used by source schedules whose PTX intentionally -// omits the acquire qualifier while retaining a typed suspendTimeHint operand. -__device__ __forceinline__ void mbarrier_wait_relaxed_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED_HINT:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra DONE_RELAXED_HINT;\n\t" - "bra LAB_WAIT_RELAXED_HINT;\n\t" - "DONE_RELAXED_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); -} - -__device__ __forceinline__ void mbarrier_wait_cluster_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" - "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" - "DONE_CLUSTER_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_suspend( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait_cluster(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); - } -} - - -__device__ __forceinline__ void tcgen05_mma_f16( - int taddr, uint64_t a_desc, uint64_t b_desc, - uint32_t i_desc, int enable_input_d) { - asm volatile( - "{\n\t" - ".reg .pred p;\n\t" - "setp.ne.b32 p, %4, 0;\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" - "}\n" - :: "r"(taddr), "l"(a_desc), "l"(b_desc), - "r"(i_desc), "r"(enable_input_d) - : "memory"); -} - - -__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { - return (x & 0x3FFFFULL) >> 4ULL; -} - - -union MmaSmemDesc { - uint64_t u64; - uint32_t u32[2]; -}; - -__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { - MmaSmemDesc tmp; - tmp.u64 = smem_desc; - tmp.u32[0] += offset; - smem_desc = tmp.u64; -} - - -__device__ __forceinline__ void elect_commit(int mbar_addr) { - asm volatile( - "{\n\t" - ".reg .pred leader;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" - ".shared::cluster.b64 [%0];\n\t" - "}\n" - :: "r"(mbar_addr)); -} - - -__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { - asm volatile( - "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" - :: "r"(mbar_addr), "r"(bytes) : "memory"); -} - - -__device__ __forceinline__ void tmem_ld_x4(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x4.b32" - " {%0, %1, %2, %3}, [%4];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ void tmem_st_x4_f32(int tmem_addr, const float* src) { - asm volatile( - "tcgen05.st.sync.aligned.32x32b.x4.b32" - " [%0], {%1, %2, %3, %4};" - :: "r"(tmem_addr), - "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3])); -} - - -__device__ __forceinline__ float approx_exp2(float x) { - float y; - asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float approx_rcp(float x) { - float y; - asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float max_noftz(float a, float b) { - float c; - asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); - return c; -} - - -__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { - asm("mul.rn.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { - asm("mul.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { - asm("add.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { - asm("add.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { - asm("sub.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { - asm("sub.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("add.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { - float2 r; - asm("sub.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("sub.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ void fma_scale_x32( - float* sv, const float2* scale2, const float2* neg_max2) -{ - float2* sv_2 = reinterpret_cast(sv); - #pragma unroll - for (int j = 0; j < 16; j++) - fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); -} - -__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("mul.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - - -// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) - -__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("add.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("add.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("add.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("add.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("add.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("add.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("add.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - - -__device__ __forceinline__ void fence_async_shared() { - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); -} - - -__device__ __forceinline__ uint64_t make_smem_desc(int addr) { - const int SBO = 1024; - return desc_encode(addr) - | (desc_encode(SBO) << 32ULL) - | (1ULL << 46ULL) - | (2ULL << 61ULL); -} - - -__device__ __forceinline__ void tma_2d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3}], [%4];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_4d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.4d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5}], [%6];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr) { - // Canonical .shared::cta form for non-multicast gather4, matching - // trtllm-gen / cuda_ptx and the PTX ISA qualifier order - // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, - // .shared::cluster is reserved for the multicast variant (ctaMask). - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { - asm volatile( - "tcgen05.commit.cta_group::1.mbarrier::arrive::one" - ".shared::cluster.b64 [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tmem_ld_x4_wait(float* dst, int addr) { - tmem_ld_x4(dst, addr); - asm volatile("tcgen05.wait::ld.sync.aligned;"); -} - - -__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { - uint32_t result; - asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" - : "=r"(result) : "r"(val)); - return result; -} - -extern "C" { - -__global__ __launch_bounds__(640, 1) void -kernel_cake_dsv4_d23652aae3370c30d978(CakeTensorMap const* tmap_q, CakeTensorMap const* tmap_swa_kv, CakeTensorMap const* tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, int* __restrict__ seq_lens, int* __restrict__ cum_seq_lens_q, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int max_q_len, int batch_size, int has_sinks) -{ - const int tid = threadIdx.x; - const int warp = make_warp_uniform(tid / 32); - const int lane = tid % 32; - - extern __shared__ __align__(1024) char smem_raw[]; - int smem; - smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); - - const int mbar_base = smem; - #define q_full_addr (mbar_base + 0) - #define q_empty_addr (mbar_base + 8) - #define kv_full_addr (mbar_base + 16) - #define kv_empty_addr (mbar_base + 64) - #define page_full_addr (mbar_base + 112) - #define page_empty_addr (mbar_base + 160) - #define score_full_addr (mbar_base + 208) - #define score_empty_addr (mbar_base + 224) - #define stats_full_addr (mbar_base + 240) - #define stats_empty_addr (mbar_base + 256) - #define p_full_addr (mbar_base + 272) - #define p_empty_addr (mbar_base + 288) - #define o_full_addr (mbar_base + 304) - #define o_empty_addr (mbar_base + 312) - #define work_full_addr (mbar_base + 320) - #define work_empty_addr (mbar_base + 336) - #define throttle_full_addr (mbar_base + 352) - #define throttle_empty_addr (mbar_base + 368) - - const int bid = blockIdx.x; - const int num_bids = gridDim.x; - - const int cta_rank = 0; - if (tid == 0) { - asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_q)) : "memory"); - asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_swa_kv)) : "memory"); - asm volatile("fence.proxy.tensormap::generic.acquire.sys [%0], 128;" :: "l"((uint64_t)(tmap_compressed_kv)) : "memory"); - } - __syncthreads(); - - - // Kernel setup ops - __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); - const int smem_q_addr = smem + 1024; - __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 9216); - const int smem_kv_addr = smem + 9216; - __nv_bfloat16* smem_v = reinterpret_cast<__nv_bfloat16*>(smem_raw + 9216); - const int smem_v_addr = smem + 9216; - __nv_bfloat16* smem_p = reinterpret_cast<__nv_bfloat16*>(smem_raw + 205824); - const int smem_p_addr = smem + 205824; - int* smem_page_offsets = reinterpret_cast(smem_raw + 209920); - const int smem_page_offsets_addr = smem + 209920; - __nv_bfloat16* smem_o = reinterpret_cast<__nv_bfloat16*>(smem_raw + 216064); - const int smem_o_addr = smem + 216064; - float* smem_softmax = reinterpret_cast(smem_raw + 218112); - const int smem_softmax_addr = smem + 218112; - unsigned int* smem_softmax_u32 = reinterpret_cast(smem_raw + 218112); - const int smem_softmax_u32_addr = smem + 218112; - float* smem_softmax_unused = reinterpret_cast(smem_raw + 218240); - const int smem_softmax_unused_addr = smem + 218240; - float* smem_corr = reinterpret_cast(smem_raw + 218368); - const int smem_corr_addr = smem + 218368; - float* smem_stats = reinterpret_cast(smem_raw + 218496); - const int smem_stats_addr = smem + 218496; - unsigned int* work_response = reinterpret_cast(smem_raw + 218688); - const int work_response_addr = smem + 218688; - asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)(tmap_q)) : "memory"); - asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)(tmap_swa_kv)) : "memory"); - asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)(tmap_compressed_kv)) : "memory"); - - // Mbarrier init (18 pipeline groups, 0 ordered-sequence groups, 48 barriers) - // Mbarriers at smem_raw[0..384) - - if (warp == 0) { - uint32_t leader = elect_sync(); - if (leader) { - // --- pipeline 'q_pipe' --- - // q_full: 1 barriers, init_count=1 - mbarrier_init(smem + 0, 1); - // q_empty: 1 barriers, init_count=1 - mbarrier_init(smem + 8, 1); - // --- pipeline 'kv_pipe' --- - // kv_full: 6 barriers, init_count=1 - mbarrier_init(smem + 16, 1); - mbarrier_init(smem + 24, 1); - mbarrier_init(smem + 32, 1); - mbarrier_init(smem + 40, 1); - mbarrier_init(smem + 48, 1); - mbarrier_init(smem + 56, 1); - // kv_empty: 6 barriers, init_count=1 - mbarrier_init(smem + 64, 1); - mbarrier_init(smem + 72, 1); - mbarrier_init(smem + 80, 1); - mbarrier_init(smem + 88, 1); - mbarrier_init(smem + 96, 1); - mbarrier_init(smem + 104, 1); - // --- pipeline 'page_pipe' --- - // page_full: 6 barriers, init_count=32 - mbarrier_init(smem + 112, 32); - mbarrier_init(smem + 120, 32); - mbarrier_init(smem + 128, 32); - mbarrier_init(smem + 136, 32); - mbarrier_init(smem + 144, 32); - mbarrier_init(smem + 152, 32); - // page_empty: 6 barriers, init_count=256 - mbarrier_init(smem + 160, 256); - mbarrier_init(smem + 168, 256); - mbarrier_init(smem + 176, 256); - mbarrier_init(smem + 184, 256); - mbarrier_init(smem + 192, 256); - mbarrier_init(smem + 200, 256); - // --- pipeline 'score_pipe' --- - // score_full: 2 barriers, init_count=1 - mbarrier_init(smem + 208, 1); - mbarrier_init(smem + 216, 1); - // score_empty: 2 barriers, init_count=128 - mbarrier_init(smem + 224, 128); - mbarrier_init(smem + 232, 128); - // --- pipeline 'stats_pipe' --- - // stats_full: 2 barriers, init_count=128 - mbarrier_init(smem + 240, 128); - mbarrier_init(smem + 248, 128); - // stats_empty: 2 barriers, init_count=128 - mbarrier_init(smem + 256, 128); - mbarrier_init(smem + 264, 128); - // --- pipeline 'p_pipe' --- - // p_full: 2 barriers, init_count=128 - mbarrier_init(smem + 272, 128); - mbarrier_init(smem + 280, 128); - // p_empty: 2 barriers, init_count=1 - mbarrier_init(smem + 288, 1); - mbarrier_init(smem + 296, 1); - // --- pipeline 'o_pipe' --- - // o_full: 1 barriers, init_count=1 - mbarrier_init(smem + 304, 1); - // o_empty: 1 barriers, init_count=128 - mbarrier_init(smem + 312, 128); - // --- pipeline 'work_pipe' --- - // work_full: 2 barriers, init_count=1 - mbarrier_init(smem + 320, 1); - mbarrier_init(smem + 328, 1); - // work_empty: 2 barriers, init_count=640 - mbarrier_init(smem + 336, 640); - mbarrier_init(smem + 344, 640); - // --- pipeline 'throttle_pipe' --- - // throttle_full: 2 barriers, init_count=256 - mbarrier_init(smem + 352, 256); - mbarrier_init(smem + 360, 256); - // throttle_empty: 2 barriers, init_count=32 - mbarrier_init(smem + 368, 32); - mbarrier_init(smem + 376, 32); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - - __syncwarp(); - - // TMEM alloc (512 columns, 512 used) - volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 384); - if (warp == 0) { - int _tmem_hold = smem + 384; - asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); - __syncwarp(); - asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); - } - - __syncthreads(); - asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); - - const int taddr = tmem_addr_storage[0]; - - // Kernel post-init ops - const int tmem_source_exact_scratch = taddr; - - // ---- Ordered hardware-WG register redistribution ---- - // Dec phase frees registers before any WG attempts inc. - if (warp >= 8 && warp <= 11) { - asm volatile("setmaxnreg.dec.sync.aligned.u32 48;"); - } - - // ---- Role: softmax ---- - if (warp <= 3) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 152;"); - { // softmax_main - unsigned int work_stage = 0; - unsigned int score_stage = 0; - unsigned int score_phase = 0; - unsigned int stats_stage = 0; - unsigned int stats_phase = 1; - unsigned int p_stage = 0; - unsigned int p_phase = 1; - unsigned int work_x = blockIdx.x; - unsigned int work_y = blockIdx.y; - unsigned int work_z = blockIdx.z; - const int warp_in_wg = warp; - const int wg_tid = warp_in_wg * 32 + lane; - const int col_pair_base = wg_tid % 4 * 2; - const int ldtm_row_base = warp_in_wg * 32 + lane / 4; - float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; - unsigned int _phase_work_full = 0; - #pragma unroll 1 - for (unsigned int _work = 0; _work < 120; _work++) { - int q_begin = cum_seq_lens_q[work_z]; - int q_end = cum_seq_lens_q[work_z + 1]; - int q_len = q_end - q_begin; - int query_idx = (unsigned int)q_begin + work_x; - int active_topk = 0; - if (work_x < (unsigned int)q_len) { - int _max_4 = ((sparse_topk_lens[query_idx] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx] + sparse_topk_lens_offset) : (0)); - int _min_3 = ((_max_4) < (sparse_topk) ? (_max_4) : (sparse_topk)); - active_topk = _min_3; - } - int num_steps = (active_topk + 128 - 1) / 128; - float row_max_pair[2]; - float row_sum_pair[2]; - row_max_pair[0] = -CAKE_INF; - row_max_pair[1] = -CAKE_INF; - row_sum_pair[0] = 0.0f; - row_sum_pair[1] = 0.0f; - asm volatile("barrier.sync 1, 128;" ::: "memory"); - uint32_t _amf_u_0 = __float_as_uint(-3.4028235e+38f); - uint32_t _amf_mask_0 = -int32_t(_amf_u_0 >> 31) | 0x80000000u; - unsigned int _amf_enc_0 = _amf_u_0 ^ _amf_mask_0; - if (wg_tid < 8) { - smem_softmax_u32[wg_tid] = _amf_enc_0; - } - asm volatile("barrier.sync 1, 128;" ::: "memory"); - #pragma unroll 1 - for (int tile = 0; tile < num_steps; tile++) { - mbarrier_wait(score_full_addr + (score_stage) * 8, score_phase); - int score_col = ((score_stage == 0) ? 0 : 8); - float scores[8]; - float score_lo[4]; - float score_hi[4]; - int tile_end = (tile + 1) * 128; - int is_full_tile = ((active_topk % 128 == 0) ? 1 : 0); - is_full_tile = is_full_tile | ((tile_end < active_topk) ? 1 : 0); - if (is_full_tile != 0) { - asm volatile( - "tcgen05.ld.sync.aligned.16x256b.x1.b32" - " {%0, %1, %2, %3}, [%4];" - : "=r"(*reinterpret_cast(&score_lo[0])), "=r"(*reinterpret_cast(&score_lo[1])), "=r"(*reinterpret_cast(&score_lo[2])), "=r"(*reinterpret_cast(&score_lo[3])) - : "r"(taddr + (unsigned int)score_col)); - asm volatile( - "tcgen05.ld.sync.aligned.16x256b.x1.b32" - " {%0, %1, %2, %3}, [%4];" - : "=r"(*reinterpret_cast(&score_hi[0])), "=r"(*reinterpret_cast(&score_hi[1])), "=r"(*reinterpret_cast(&score_hi[2])), "=r"(*reinterpret_cast(&score_hi[3])) - : "r"(taddr + (unsigned int)score_col + 1048576)); - #pragma unroll - for (int c = 0; c < 4; c++) { - scores[c] = score_lo[c]; - scores[c + 4] = score_hi[c]; - } - } else { - asm volatile("" ::: "memory"); - asm volatile( - "tcgen05.ld.sync.aligned.16x256b.x1.b32" - " {%0, %1, %2, %3}, [%4];" - : "=r"(*reinterpret_cast(&score_lo[0])), "=r"(*reinterpret_cast(&score_lo[1])), "=r"(*reinterpret_cast(&score_lo[2])), "=r"(*reinterpret_cast(&score_lo[3])) - : "r"(taddr + (unsigned int)score_col)); - asm volatile( - "tcgen05.ld.sync.aligned.16x256b.x1.b32" - " {%0, %1, %2, %3}, [%4];" - : "=r"(*reinterpret_cast(&score_hi[0])), "=r"(*reinterpret_cast(&score_hi[1])), "=r"(*reinterpret_cast(&score_hi[2])), "=r"(*reinterpret_cast(&score_hi[3])) - : "r"(taddr + (unsigned int)score_col + 1048576)); - #pragma unroll - for (int c_1 = 0; c_1 < 4; c_1++) { - scores[c_1] = score_lo[c_1]; - scores[c_1 + 4] = score_hi[c_1]; - } - if (active_topk <= tile * 128 + ldtm_row_base) { - scores[0] = -CAKE_INF; - scores[1] = -CAKE_INF; - } - if (active_topk <= tile * 128 + ldtm_row_base + 8) { - scores[2] = -CAKE_INF; - scores[3] = -CAKE_INF; - } - if (active_topk <= tile * 128 + ldtm_row_base + 16) { - scores[4] = -CAKE_INF; - scores[5] = -CAKE_INF; - } - if (active_topk <= tile * 128 + ldtm_row_base + 24) { - scores[6] = -CAKE_INF; - scores[7] = -CAKE_INF; - } - asm volatile("" ::: "memory"); - } - float pair_max[2]; - float _max_5 = max_noftz(scores[0], scores[2]); - float _max_6 = max_noftz(scores[4], scores[6]); - float _max_7 = max_noftz(_max_5, _max_6); - pair_max[0] = _max_7; - float _max_8 = max_noftz(scores[1], scores[3]); - float _max_9 = max_noftz(scores[5], scores[7]); - float _max_10 = max_noftz(_max_8, _max_9); - pair_max[1] = _max_10; - #pragma unroll - for (int c_2 = 0; c_2 < 2; c_2++) { - float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, pair_max[c_2], 16); - float _max_11 = max_noftz(pair_max[c_2], _shfl_xor_0); - pair_max[c_2] = _max_11; - float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, pair_max[c_2], 8); - float _max_12 = max_noftz(pair_max[c_2], _shfl_xor_1); - pair_max[c_2] = _max_12; - } - float old_max_pair[2]; - float new_max_pair[2]; - #pragma unroll - for (int c_3 = 0; c_3 < 2; c_3++) { - old_max_pair[c_3] = row_max_pair[c_3]; - float _max_13 = max_noftz(row_max_pair[c_3], pair_max[c_3]); - new_max_pair[c_3] = _max_13; - } - if (lane < 8) { - uint32_t _amf_u_1 = __float_as_uint(new_max_pair[0]); - uint32_t _amf_mask_1 = -int32_t(_amf_u_1 >> 31) | 0x80000000u; - unsigned int _amf_enc_1 = _amf_u_1 ^ _amf_mask_1; - uint32_t _amf_u_2 = __float_as_uint(new_max_pair[1]); - uint32_t _amf_mask_2 = -int32_t(_amf_u_2 >> 31) | 0x80000000u; - unsigned int _amf_enc_2 = _amf_u_2 ^ _amf_mask_2; - atomicMax(&smem_softmax_u32[col_pair_base], _amf_enc_1); - atomicMax(&smem_softmax_u32[col_pair_base + 1], _amf_enc_2); - } - asm volatile("barrier.sync 1, 128;" ::: "memory"); - uint32_t _amf_u_3 = smem_softmax_u32[col_pair_base]; - uint32_t _amf_mask_3 = ((_amf_u_3 >> 31) - 1u) | 0x80000000u; - float _amf_dec_0 = __uint_as_float(_amf_u_3 ^ _amf_mask_3); - new_max_pair[0] = _amf_dec_0; - uint32_t _amf_u_4 = smem_softmax_u32[col_pair_base + 1]; - uint32_t _amf_mask_4 = ((_amf_u_4 >> 31) - 1u) | 0x80000000u; - float _amf_dec_1 = __uint_as_float(_amf_u_4 ^ _amf_mask_4); - new_max_pair[1] = _amf_dec_1; - mbarrier_wait(stats_empty_addr + (stats_stage) * 8, stats_phase); - float stats_pair[4]; - stats_pair[0] = old_max_pair[0]; - stats_pair[1] = old_max_pair[1]; - stats_pair[2] = new_max_pair[0]; - stats_pair[3] = new_max_pair[1]; - int stats_col = ((stats_stage == 0) ? 16 : 48); - tmem_st_x4_f32(taddr + (unsigned int)stats_col, stats_pair); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(stats_full_addr + (stats_stage) * 8); - stats_stage += 1; - if (stats_stage == 2) { stats_stage = 0; stats_phase ^= 1; } - float acc_scale_pair[2]; - #pragma unroll - for (int c_4 = 0; c_4 < 2; c_4++) { - float delta = softmax_scale_log2 * (old_max_pair[c_4] - new_max_pair[c_4]); - float _exp2_0 = approx_exp2(delta); - acc_scale_pair[c_4] = ((old_max_pair[c_4] > -CAKE_INF) ? _exp2_0 : 1.0f); - row_max_pair[c_4] = new_max_pair[c_4]; - } - asm volatile("" ::: "memory"); - float exp_values[8]; - #pragma unroll - for (int c_5 = 0; c_5 < 8; c_5++) { - float safe_max = ((new_max_pair[c_5 % 2] == -CAKE_INF) ? 0.0f : new_max_pair[c_5 % 2]); - float _exp2_1 = approx_exp2(scores[c_5] * softmax_scale_log2 - safe_max * softmax_scale_log2); - exp_values[c_5] = _exp2_1; - } - mbarrier_wait(p_empty_addr + (p_stage) * 8, p_phase); - unsigned int regs_p[4]; - #pragma unroll - for (int _lp = 0; _lp < 4; _lp++) { - __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(exp_values[_lp*2 + 0], exp_values[_lp*2+1 + 0])); - regs_p[_lp] = *(uint32_t*)&_bf2; - } - int slice_idx = warp_in_wg / 2; - int warp_idx_in_slice = warp_in_wg % 2; - int mtx_idx = lane / 8; - int thr_row_idx = lane % 8; - int seg_col_idx = warp_idx_in_slice * 4 + mtx_idx ^ thr_row_idx; - int stsm_offset = slice_idx * 8 * 128 + thr_row_idx * 128 + seg_col_idx * 16; - asm volatile("stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 [%0], {%1, %2, %3, %4};\n" - :: "l"(__cvta_generic_to_shared(reinterpret_cast(smem_p) + (p_stage * 2048 + (unsigned int)stsm_offset))), "r"(*reinterpret_cast(®s_p[0])), "r"(*reinterpret_cast(®s_p[1])), "r"(*reinterpret_cast(®s_p[2])), "r"(*reinterpret_cast(®s_p[3])) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(p_full_addr + (p_stage) * 8); - p_stage += 1; - if (p_stage == 2) { p_stage = 0; p_phase ^= 1; } - float2 _f2_0 = make_float2(exp_values[0], exp_values[1]); - float2 _f2_1 = make_float2(exp_values[2], exp_values[3]); - float2 _f2_2 = make_float2(exp_values[4], exp_values[5]); - float2 _f2_3 = make_float2(exp_values[6], exp_values[7]); - float2 _f2_4 = make_float2(row_sum_pair[0], row_sum_pair[1]); - float2 _f2_5 = make_float2(acc_scale_pair[0], acc_scale_pair[1]); - float2 sum_p0_pair = fma_f32x2_rn_ftz(_f2_4, _f2_5, _f2_0); - float2 sum_p1_pair = add_f32x2(sum_p0_pair, _f2_1); - float2 sum_p2_pair = add_f32x2(sum_p1_pair, _f2_2); - float2 next_sum_pair = add_f32x2(sum_p2_pair, _f2_3); - row_sum_pair[0] = next_sum_pair.x; - row_sum_pair[1] = next_sum_pair.y; - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - mbarrier_arrive(score_empty_addr + (score_stage) * 8); - score_stage += 1; - if (score_stage == 2) { score_stage = 0; score_phase ^= 1; } - } - if (num_steps > 0) { - mbarrier_wait(stats_empty_addr + (stats_stage) * 8, stats_phase); - float final_stats_pair[4]; - final_stats_pair[0] = row_sum_pair[0]; - final_stats_pair[1] = row_sum_pair[1]; - final_stats_pair[2] = row_max_pair[0]; - final_stats_pair[3] = row_max_pair[1]; - int final_stats_col = ((stats_stage == 0) ? 16 : 48); - tmem_st_x4_f32(taddr + (unsigned int)final_stats_col, final_stats_pair); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(stats_full_addr + (stats_stage) * 8); - stats_stage += 1; - if (stats_stage == 2) { stats_stage = 0; stats_phase ^= 1; } - } - mbarrier_wait(work_full_addr + (work_stage) * 8, _phase_work_full); - uint32_t _clc_valid_3 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_3) - : "r"(work_response_addr + work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_9 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_9) - : "r"(work_response_addr + work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_10 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_10) - : "r"(work_response_addr + work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_11 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_11) - : "r"(work_response_addr + work_stage * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(work_empty_addr + work_stage * 8), "r"(0) : "memory"); - work_stage += 1; - if (work_stage == 2) { work_stage = 0; _phase_work_full ^= 1; } - if (_clc_valid_3 == 0) { - break; - } - work_x = _clc_ctaid_9; - work_y = _clc_ctaid_10; - work_z = _clc_ctaid_11; - } - } - } - // ---- Role: correction ---- - if (warp >= 4 && warp <= 7) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 136;"); - { // correction_main - unsigned int work_stage_1 = 0; - unsigned int stats_stage_1 = 0; - unsigned int stats_phase_1 = 0; - unsigned int work_x_1 = blockIdx.x; - unsigned int work_y_1 = blockIdx.y; - unsigned int work_z_1 = blockIdx.z; - const int corr_rank = warp - 4; - const int corr_tid = corr_rank * 32 + lane; - const int col_pair_base_1 = corr_tid % 4 * 2; - float softmax_scale_log2_1 = bmm1_scale[0] * 1.4426950408889634f; - float output_scale = bmm2_scale[0]; - unsigned int _phase_o_full_0 = 0; - unsigned int _phase_work_full_1 = 0; - #pragma unroll 1 - for (unsigned int _work_1 = 0; _work_1 < 120; _work_1++) { - int q_begin_1 = cum_seq_lens_q[work_z_1]; - int q_end_1 = cum_seq_lens_q[work_z_1 + 1]; - int q_len_1 = q_end_1 - q_begin_1; - int query_idx_1 = (unsigned int)q_begin_1 + work_x_1; - int active_topk_1 = 0; - if (work_x_1 < (unsigned int)q_len_1) { - int _max_14 = ((sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) : (0)); - int _min_4 = ((_max_14) < (sparse_topk) ? (_max_14) : (sparse_topk)); - active_topk_1 = _min_4; - } - int num_steps_1 = (active_topk_1 + 128 - 1) / 128; - int head_group = work_y_1 / 4; - int value_quarter = work_y_1 % 4; - if (num_steps_1 > 0) { - mbarrier_wait(stats_full_addr + (stats_stage_1) * 8, stats_phase_1); - mbarrier_arrive(stats_empty_addr + (stats_stage_1) * 8); - stats_stage_1 += 1; - if (stats_stage_1 == 2) { stats_stage_1 = 0; stats_phase_1 ^= 1; } - } - #pragma unroll 1 - for (int tile_1 = 1; tile_1 < num_steps_1; tile_1++) { - mbarrier_wait(stats_full_addr + (stats_stage_1) * 8, stats_phase_1); - int stats_col_1 = ((stats_stage_1 == 0) ? 16 : 48); - float _tmem_load_0[4]; - tmem_ld_x4(&_tmem_load_0[0], taddr + (unsigned int)stats_col_1); - asm volatile("tcgen05.wait::ld.sync.aligned;"); - float scale_pair[2]; - #pragma unroll - for (int c_6 = 0; c_6 < 2; c_6++) { - float max_diff = _tmem_load_0[c_6] - _tmem_load_0[c_6 + 2]; - float _exp2_2 = approx_exp2(softmax_scale_log2_1 * max_diff); - scale_pair[c_6] = ((max_diff != 0.0f) ? _exp2_2 : 1.0f); - } - mbarrier_arrive(stats_empty_addr + (stats_stage_1) * 8); - stats_stage_1 += 1; - if (stats_stage_1 == 2) { stats_stage_1 = 0; stats_phase_1 ^= 1; } - mbarrier_wait(o_full_addr, _phase_o_full_0); - _phase_o_full_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - int needs_scale = ((scale_pair[0] != 1.0f) ? 1 : 0); - needs_scale = needs_scale | ((scale_pair[1] != 1.0f) ? 1 : 0); - int _vote_0 = __any_sync(0xFFFFFFFF, needs_scale != 0); - needs_scale = _vote_0; - if (needs_scale != 0) { - float values_lo[4]; - float values_hi[4]; - asm volatile( - "tcgen05.ld.sync.aligned.16x256b.x1.b32" - " {%0, %1, %2, %3}, [%4];" - : "=r"(*reinterpret_cast(&values_lo[0])), "=r"(*reinterpret_cast(&values_lo[1])), "=r"(*reinterpret_cast(&values_lo[2])), "=r"(*reinterpret_cast(&values_lo[3])) - : "r"(taddr + 80)); - asm volatile( - "tcgen05.ld.sync.aligned.16x256b.x1.b32" - " {%0, %1, %2, %3}, [%4];" - : "=r"(*reinterpret_cast(&values_hi[0])), "=r"(*reinterpret_cast(&values_hi[1])), "=r"(*reinterpret_cast(&values_hi[2])), "=r"(*reinterpret_cast(&values_hi[3])) - : "r"(taddr + 80 + 1048576)); - #pragma unroll - for (int pair = 0; pair < 4; pair++) { - asm volatile("" ::: "memory"); - values_lo[pair] = values_lo[pair] * scale_pair[pair % 2]; - values_hi[pair] = values_hi[pair] * scale_pair[pair % 2]; - } - asm volatile( - "tcgen05.st.sync.aligned.16x256b.x1.b32" - " [%0], {%1, %2, %3, %4};" - :: "r"(taddr + 80), "r"(*reinterpret_cast(&values_lo[0])), "r"(*reinterpret_cast(&values_lo[1])), "r"(*reinterpret_cast(&values_lo[2])), "r"(*reinterpret_cast(&values_lo[3]))); - asm volatile( - "tcgen05.st.sync.aligned.16x256b.x1.b32" - " [%0], {%1, %2, %3, %4};" - :: "r"(taddr + 80 + 1048576), "r"(*reinterpret_cast(&values_hi[0])), "r"(*reinterpret_cast(&values_hi[1])), "r"(*reinterpret_cast(&values_hi[2])), "r"(*reinterpret_cast(&values_hi[3]))); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - } - mbarrier_arrive(o_empty_addr); - } - if (num_steps_1 > 0) { - mbarrier_wait(stats_full_addr + (stats_stage_1) * 8, stats_phase_1); - int final_stats_col_1 = ((stats_stage_1 == 0) ? 16 : 48); - float _tmem_load_1[4]; - tmem_ld_x4(&_tmem_load_1[0], taddr + (unsigned int)final_stats_col_1); - asm volatile("tcgen05.wait::ld.sync.aligned;"); - mbarrier_arrive(stats_empty_addr + (stats_stage_1) * 8); - stats_stage_1 += 1; - if (stats_stage_1 == 2) { stats_stage_1 = 0; stats_phase_1 ^= 1; } - float pair_sum[2]; - pair_sum[0] = _tmem_load_1[0]; - pair_sum[1] = _tmem_load_1[1]; - float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, pair_sum[0], 16); - pair_sum[0] = pair_sum[0] + _shfl_xor_2; - float _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, pair_sum[1], 16); - pair_sum[1] = pair_sum[1] + _shfl_xor_3; - float _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, pair_sum[0], 8); - pair_sum[0] = pair_sum[0] + _shfl_xor_4; - float _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, pair_sum[1], 8); - pair_sum[1] = pair_sum[1] + _shfl_xor_5; - float _shfl_xor_6 = __shfl_xor_sync(0xFFFFFFFF, pair_sum[0], 4); - pair_sum[0] = pair_sum[0] + _shfl_xor_6; - float _shfl_xor_7 = __shfl_xor_sync(0xFFFFFFFF, pair_sum[1], 4); - pair_sum[1] = pair_sum[1] + _shfl_xor_7; - if (lane < 4) { - smem_corr[corr_rank * 8 + col_pair_base_1] = pair_sum[0]; - smem_corr[corr_rank * 8 + col_pair_base_1 + 1] = pair_sum[1]; - } - asm volatile("barrier.sync 2, 128;" ::: "memory"); - float sum0 = smem_corr[col_pair_base_1] + smem_corr[8 + col_pair_base_1] + smem_corr[16 + col_pair_base_1] + smem_corr[24 + col_pair_base_1]; - float sum1 = smem_corr[col_pair_base_1 + 1] + smem_corr[8 + col_pair_base_1 + 1] + smem_corr[16 + col_pair_base_1 + 1] + smem_corr[24 + col_pair_base_1 + 1]; - if (has_sinks != 0) { - float _exp2_3 = approx_exp2(sinks[head_group * 8 + col_pair_base_1] * 1.4426950408889634f - _tmem_load_1[2] * softmax_scale_log2_1); - sum0 = sum0 + _exp2_3; - float _exp2_4 = approx_exp2(sinks[head_group * 8 + col_pair_base_1 + 1] * 1.4426950408889634f - _tmem_load_1[3] * softmax_scale_log2_1); - sum1 = sum1 + _exp2_4; - } - float _rcp_0 = approx_rcp(sum0); - float scale0 = ((sum0 > 0.0f) ? output_scale * _rcp_0 : 0.0f); - float _rcp_1 = approx_rcp(sum1); - float scale1 = ((sum1 > 0.0f) ? output_scale * _rcp_1 : 0.0f); - mbarrier_wait(o_full_addr, _phase_o_full_0); - _phase_o_full_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - float output_lo[4]; - float output_hi[4]; - asm volatile( - "tcgen05.ld.sync.aligned.16x256b.x1.b32" - " {%0, %1, %2, %3}, [%4];" - : "=r"(*reinterpret_cast(&output_lo[0])), "=r"(*reinterpret_cast(&output_lo[1])), "=r"(*reinterpret_cast(&output_lo[2])), "=r"(*reinterpret_cast(&output_lo[3])) - : "r"(taddr + 80)); - asm volatile( - "tcgen05.ld.sync.aligned.16x256b.x1.b32" - " {%0, %1, %2, %3}, [%4];" - : "=r"(*reinterpret_cast(&output_hi[0])), "=r"(*reinterpret_cast(&output_hi[1])), "=r"(*reinterpret_cast(&output_hi[2])), "=r"(*reinterpret_cast(&output_hi[3])) - : "r"(taddr + 80 + 1048576)); - unsigned int packed_o[4]; - #pragma unroll - for (int pair_1 = 0; pair_1 < 4; pair_1++) { - asm volatile("" ::: "memory"); - output_lo[pair_1] = output_lo[pair_1] * ((pair_1 % 2 == 0) ? scale0 : scale1); - output_hi[pair_1] = output_hi[pair_1] * ((pair_1 % 2 == 0) ? scale0 : scale1); - } - float output_values[8]; - #pragma unroll - for (int c_7 = 0; c_7 < 4; c_7++) { - output_values[c_7] = output_lo[c_7]; - output_values[c_7 + 4] = output_hi[c_7]; - } - #pragma unroll - for (int _lp = 0; _lp < 4; _lp++) { - __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(output_values[_lp*2 + 0], output_values[_lp*2+1 + 0])); - packed_o[_lp] = *(uint32_t*)&_bf2; - } - int slice_idx_1 = corr_rank / 2; - int warp_idx_in_slice_1 = corr_rank % 2; - int mtx_idx_1 = lane / 8; - int thr_row_idx_1 = lane % 8; - int seg_col_idx_1 = warp_idx_in_slice_1 * 4 + mtx_idx_1 ^ thr_row_idx_1; - int stsm_offset_1 = slice_idx_1 * 8 * 128 + thr_row_idx_1 * 128 + seg_col_idx_1 * 16; - asm volatile("stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 [%0], {%1, %2, %3, %4};\n" - :: "l"(__cvta_generic_to_shared(reinterpret_cast(smem_o) + stsm_offset_1)), "r"(*reinterpret_cast(&packed_o[0])), "r"(*reinterpret_cast(&packed_o[1])), "r"(*reinterpret_cast(&packed_o[2])), "r"(*reinterpret_cast(&packed_o[3])) - : "memory"); - asm volatile("barrier.sync 2, 128;" ::: "memory"); - int copy_base = corr_tid * 16; - int smem_row = copy_base / 128; - int copy_smem_offset = copy_base ^ smem_row % 8 * 16; - int copy_row = smem_row % 8; - int copy_col = (smem_row / 8 * 128 + copy_base % 128) / 2; - unsigned int copy_vec[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(©_vec[0])), "=r"(*reinterpret_cast(©_vec[(0) + 1])), "=r"(*reinterpret_cast(©_vec[(0) + 2])), "=r"(*reinterpret_cast(©_vec[(0) + 3])) - : "r"(smem_o_addr + (unsigned int)copy_smem_offset)); - int head_groups = gridDim.y / 4; - int output_base = query_idx_1 * head_groups * 8 * 512 + head_group * 8 * 512 + value_quarter * 128; - reinterpret_cast(O + (output_base + copy_row * 512 + copy_col))[0] = reinterpret_cast(copy_vec)[0]; - mbarrier_arrive(o_empty_addr); - } - mbarrier_wait(work_full_addr + (work_stage_1) * 8, _phase_work_full_1); - uint32_t _clc_valid_4 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_4) - : "r"(work_response_addr + work_stage_1 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_12 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_12) - : "r"(work_response_addr + work_stage_1 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_13 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_13) - : "r"(work_response_addr + work_stage_1 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_14 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_14) - : "r"(work_response_addr + work_stage_1 * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(work_empty_addr + work_stage_1 * 8), "r"(0) : "memory"); - work_stage_1 += 1; - if (work_stage_1 == 2) { work_stage_1 = 0; _phase_work_full_1 ^= 1; } - if (_clc_valid_4 == 0) { - break; - } - work_x_1 = _clc_ctaid_12; - work_y_1 = _clc_ctaid_13; - work_z_1 = _clc_ctaid_14; - } - asm volatile("barrier.sync 11, 128;" ::: "memory"); - if (corr_rank == 0) { - int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); - asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); - } - } - } - // ---- Role: mma ---- - if (warp == 8) { - { // mma_main - unsigned int work_stage_2 = 0; - unsigned int q_stage = 0; - unsigned int q_phase = 0; - unsigned int kv_stage = 0; - unsigned int kv_phase = 0; - unsigned int score_stage_1 = 0; - unsigned int score_phase_1 = 1; - unsigned int p_stage_1 = 0; - unsigned int p_phase_1 = 0; - unsigned int work_x_2 = blockIdx.x; - unsigned int work_y_2 = blockIdx.y; - unsigned int work_z_2 = blockIdx.z; - unsigned int _phase_o_empty_0 = 1; - unsigned int _phase_work_full_2 = 0; - #pragma unroll 1 - for (unsigned int _work_2 = 0; _work_2 < 120; _work_2++) { - int q_begin_2 = cum_seq_lens_q[work_z_2]; - int q_end_2 = cum_seq_lens_q[work_z_2 + 1]; - int q_len_2 = q_end_2 - q_begin_2; - int query_idx_2 = (unsigned int)q_begin_2 + work_x_2; - int active_topk_2 = 0; - if (work_x_2 < (unsigned int)q_len_2) { - int _max_3 = ((sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) : (0)); - int _min_2 = ((_max_3) < (sparse_topk) ? (_max_3) : (sparse_topk)); - active_topk_2 = _min_2; - } - int num_steps_2 = (active_topk_2 + 128 - 1) / 128; - int value_quarter_1 = work_y_2 % 4; - if (num_steps_2 > 0) { - mbarrier_wait(q_full_addr + (q_stage) * 8, q_phase); - int first_output = 1; - mbarrier_wait(score_empty_addr + (score_stage_1) * 8, score_phase_1); - int score_col_1 = ((score_stage_1 == 0) ? 0 : 8); - #pragma unroll - for (int hs = 0; hs < 3; hs++) { - int hs_sel = ((value_quarter_1 > hs) ? hs : hs + 1); - mbarrier_wait(kv_full_addr + (kv_stage) * 8, kv_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - int _mma_a_lo_0 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage) * 2048); - int _mma_b_lo_0 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (hs_sel) * 128); - { - uint64_t _mma_ss_a_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_0); - uint64_t _mma_ss_b_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_0); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134349968, ((hs == 0) ? 0 : 1)); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134349968, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134349968, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134349968, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 1018U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 58U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134349968, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134349968, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134349968, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 134349968, 1); - } - } - elect_commit(kv_empty_addr + (kv_stage) * 8); - kv_stage += 1; - if (kv_stage == 6) { kv_stage = 0; kv_phase ^= 1; } - } - mbarrier_wait(kv_full_addr + (kv_stage) * 8, kv_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - int _mma_a_lo_1 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage) * 2048); - int _mma_b_lo_1 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (value_quarter_1) * 128); - { - uint64_t _mma_ss_a_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_1); - uint64_t _mma_ss_b_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_1); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134349968, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134349968, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134349968, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134349968, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 1018U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 58U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134349968, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134349968, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134349968, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 134349968, 1); - } - } - kv_stage += 1; - if (kv_stage == 6) { kv_stage = 0; kv_phase ^= 1; } - elect_commit(score_full_addr + (score_stage_1) * 8); - score_stage_1 += 1; - if (score_stage_1 == 2) { score_stage_1 = 0; score_phase_1 ^= 1; } - #pragma unroll 1 - for (int _tile = 0; _tile < num_steps_2 - 1; _tile++) { - mbarrier_wait(score_empty_addr + (score_stage_1) * 8, score_phase_1); - int score_col_0 = ((score_stage_1 == 0) ? 0 : 8); - #pragma unroll - for (int hs_1 = 0; hs_1 < 3; hs_1++) { - int hs_sel_1 = ((value_quarter_1 > hs_1) ? hs_1 : hs_1 + 1); - mbarrier_wait(kv_full_addr + (kv_stage) * 8, kv_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - int _mma_a_lo_2 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage) * 2048); - int _mma_b_lo_2 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (hs_sel_1) * 128); - { - uint64_t _mma_ss_a_desc_2 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_2); - uint64_t _mma_ss_b_desc_2 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_2); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134349968, ((hs_1 == 0) ? 0 : 1)); - } - incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134349968, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134349968, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134349968, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_2, 1018U); - incr_smem_desc_lo(_mma_ss_b_desc_2, 58U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134349968, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134349968, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134349968, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 134349968, 1); - } - } - elect_commit(kv_empty_addr + (kv_stage) * 8); - kv_stage += 1; - if (kv_stage == 6) { kv_stage = 0; kv_phase ^= 1; } - } - mbarrier_wait(kv_full_addr + (kv_stage) * 8, kv_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - int _mma_a_lo_3 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (kv_stage) * 2048); - int _mma_b_lo_3 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (value_quarter_1) * 128); - { - uint64_t _mma_ss_a_desc_3 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_3); - uint64_t _mma_ss_b_desc_3 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_3); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134349968, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134349968, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134349968, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134349968, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_3, 1018U); - incr_smem_desc_lo(_mma_ss_b_desc_3, 58U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134349968, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134349968, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134349968, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (score_col_0)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 134349968, 1); - } - } - kv_stage += 1; - if (kv_stage == 6) { kv_stage = 0; kv_phase ^= 1; } - elect_commit(score_full_addr + (score_stage_1) * 8); - score_stage_1 += 1; - if (score_stage_1 == 2) { score_stage_1 = 0; score_phase_1 ^= 1; } - unsigned int v_cur = (kv_stage + 1) % 6; - mbarrier_wait(p_full_addr + (p_stage_1) * 8, p_phase_1); - mbarrier_wait(o_empty_addr, _phase_o_empty_0); - _phase_o_empty_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - int _mma_a_lo_4 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (v_cur) * 2048); - int _mma_b_lo_4 = make_warp_uniform(((((smem_p_addr) >> 4) & 0x3FFF) | 0x400000) + (p_stage_1) * 128); - { - uint64_t _mma_ss_a_desc_4 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_4); - uint64_t _mma_ss_b_desc_4 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_4); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134382736, ((first_output) ? 0 : 1)); - } - incr_smem_desc_lo(_mma_ss_a_desc_4, 128U); - incr_smem_desc_lo(_mma_ss_b_desc_4, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134382736, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_4, 128U); - incr_smem_desc_lo(_mma_ss_b_desc_4, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134382736, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_4, 128U); - incr_smem_desc_lo(_mma_ss_b_desc_4, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134382736, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_4, 128U); - incr_smem_desc_lo(_mma_ss_b_desc_4, 58U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134382736, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_4, 128U); - incr_smem_desc_lo(_mma_ss_b_desc_4, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134382736, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_4, 128U); - incr_smem_desc_lo(_mma_ss_b_desc_4, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134382736, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_4, 128U); - incr_smem_desc_lo(_mma_ss_b_desc_4, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_4, _mma_ss_b_desc_4, 134382736, 1); - } - } - elect_commit(o_full_addr); - elect_commit(kv_empty_addr + (v_cur) * 8); - elect_commit(p_empty_addr + (p_stage_1) * 8); - p_stage_1 += 1; - if (p_stage_1 == 2) { p_stage_1 = 0; p_phase_1 ^= 1; } - first_output = 0; - } - unsigned int v_tail = (kv_stage + 5) % 6; - mbarrier_wait(p_full_addr + (p_stage_1) * 8, p_phase_1); - mbarrier_wait(o_empty_addr, _phase_o_empty_0); - _phase_o_empty_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - int _mma_a_lo_5 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (v_tail) * 2048); - int _mma_b_lo_5 = make_warp_uniform(((((smem_p_addr) >> 4) & 0x3FFF) | 0x400000) + (p_stage_1) * 128); - { - uint64_t _mma_ss_a_desc_5 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_5); - uint64_t _mma_ss_b_desc_5 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_5); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134382736, ((first_output) ? 0 : 1)); - } - incr_smem_desc_lo(_mma_ss_a_desc_5, 128U); - incr_smem_desc_lo(_mma_ss_b_desc_5, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134382736, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_5, 128U); - incr_smem_desc_lo(_mma_ss_b_desc_5, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134382736, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_5, 128U); - incr_smem_desc_lo(_mma_ss_b_desc_5, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134382736, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_5, 128U); - incr_smem_desc_lo(_mma_ss_b_desc_5, 58U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134382736, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_5, 128U); - incr_smem_desc_lo(_mma_ss_b_desc_5, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134382736, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_5, 128U); - incr_smem_desc_lo(_mma_ss_b_desc_5, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134382736, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_5, 128U); - incr_smem_desc_lo(_mma_ss_b_desc_5, 2U); - if (elect_sync()) { - tcgen05_mma_f16((tmem_source_exact_scratch + (80)), _mma_ss_a_desc_5, _mma_ss_b_desc_5, 134382736, 1); - } - } - elect_commit(o_full_addr); - elect_commit(kv_empty_addr + (v_tail) * 8); - elect_commit(p_empty_addr + (p_stage_1) * 8); - p_stage_1 += 1; - if (p_stage_1 == 2) { p_stage_1 = 0; p_phase_1 ^= 1; } - elect_commit(q_empty_addr + (q_stage) * 8); - q_phase ^= 1; - } - mbarrier_wait(work_full_addr + (work_stage_2) * 8, _phase_work_full_2); - uint32_t _clc_valid_2 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_2) - : "r"(work_response_addr + work_stage_2 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_6 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_6) - : "r"(work_response_addr + work_stage_2 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_7 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_7) - : "r"(work_response_addr + work_stage_2 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_8 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_8) - : "r"(work_response_addr + work_stage_2 * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(work_empty_addr + work_stage_2 * 8), "r"(0) : "memory"); - work_stage_2 += 1; - if (work_stage_2 == 2) { work_stage_2 = 0; _phase_work_full_2 ^= 1; } - if (_clc_valid_2 == 0) { - break; - } - work_x_2 = _clc_ctaid_6; - work_y_2 = _clc_ctaid_7; - work_z_2 = _clc_ctaid_8; - } - #pragma unroll - for (int tail = 0; tail < 2; tail++) { - elect_commit(score_full_addr + (tail) * 8); - } - } - } - // ---- Role: page_loader ---- - if (warp == 9) { - { // page_loader_main - unsigned int work_stage_3 = 0; - unsigned int page_stage = 0; - unsigned int page_phase = 1; - unsigned int work_x_3 = blockIdx.x; - unsigned int work_y_3 = blockIdx.y; - unsigned int work_z_3 = blockIdx.z; - unsigned int _phase_work_full_3 = 0; - #pragma unroll 1 - for (unsigned int _work_3 = 0; _work_3 < 120; _work_3++) { - int q_begin_3 = cum_seq_lens_q[work_z_3]; - int q_end_3 = cum_seq_lens_q[work_z_3 + 1]; - int q_len_3 = q_end_3 - q_begin_3; - int query_idx_3 = (unsigned int)q_begin_3 + work_x_3; - int active_topk_3 = 0; - if (work_x_3 < (unsigned int)q_len_3) { - int _max_0 = ((sparse_topk_lens[query_idx_3] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_3] + sparse_topk_lens_offset) : (0)); - int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); - active_topk_3 = _min_0; - } - int num_steps_3 = (active_topk_3 + 128 - 1) / 128; - int num_pairs = (num_steps_3 + 1) / 2; - int _max_1 = ((active_topk_3 - 1 & -4) > (0) ? (active_topk_3 - 1 & -4) : (0)); - int last_aligned = _max_1; - int* swa_row = swa_indices + (query_idx_3 * swa_index_stride); - int* compressed_row = compressed_indices + (query_idx_3 * compressed_index_stride); - #pragma unroll 1 - for (int pair_2 = 0; pair_2 < num_pairs; pair_2++) { - int lane_base = lane * 4; - #pragma unroll - for (int duplicate = 0; duplicate < 2; duplicate++) { - mbarrier_wait(page_empty_addr + (page_stage) * 8, page_phase); - int page0 = pair_2 * 2 * 128 + lane_base; - int page1 = page0 + 128; - int safe0 = ((page0 <= last_aligned) ? page0 : last_aligned); - int safe1 = ((page1 <= last_aligned) ? page1 : last_aligned); - int* src0 = ((safe0 < 128) ? swa_row + safe0 : compressed_row + (safe0 - 128)); - int* src1 = ((safe1 < 128) ? swa_row + safe1 : compressed_row + (safe1 - 128)); - asm volatile("cp.async.ca.shared::cta.global [%0], [%1], 16;" - :: "r"(smem_page_offsets_addr + page_stage * 1024 + (unsigned int)(lane_base * 4)), "l"(src0)); - asm volatile("cp.async.ca.shared::cta.global [%0], [%1], 16;" - :: "r"(smem_page_offsets_addr + page_stage * 1024 + (unsigned int)((128 + lane_base) * 4)), "l"(src1)); - asm volatile( - "{\n\t" - "cp.async.mbarrier.arrive.shared::cta.b64 [%0];\n\t" - "}" - :: "r"(page_full_addr + (page_stage) * 8) : "memory"); - mbarrier_arrive(page_full_addr + (page_stage) * 8); - page_stage += 1; - if (page_stage == 6) { page_stage = 0; page_phase ^= 1; } - } - } - mbarrier_wait(work_full_addr + (work_stage_3) * 8, _phase_work_full_3); - uint32_t _clc_valid_0 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_0) - : "r"(work_response_addr + work_stage_3 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_0 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_0) - : "r"(work_response_addr + work_stage_3 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_1 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_1) - : "r"(work_response_addr + work_stage_3 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_2 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_2) - : "r"(work_response_addr + work_stage_3 * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(work_empty_addr + work_stage_3 * 8), "r"(0) : "memory"); - work_stage_3 += 1; - if (work_stage_3 == 2) { work_stage_3 = 0; _phase_work_full_3 ^= 1; } - if (_clc_valid_0 == 0) { - break; - } - work_x_3 = _clc_ctaid_0; - work_y_3 = _clc_ctaid_1; - work_z_3 = _clc_ctaid_2; - } - } - } - // ---- Role: scheduler ---- - if (warp == 10) { - { // scheduler_main - asm volatile("griddepcontrol.wait;" ::: "memory"); - unsigned int work_stage_4 = 0; - unsigned int throttle_stage = 0; - unsigned int throttle_phase = 0; - unsigned int _phase_work_empty = 1; - unsigned int _phase_work_full_4 = 0; - #pragma unroll 1 - for (unsigned int _work_4 = 0; _work_4 < 120; _work_4++) { - mbarrier_wait(throttle_full_addr + (throttle_stage) * 8, throttle_phase); - mbarrier_arrive(throttle_empty_addr + (throttle_stage) * 8); - throttle_stage += 1; - if (throttle_stage == 2) { throttle_stage = 0; throttle_phase ^= 1; } - mbarrier_wait(work_empty_addr + (work_stage_4) * 8, _phase_work_empty); - if (elect_sync()) { - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [remAddr32], %2;\n\t" - "}" - :: "r"(work_full_addr + work_stage_4 * 8), "r"(0), "r"((uint32_t)(16)) : "memory"); - asm volatile( - "fence.proxy.async.shared::cta;\n\t" - "clusterlaunchcontrol.try_cancel.async.shared::cta" - ".mbarrier::complete_tx::bytes.b128" - " [%0], [%1];" - :: "r"(work_response_addr + work_stage_4 * 16 + 0 * 16), "r"(work_full_addr + work_stage_4 * 8) - : "memory"); - } - mbarrier_wait(work_full_addr + (work_stage_4) * 8, _phase_work_full_4); - uint32_t _clc_valid_5 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_5) - : "r"(work_response_addr + work_stage_4 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_15 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_15) - : "r"(work_response_addr + work_stage_4 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_16 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_16) - : "r"(work_response_addr + work_stage_4 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_17 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_17) - : "r"(work_response_addr + work_stage_4 * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(work_empty_addr + work_stage_4 * 8), "r"(0) : "memory"); - work_stage_4 += 1; - if (work_stage_4 == 2) { work_stage_4 = 0; _phase_work_empty ^= 1; _phase_work_full_4 ^= 1; } - if (_clc_valid_5 == 0) { - break; - } - } - } - } - // ---- Role: padding ---- - if (warp == 11) { - { // padding_main - unsigned int work_stage_5 = 0; - unsigned int work_x_4 = blockIdx.x; - unsigned int work_y_4 = blockIdx.y; - unsigned int work_z_4 = blockIdx.z; - unsigned int _phase_work_full_5 = 0; - #pragma unroll 1 - for (unsigned int _work_5 = 0; _work_5 < 120; _work_5++) { - mbarrier_wait(work_full_addr + (work_stage_5) * 8, _phase_work_full_5); - uint32_t _clc_valid_6 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_6) - : "r"(work_response_addr + work_stage_5 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_18 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_18) - : "r"(work_response_addr + work_stage_5 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_19 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_19) - : "r"(work_response_addr + work_stage_5 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_20 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_20) - : "r"(work_response_addr + work_stage_5 * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(work_empty_addr + work_stage_5 * 8), "r"(0) : "memory"); - work_stage_5 += 1; - if (work_stage_5 == 2) { work_stage_5 = 0; _phase_work_full_5 ^= 1; } - if (_clc_valid_6 == 0) { - break; - } - work_x_4 = _clc_ctaid_18; - work_y_4 = _clc_ctaid_19; - work_z_4 = _clc_ctaid_20; - } - } - } - // ---- Role: loader ---- - if (warp >= 12 && warp <= 19) { - asm volatile("setmaxnreg.dec.sync.aligned.u32 72;"); - { // loader_main - asm volatile("griddepcontrol.wait;" ::: "memory"); - const int load_rank = warp - 12; - unsigned int work_stage_6 = 0; - unsigned int throttle_stage_1 = 0; - unsigned int throttle_phase_1 = 1; - unsigned int q_stage_1 = 0; - unsigned int q_phase_1 = 1; - unsigned int kv_stage_1 = 0; - unsigned int kv_phase_1 = 1; - unsigned int page_k_stage = 0; - unsigned int page_k_phase = 0; - unsigned int page_v_stage = 1; - unsigned int page_v_phase = 0; - unsigned int work_x_5 = blockIdx.x; - unsigned int work_y_5 = blockIdx.y; - unsigned int work_z_5 = blockIdx.z; - unsigned int _phase_work_full_6 = 0; - #pragma unroll 1 - for (unsigned int _work_6 = 0; _work_6 < 120; _work_6++) { - int q_begin_4 = cum_seq_lens_q[work_z_5]; - int q_end_4 = cum_seq_lens_q[work_z_5 + 1]; - int q_len_4 = q_end_4 - q_begin_4; - int query_idx_4 = (unsigned int)q_begin_4 + work_x_5; - int active_topk_4 = 0; - if (work_x_5 < (unsigned int)q_len_4) { - int _max_2 = ((sparse_topk_lens[query_idx_4] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_4] + sparse_topk_lens_offset) : (0)); - int _min_1 = ((_max_2) < (sparse_topk) ? (_max_2) : (sparse_topk)); - active_topk_4 = _min_1; - } - int num_steps_4 = (active_topk_4 + 128 - 1) / 128; - int head_group_1 = work_y_5 / 4; - int value_quarter_2 = work_y_5 % 4; - mbarrier_wait(throttle_empty_addr + (throttle_stage_1) * 8, throttle_phase_1); - mbarrier_arrive(throttle_full_addr + (throttle_stage_1) * 8); - throttle_stage_1 += 1; - if (throttle_stage_1 == 2) { throttle_stage_1 = 0; throttle_phase_1 ^= 1; } - if (num_steps_4 > 0) { - mbarrier_wait(q_empty_addr + (q_stage_1) * 8, q_phase_1); - if (load_rank == 0) { - if (elect_sync()) { - mbarrier_arrive_expect_tx(q_full_addr + (q_stage_1) * 8, 8192); - #pragma unroll - for (int hs_2 = 0; hs_2 < 4; hs_2++) { - tma_4d_gmem2smem(smem_q_addr + (unsigned int)(hs_2 * 2048), tmap_q, 0, head_group_1 * 8, hs_2 * 2, query_idx_4, q_full_addr + (q_stage_1) * 8); - } - } - } - mbarrier_wait(page_full_addr + (page_k_stage) * 8, page_k_phase); - int page_base = 0; - #pragma unroll - for (int hs_3 = 0; hs_3 < 4; hs_3++) { - int hs_sel_2 = ((hs_3 == 3) ? value_quarter_2 : ((value_quarter_2 > hs_3) ? hs_3 : hs_3 + 1)); - mbarrier_wait(kv_empty_addr + (kv_stage_1) * 8, kv_phase_1); - if (load_rank == 0) { - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_1) * 8, 32768); - } - } - asm volatile("barrier.sync 9, 256;" ::: "memory"); - #pragma unroll - for (int group = 0; group < 4; group++) { - const int row_base = load_rank * 4 + group * 32; - int rows[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&rows[0])), "=r"(*reinterpret_cast(&rows[(0) + 1])), "=r"(*reinterpret_cast(&rows[(0) + 2])), "=r"(*reinterpret_cast(&rows[(0) + 3])) - : "r"(smem_page_offsets_addr + page_k_stage * 1024 + (unsigned int)((page_base + row_base) * 4))); - int dst_k = smem_kv_addr + kv_stage_1 * 32768 + (unsigned int)(row_base * 64 * 2); - if (elect_sync()) { - { - tma_gather4_gmem2smem(dst_k, tmap_swa_kv, hs_sel_2 * 128, rows[0], rows[1], rows[2], rows[3], kv_full_addr + (kv_stage_1) * 8); - tma_gather4_gmem2smem(dst_k + 16384, tmap_swa_kv, hs_sel_2 * 128 + 64, rows[0], rows[1], rows[2], rows[3], kv_full_addr + (kv_stage_1) * 8); - } - } - } - kv_stage_1 += 1; - if (kv_stage_1 == 6) { kv_stage_1 = 0; kv_phase_1 ^= 1; } - } - #pragma unroll 1 - for (int tile_2 = 0; tile_2 < num_steps_4 - 1; tile_2++) { - int next_tile = tile_2 + 1; - if ((next_tile & 1) == 0) { - mbarrier_wait(page_full_addr + (page_k_stage) * 8, page_k_phase); - } - int page_base_0 = (next_tile & 1) * 128; - #pragma unroll - for (int hs_4 = 0; hs_4 < 4; hs_4++) { - int hs_sel_3 = ((hs_4 == 3) ? value_quarter_2 : ((value_quarter_2 > hs_4) ? hs_4 : hs_4 + 1)); - mbarrier_wait(kv_empty_addr + (kv_stage_1) * 8, kv_phase_1); - if (load_rank == 0) { - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (kv_stage_1) * 8, 32768); - } - } - asm volatile("barrier.sync 9, 256;" ::: "memory"); - #pragma unroll - for (int group_1 = 0; group_1 < 4; group_1++) { - const int row_base_1 = load_rank * 4 + group_1 * 32; - int rows_1[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&rows_1[0])), "=r"(*reinterpret_cast(&rows_1[(0) + 1])), "=r"(*reinterpret_cast(&rows_1[(0) + 2])), "=r"(*reinterpret_cast(&rows_1[(0) + 3])) - : "r"(smem_page_offsets_addr + page_k_stage * 1024 + (unsigned int)((page_base_0 + row_base_1) * 4))); - int dst_k_1 = smem_kv_addr + kv_stage_1 * 32768 + (unsigned int)(row_base_1 * 64 * 2); - if (elect_sync()) { - if (next_tile == 0) { - tma_gather4_gmem2smem(dst_k_1, tmap_swa_kv, hs_sel_3 * 128, rows_1[0], rows_1[1], rows_1[2], rows_1[3], kv_full_addr + (kv_stage_1) * 8); - tma_gather4_gmem2smem(dst_k_1 + 16384, tmap_swa_kv, hs_sel_3 * 128 + 64, rows_1[0], rows_1[1], rows_1[2], rows_1[3], kv_full_addr + (kv_stage_1) * 8); - } else { - tma_gather4_gmem2smem(dst_k_1, tmap_compressed_kv, hs_sel_3 * 128, rows_1[0], rows_1[1], rows_1[2], rows_1[3], kv_full_addr + (kv_stage_1) * 8); - tma_gather4_gmem2smem(dst_k_1 + 16384, tmap_compressed_kv, hs_sel_3 * 128 + 64, rows_1[0], rows_1[1], rows_1[2], rows_1[3], kv_full_addr + (kv_stage_1) * 8); - } - } - } - kv_stage_1 += 1; - if (kv_stage_1 == 6) { kv_stage_1 = 0; kv_phase_1 ^= 1; } - } - if ((next_tile & 1) != 0 && next_tile < num_steps_4 - 1) { - mbarrier_arrive(page_empty_addr + (page_k_stage) * 8); - page_k_stage += 1; - if (page_k_stage == 6) { page_k_stage = 0; page_k_phase ^= 1; } - page_k_stage += 1; - if (page_k_stage == 6) { page_k_stage = 0; page_k_phase ^= 1; } - } - if ((tile_2 & 1) == 0) { - mbarrier_wait(page_full_addr + (page_v_stage) * 8, page_v_phase); - } - if ((tile_2 & 1) != 0) { - mbarrier_arrive(page_empty_addr + (page_v_stage) * 8); - page_v_stage += 1; - if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } - page_v_stage += 1; - if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } - } - } - int last_tile = num_steps_4 - 1; - mbarrier_arrive(page_empty_addr + (page_k_stage) * 8); - page_k_stage += 1; - if (page_k_stage == 6) { page_k_stage = 0; page_k_phase ^= 1; } - page_k_stage += 1; - if (page_k_stage == 6) { page_k_stage = 0; page_k_phase ^= 1; } - if ((last_tile & 1) == 0) { - mbarrier_wait(page_full_addr + (page_v_stage) * 8, page_v_phase); - } - mbarrier_arrive(page_empty_addr + (page_v_stage) * 8); - page_v_stage += 1; - if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } - page_v_stage += 1; - if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } - } - mbarrier_wait(work_full_addr + (work_stage_6) * 8, _phase_work_full_6); - uint32_t _clc_valid_1 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_1) - : "r"(work_response_addr + work_stage_6 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_3 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_3) - : "r"(work_response_addr + work_stage_6 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_4 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_4) - : "r"(work_response_addr + work_stage_6 * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_5 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_5) - : "r"(work_response_addr + work_stage_6 * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - asm volatile( - "{\n\t" - ".reg .b32 remAddr32;\n\t" - "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" - "mbarrier.arrive.release.cta.shared::cluster.b64 _, [remAddr32];\n\t" - "}" - :: "r"(work_empty_addr + work_stage_6 * 8), "r"(0) : "memory"); - work_stage_6 += 1; - if (work_stage_6 == 2) { work_stage_6 = 0; _phase_work_full_6 ^= 1; } - if (_clc_valid_1 == 0) { - break; - } - work_x_5 = _clc_ctaid_3; - work_y_5 = _clc_ctaid_4; - work_z_5 = _clc_ctaid_5; - } - } - } - - // Cleanup -} - -} // extern "C" diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_f55f69afe4ac3ec96817_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_f55f69afe4ac3ec96817_binding.cu deleted file mode 100644 index d0f749f0def..00000000000 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_f55f69afe4ac3ec96817_binding.cu +++ /dev/null @@ -1,441 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -// Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_f55f69afe4ac3ec96817'. -#include -#include -#include -#include - -#include "tvm_ffi_utils.h" - -#include -#include -#include -#include -#include -#include -#include - -extern "C" __global__ void kernel_cake_dsv4_f55f69afe4ac3ec96817(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ cum_seq_lens_q, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); - - -namespace cake_host_shim_90ab13be7c9c8334 { - -using tvm::ffi::Optional; -using tvm::ffi::TensorView; - -class ScopedCudaDevice { - public: - explicit ScopedCudaDevice(int device_id) { - cudaError_t error = cudaGetDevice(&previous_device_); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed before host-shim launch: cudaError=" - << static_cast(error); - if (previous_device_ != device_id) { - error = cudaSetDevice(device_id); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaSetDevice failed before host-shim launch for cuda:" - << device_id << ": cudaError=" << static_cast(error); - restore_ = true; - } - } - - ScopedCudaDevice(const ScopedCudaDevice&) = delete; - ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; - - ~ScopedCudaDevice() noexcept { - if (restore_) { - (void)cudaSetDevice(previous_device_); - } - } - - private: - int previous_device_ = -1; - bool restore_ = false; -}; - -inline int64_t CakeDeviceMultiprocessorCount(int device_id) { - constexpr int kMaxCachedCudaDevices = 64; - TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) - << "physical-SM-count cache does not cover cuda:" << device_id; - static std::atomic count_by_device[kMaxCachedCudaDevices]{}; - int cached = count_by_device[device_id].load(std::memory_order_acquire); - if (cached > 0) return cached; - - int count = 0; - cudaError_t error = cudaDeviceGetAttribute( - &count, cudaDevAttrMultiProcessorCount, device_id); - TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) - << "querying multiProcessorCount failed for cuda:" << device_id - << ": cudaError=" << static_cast(error) << ", count=" << count; - // Concurrent first launches may repeat the immutable device query, but all - // publication is atomic and every later hot-path lookup is one acquire load. - count_by_device[device_id].store(count, std::memory_order_release); - return count; -} - -inline void CheckCudaTensor(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) - << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; -} - -inline void CheckSameCudaDevice( - const TensorView& t, - const TensorView& reference, - const char* name, - const char* reference_name) { - TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) - << name << " must be on the same CUDA device as " << reference_name - << ": got cuda:" << t.device().device_id - << " versus cuda:" << reference.device().device_id; -} - -inline void CheckCurrentCudaDevice( - const TensorView& reference, - const char* reference_name) { - int current_device = -1; - cudaError_t error = cudaGetDevice(¤t_device); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed while validating " << reference_name - << ": cudaError=" << static_cast(error); - TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) - << "current CUDA device must match " << reference_name - << ": current=cuda:" << current_device - << ", tensor=cuda:" << reference.device().device_id; -} - -inline void CheckContiguous(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; -} - -inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { - DLDataType d = t.dtype(); - TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) - << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits - << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits - << ", lanes=" << (int)d.lanes << ")"; -} - -// A logical axis.outer(trailing) folds every source dim above the trailing -// dimensions. Shape products are independent of physical strides, so verify -// the leading dimensions form one dense row-major chain instead of inventing -// a "folded stride". The descriptor reads its exact adjacent physical step -// separately through stride[-(trailing + 1)]. -inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { - TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) - << name << " cannot fold leading dimensions above " << trailing - << " trailing dims from ndim=" << t.ndim(); - int outer_last = t.ndim() - trailing - 1; - if (outer_last <= 0) { - return; - } - int64_t step = t.stride(outer_last); - TVM_FFI_CHECK(step > 0, ValueError) - << name << " physical strides must be positive"; - int64_t expected = step; - for (int axis = outer_last - 1; axis >= 0; --axis) { - expected *= t.size(axis + 1); - if (t.size(axis) > 1) { - TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) - << name << " leading dims are not physically foldable above " << trailing - << " trailing dims: stride(" << axis << ")=" << t.stride(axis) - << ", expected " << expected; - } - } -} - -// 4D TMA descriptor for buffer 'tmap_q' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 3, ValueError) - << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - int64_t d2 = t.size(t.ndim() - 2); - int64_t d3 = t.size(t.ndim() - 3); - TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) - << "TMA source 'tmap_q' trailing dims must be positive"; - TVM_FFI_CHECK(d2 % 64 == 0, ValueError) - << "TMA source 'tmap_q' extent " << d2 - << " must divide exactly by " << 64; - uint64_t global_dim[4] = {(uint64_t)(d1), (uint64_t)(64), (uint64_t)((d2 / 64)), (uint64_t)(d3)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; - TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) - << "TMA box (128, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; - int64_t carrier_stride_0 = d1; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; - int64_t carrier_stride_1 = (64 * d1); - TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; - TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; - TVM_FFI_CHECK((carrier_stride_1 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; - int64_t carrier_stride_2 = (d2 * d1); - TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; - TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; - TVM_FFI_CHECK((carrier_stride_2 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; - uint64_t global_strides[3] = { - (uint64_t)((carrier_stride_0 * 8) / 8), - (uint64_t)((carrier_stride_1 * 8) / 8), - (uint64_t)((carrier_stride_2 * 8) / 8), - }; - uint32_t box_dim[4] = {128u, 64u, 1u, 1u}; - uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_L2_128B, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_swa_kv' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_swa_kv' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; - TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (128, 1) exceeds resolved global dims for 'tmap_swa_kv'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 8) / 8), - }; - uint32_t box_dim[2] = {128u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_compressed_kv' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_compressed_kv(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_compressed_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_compressed_kv' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_compressed_kv' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_compressed_kv"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_compressed_kv' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved a non-positive global dim"; - TVM_FFI_CHECK(128u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (128, 1) exceeds resolved global dims for 'tmap_compressed_kv'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 8) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_compressed_kv' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 8) / 8), - }; - uint32_t box_dim[2] = {128u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_kv') failed: CUresult=" << (int)r; - return tm; -} - -void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_tmap_compressed_kv, TensorView arg_O, TensorView arg_cum_seq_lens_q, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, int64_t grid_x, int64_t grid_y, int64_t grid_z) { - DLDevice dev = arg_tmap_q.device(); - ScopedCudaDevice device_guard(dev.device_id); - CheckCudaTensor(arg_tmap_q, "tmap_q"); - CheckDtype(arg_tmap_q, "tmap_q", 1, 8, 1); - CheckContiguous(arg_tmap_q, "tmap_q"); - CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); - CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 1, 8, 1); - CheckCudaTensor(arg_tmap_compressed_kv, "tmap_compressed_kv"); - CheckDtype(arg_tmap_compressed_kv, "tmap_compressed_kv", 1, 8, 1); - CheckCudaTensor(arg_O, "O"); - CheckDtype(arg_O, "O", 4, 16, 1); - CheckContiguous(arg_O, "O"); - CheckCudaTensor(arg_cum_seq_lens_q, "cum_seq_lens_q"); - CheckDtype(arg_cum_seq_lens_q, "cum_seq_lens_q", 0, 32, 1); - CheckContiguous(arg_cum_seq_lens_q, "cum_seq_lens_q"); - CheckCudaTensor(arg_swa_indices, "swa_indices"); - CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); - CheckContiguous(arg_swa_indices, "swa_indices"); - CheckCudaTensor(arg_compressed_indices, "compressed_indices"); - CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); - CheckContiguous(arg_compressed_indices, "compressed_indices"); - CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); - CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckCudaTensor(arg_sinks, "sinks"); - CheckDtype(arg_sinks, "sinks", 2, 32, 1); - CheckContiguous(arg_sinks, "sinks"); - CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); - CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); - CheckContiguous(arg_bmm1_scale, "bmm1_scale"); - CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); - CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); - CheckContiguous(arg_bmm2_scale, "bmm2_scale"); - TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) - << "scalar 'num_heads' value " << arg_num_heads - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) - << "scalar 'swa_index_stride' value " << arg_swa_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) - << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) - << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) - << "scalar 'sparse_topk' value " << arg_sparse_topk - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) - << "scalar 'has_sinks' value " << arg_has_sinks - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) - << "scalar 'total_work_items' value " << arg_total_work_items - << " is outside i32 range [-2147483648, 2147483647]"; - CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); - CheckSameCudaDevice(arg_tmap_compressed_kv, arg_tmap_q, "tmap_compressed_kv", "tmap_q"); - CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); - CheckSameCudaDevice(arg_cum_seq_lens_q, arg_tmap_q, "cum_seq_lens_q", "tmap_q"); - CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); - CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); - CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); - CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); - CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); - CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); - TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) - << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y - << ", " << grid_z << ")"; - - cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); - CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); - CUtensorMap p_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); - CUtensorMap p_tmap_compressed_kv = EncodeTma_tmap_compressed_kv(arg_tmap_compressed_kv); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); - int* p_cum_seq_lens_q = static_cast(arg_cum_seq_lens_q.data_ptr()); - int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); - int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); - int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); - float* p_sinks = static_cast(arg_sinks.data_ptr()); - float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); - float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); - int32_t v_num_heads = (int32_t)arg_num_heads; - int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; - int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; - int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; - int32_t v_sparse_topk = (int32_t)arg_sparse_topk; - int32_t v_has_sinks = (int32_t)arg_has_sinks; - int32_t v_total_work_items = (int32_t)arg_total_work_items; - void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_tmap_compressed_kv, &p_O, &p_cum_seq_lens_q, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_has_sinks, &v_total_work_items}; - - static std::mutex smem_mu; - static auto* smem_status_by_device = - new std::unordered_map(); - cudaError_t smem_status = cudaSuccess; - { - std::lock_guard lock(smem_mu); - auto it = smem_status_by_device->find(dev.device_id); - if (it == smem_status_by_device->end()) { - int current_device = -1; - smem_status = cudaGetDevice(¤t_device); - if (smem_status == cudaSuccess && current_device != dev.device_id) { - smem_status = cudaErrorInvalidDevice; - } - if (smem_status == cudaSuccess) { - smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_f55f69afe4ac3ec96817), - cudaFuncAttributeMaxDynamicSharedMemorySize, - 203136); - } - smem_status_by_device->emplace(dev.device_id, smem_status); - } else { - smem_status = it->second; - } - } - TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_f55f69afe4ac3ec96817 failed: " - << cudaGetErrorString(smem_status); - dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); - dim3 block(512u, 1u, 1u); - - cudaLaunchConfig_t config{}; - config.gridDim = grid; - config.blockDim = block; - config.dynamicSmemBytes = 203136u; - config.stream = stream; - cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_f55f69afe4ac3ec96817), kargs); - TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_f55f69afe4ac3ec96817 failed: " - << cudaGetErrorString(launch_status); -} - -} // namespace cake_host_shim_90ab13be7c9c8334 - -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_90ab13be7c9c8334::Run); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_f55f69afe4ac3ec96817_kernel.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_f55f69afe4ac3ec96817_kernel.cu deleted file mode 100644 index 53eb116e5d9..00000000000 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_f55f69afe4ac3ec96817_kernel.cu +++ /dev/null @@ -1,2972 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -typedef signed char int8_t; -typedef unsigned char uint8_t; -typedef unsigned short uint16_t; -typedef unsigned int uint32_t; -#if defined(__CUDACC_RTC__) -typedef unsigned long long uint64_t; -#else -typedef unsigned long uint64_t; -#endif -static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); -typedef signed int int32_t; -typedef short int int16_t; -struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; -struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; -static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); -static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); - -#if defined(__CUDACC_RTC__) -typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; -#else -#include -#endif - -static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); -static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); -#include -#include - -#define CAKE_INF CUDART_INF_F -#define TMEM_NCOLS 512 -#define TMEM_TMEM_OFFSET 0 -#define NUM_Q_PIPE_STAGES 2 -#define NUM_KV_PIPE_STAGES 4 -#define NUM_PAGE_PIPE_STAGES 6 -#define NUM_SCORE_PIPE_STAGES 2 -#define NUM_STATS_PIPE_STAGES 2 -#define NUM_WORK_PIPE_STAGES 2 -#define NUM_THROTTLE_PIPE_STAGES 2 -#define SMEM_SMEM_Q_OFF 0 -#define SMEM_SMEM_Q_STAGE_BYTES 32768 -#define SMEM_SMEM_Q_STRIDE 32768 -#define SMEM_SMEM_K_OFF 65536 -#define SMEM_SMEM_K_STAGE_BYTES 32768 -#define SMEM_SMEM_K_STRIDE 32768 -#define SMEM_SMEM_V_OFF 65536 -#define SMEM_SMEM_V_STAGE_BYTES 32768 -#define SMEM_SMEM_V_STRIDE 32768 -#define SMEM_SMEM_PAGE_OFFSETS_OFF 196608 -#define SMEM_SMEM_PAGE_OFFSETS_STAGE_BYTES 1024 -#define SMEM_SMEM_PAGE_OFFSETS_STRIDE 1024 -#define SMEM_WORK_RESPONSE_OFF 202752 -#define SMEM_WORK_RESPONSE_STAGE_BYTES 16 -#define SMEM_WORK_RESPONSE_STRIDE 16 -#define SMEM_TOTAL 203136 -#define THREADS 512 - -#include - -__device__ __forceinline__ uint32_t elect_sync() { - uint32_t pred = 0; - asm volatile( - "{\n\t" - ".reg .pred %%px;\n\t" - "elect.sync _|%%px, %1;\n\t" - "@%%px mov.s32 %0, 1;\n\t" - "}\n" - : "+r"(pred) - : "r"(0xFFFFFFFF)); - return pred; -} - - -__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { - asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" - :: "r"(mbar_addr), "r"(count) : "memory"); -} - -__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { - asm volatile("mbarrier.init.b64 [%0], %1;" - :: "l"(mbar_addr), "r"(count) : "memory"); -} - - -__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - - -// CTA-local pipelines have short, resident producer/consumer edges. Omitting -// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the -// explicit loop still makes this helper blocking until acquire succeeds. -__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE;\n\t" - "bra.uni LAB_WAIT;\n\t" - "DONE:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Source-faithful relaxed CTA wait used only by a typed protocol that does -// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. -__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, 10000000;\n\t" - "@P1 bra.uni DONE_RELAXED;\n\t" - "bra.uni LAB_WAIT_RELAXED;\n\t" - "DONE_RELAXED:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Exact source ports may request the PTX suspendTimeHint operand explicitly. -// The hint is expressed in nanoseconds and is kept separate from the canonical -// no-hint CTA helper so unrelated schedules retain their existing retry path. -__device__ __forceinline__ void mbarrier_wait_suspend( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_SUSPEND:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_SUSPEND;\n\t" - "bra.uni LAB_WAIT_SUSPEND;\n\t" - "DONE_SUSPEND:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE_CLUSTER;\n\t" - "bra.uni LAB_WAIT_CLUSTER;\n\t" - "DONE_CLUSTER:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - ".reg .u32 WAIT_ADDR;\n\t" - "mov.u32 WAIT_ADDR, %0;\n\t" - "LAB_WAIT_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [WAIT_ADDR], %1, %2;\n\t" - "@P1 bra.uni DONE_HINT;\n\t" - "bra.uni LAB_WAIT_HINT;\n\t" - "DONE_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -// Exact unqualified CTA wait used by source schedules whose PTX intentionally -// omits the acquire qualifier while retaining a typed suspendTimeHint operand. -__device__ __forceinline__ void mbarrier_wait_relaxed_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED_HINT:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra DONE_RELAXED_HINT;\n\t" - "bra LAB_WAIT_RELAXED_HINT;\n\t" - "DONE_RELAXED_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); -} - -__device__ __forceinline__ void mbarrier_wait_cluster_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" - "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" - "DONE_CLUSTER_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_suspend( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait_cluster(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); - } -} - - -__device__ __forceinline__ void tcgen05_mma_f8f6f4( - int taddr, uint64_t a_desc, uint64_t b_desc, - uint32_t i_desc, int enable_input_d) { - asm volatile( - "{\n\t" - ".reg .pred p;\n\t" - "setp.ne.b32 p, %4, 0;\n\t" - "tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], %1, %2, %3, p;\n\t" - "}\n" - :: "r"(taddr), "l"(a_desc), "l"(b_desc), - "r"(i_desc), "r"(enable_input_d) - : "memory"); -} - - -__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { - return (x & 0x3FFFFULL) >> 4ULL; -} - - -union MmaSmemDesc { - uint64_t u64; - uint32_t u32[2]; -}; - -__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { - MmaSmemDesc tmp; - tmp.u64 = smem_desc; - tmp.u32[0] += offset; - smem_desc = tmp.u64; -} - - -__device__ __forceinline__ void elect_commit(int mbar_addr) { - asm volatile( - "{\n\t" - ".reg .pred leader;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" - ".shared::cluster.b64 [%0];\n\t" - "}\n" - :: "r"(mbar_addr)); -} - - -__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { - asm volatile( - "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" - :: "r"(mbar_addr), "r"(bytes) : "memory"); -} - - -__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7," - " %8, %9, %10, %11, %12, %13, %14, %15," - " %16, %17, %18, %19, %20, %21, %22, %23," - " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), - "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), - "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), - "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), - "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), - "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), - "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ float approx_exp2(float x) { - float y; - asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float approx_rcp(float x) { - float y; - asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float max_noftz(float a, float b) { - float c; - asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); - return c; -} - - -__device__ __forceinline__ float warp_reduce_max(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); - return val; -} - - -__device__ __forceinline__ float warp_reduce_sum(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val += __shfl_xor_sync(0xFFFFFFFF, val, offset); - return val; -} - - -__device__ __forceinline__ float row_max_reduce(float2 acc) { - return max_noftz(acc.x, acc.y); -} - - -__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { - #pragma unroll - for (int j = 0; j < 16; j++) { - if (j % 2 == 0) - acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); - else - acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); - } -} - - -__device__ __forceinline__ float2 ex2_emulation_f32x2_value(float2 value) { - const float c0 = 1.0f, c1 = 0.695146143436431884765625f; - const float c2 = 0.227564394474029541015625f, c3 = 0.077119089663028717041015625f; - const float magic = 12582912.0f; - float x0 = max_noftz(value.x, -127.0f), x1 = max_noftz(value.y, -127.0f); - float2 xc2 = make_float2(x0, x1), magic2 = make_float2(magic, magic); - float2 xr2; - asm("add.rm.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xr2) - : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&magic2)); - float2 c3_2 = make_float2(c3, c3), c2_2 = make_float2(c2, c2); - float2 c1_2 = make_float2(c1, c1), c0_2 = make_float2(c0, c0); - float2 xrb2, xfrac2; - asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xrb2) - : "l"(*(unsigned long long*)&xr2), "l"(*(unsigned long long*)&magic2)); - asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xfrac2) - : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&xrb2)); - float2 poly2; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&c3_2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c2_2)); - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c1_2)); - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) - : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c0_2)); - int x0r_i, x1r_i, p0_i, p1_i; - asm("mov.b64 {%0, %1}, %2;" : "=r"(x0r_i), "=r"(x1r_i) : "l"(*(unsigned long long*)&xr2)); - asm("mov.b64 {%0, %1}, %2;" : "=r"(p0_i), "=r"(p1_i) : "l"(*(unsigned long long*)&poly2)); - float r0, r1; - asm("mov.b32 %0, %1;" : "=f"(r0) : "r"((x0r_i << 23) + p0_i)); - asm("mov.b32 %0, %1;" : "=f"(r1) : "r"((x1r_i << 23) + p1_i)); - return make_float2(r0, r1); -} - -__device__ __forceinline__ void ex2_emulation_f32x2(float* x0_ptr, float* x1_ptr) { - float2 result = ex2_emulation_f32x2_value(make_float2(*x0_ptr, *x1_ptr)); - *x0_ptr = result.x; *x1_ptr = result.y; -} - -__device__ __forceinline__ void softmax_frag_exp2_cast( - float* sv, uint32_t* pv, int use_emu) -{ - #pragma unroll - for (int j = 0; j < 16; j++) { - if (use_emu && j >= 12) - ex2_emulation_f32x2(&sv[j*2], &sv[j*2+1]); - else { - sv[j*2] = approx_exp2(sv[j*2]); - sv[j*2+1] = approx_exp2(sv[j*2+1]); - } - } - #pragma unroll - for (int j = 0; j < 16; j++) { - __nv_bfloat162 bf = __float22bfloat162_rn({sv[j*2], sv[j*2+1]}); - pv[j] = reinterpret_cast(bf); - } -} - - - -__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { - const float2* sv2 = reinterpret_cast(sv); - #pragma unroll - for (int j = 0; j < 16; j++) { - asm("add.f32x2 %0, %1, %2;" - : "+l"(reinterpret_cast(*acc)) - : "l"(reinterpret_cast(*acc)), - "l"(reinterpret_cast(sv2[j]))); - } -} - - -__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { - asm("mul.rn.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { - asm("mul.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { - asm("add.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { - asm("add.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { - asm("sub.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { - asm("sub.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("add.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { - float2 r; - asm("sub.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("sub.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ void fma_scale_x32( - float* sv, const float2* scale2, const float2* neg_max2) -{ - float2* sv_2 = reinterpret_cast(sv); - #pragma unroll - for (int j = 0; j < 16; j++) - fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); -} - -__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("mul.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - - -// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) - -__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("add.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("add.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("add.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("add.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("add.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("add.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("add.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - - -__device__ __forceinline__ void fence_async_shared() { - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); -} - - -__device__ __forceinline__ uint64_t make_smem_desc(int addr) { - const int SBO = 1024; - return desc_encode(addr) - | (desc_encode(SBO) << 32ULL) - | (1ULL << 46ULL) - | (2ULL << 61ULL); -} - - -__device__ __forceinline__ void tma_2d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3}], [%4];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_4d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.4d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5}], [%6];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr) { - // Canonical .shared::cta form for non-multicast gather4, matching - // trtllm-gen / cuda_ptx and the PTX ISA qualifier order - // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, - // .shared::cluster is reserved for the multicast variant (ctaMask). - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { - asm volatile( - "tcgen05.commit.cta_group::1.mbarrier::arrive::one" - ".shared::cluster.b64 [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x16.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7," - " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), - "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), - "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { - tmem_ld_x16(dst, addr); - asm volatile("tcgen05.wait::ld.sync.aligned;"); -} - - -__device__ __forceinline__ void tmem_ld_x8(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x8.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { - uint32_t result; - asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" - : "=r"(result) : "r"(val)); - return result; -} - - -template -struct CakePipelineState { - uint32_t index_; - uint32_t phase_; - uint32_t count_; - - __device__ __forceinline__ uint32_t index() const { return index_; } - __device__ __forceinline__ uint32_t phase() const { return phase_; } - __device__ __forceinline__ uint32_t count() const { return count_; } - __device__ __forceinline__ CakePipelineState& operator++() { - ++count_; - ++index_; - if (index_ == Stages) { - index_ = 0; - phase_ ^= 1; - } - return *this; - } -}; - -template -__device__ __forceinline__ CakePipelineState -cake_pipeline_state_at(CakePipelineState state) { - #pragma unroll - for (int step = 0; step < Offset; ++step) { - ++state; - } - return state; -} - -extern "C" { - -__global__ __launch_bounds__(512, 1) void -kernel_cake_dsv4_f55f69afe4ac3ec96817(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, const __grid_constant__ CUtensorMap tmap_compressed_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ cum_seq_lens_q, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items) -{ - const int tid = threadIdx.x; - const int warp = make_warp_uniform(tid / 32); - const int lane = tid % 32; - - extern __shared__ __align__(1024) char smem_raw[]; - int smem; - smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); - - const int mbar_base = smem + 202784; - #define q_full_addr (mbar_base + 0) - #define q_empty_addr (mbar_base + 16) - #define kv_full_addr (mbar_base + 32) - #define kv_empty_addr (mbar_base + 64) - #define page_full_addr (mbar_base + 96) - #define page_empty_addr (mbar_base + 144) - #define score_full_addr (mbar_base + 192) - #define score_empty_addr (mbar_base + 208) - #define stats_full_addr (mbar_base + 224) - #define stats_empty_addr (mbar_base + 240) - #define o_full_addr (mbar_base + 256) - #define o_empty_addr (mbar_base + 264) - #define work_full_addr (mbar_base + 272) - #define work_empty_addr (mbar_base + 288) - #define throttle_full_addr (mbar_base + 304) - #define throttle_empty_addr (mbar_base + 320) - - const int bid = blockIdx.x; - const int num_bids = gridDim.x; - - const int cta_rank = 0; - - // Kernel setup ops - uint8_t* smem_q = reinterpret_cast(smem_raw + 0); - const int smem_q_addr = smem + 0; - uint8_t* smem_k = reinterpret_cast(smem_raw + 65536); - const int smem_k_addr = smem + 65536; - uint8_t* smem_v = reinterpret_cast(smem_raw + 65536); - const int smem_v_addr = smem + 65536; - int* smem_page_offsets = reinterpret_cast(smem_raw + 196608); - const int smem_page_offsets_addr = smem + 196608; - unsigned int* work_response = reinterpret_cast(smem_raw + 202752); - const int work_response_addr = smem + 202752; - asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_q))) : "memory"); - asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_swa_kv))) : "memory"); - asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_compressed_kv))) : "memory"); - - // Mbarrier init (16 pipeline groups, 0 ordered-sequence groups, 42 barriers) - // Mbarriers at smem_raw[202784..203120) - - if (warp == 12) { - uint32_t leader = elect_sync(); - if (leader) { - // --- pipeline 'q_pipe' --- - // q_full: 2 barriers, init_count=1 - mbarrier_init(smem + 202784, 1); - mbarrier_init(smem + 202792, 1); - // q_empty: 2 barriers, init_count=1 - mbarrier_init(smem + 202800, 1); - mbarrier_init(smem + 202808, 1); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - if (warp == 0) { - uint32_t leader = elect_sync(); - if (leader) { - // --- pipeline 'kv_pipe' --- - // kv_full: 4 barriers, init_count=4 - mbarrier_init(smem + 202816, 4); - mbarrier_init(smem + 202824, 4); - mbarrier_init(smem + 202832, 4); - mbarrier_init(smem + 202840, 4); - // kv_empty: 4 barriers, init_count=1 - mbarrier_init(smem + 202848, 1); - mbarrier_init(smem + 202856, 1); - mbarrier_init(smem + 202864, 1); - mbarrier_init(smem + 202872, 1); - // --- pipeline 'page_pipe' --- - // page_full: 6 barriers, init_count=32 - mbarrier_init(smem + 202880, 32); - mbarrier_init(smem + 202888, 32); - mbarrier_init(smem + 202896, 32); - mbarrier_init(smem + 202904, 32); - mbarrier_init(smem + 202912, 32); - mbarrier_init(smem + 202920, 32); - // page_empty: 6 barriers, init_count=128 - mbarrier_init(smem + 202928, 128); - mbarrier_init(smem + 202936, 128); - mbarrier_init(smem + 202944, 128); - mbarrier_init(smem + 202952, 128); - mbarrier_init(smem + 202960, 128); - mbarrier_init(smem + 202968, 128); - // --- pipeline 'score_pipe' --- - // score_full: 2 barriers, init_count=1 - mbarrier_init(smem + 202976, 1); - mbarrier_init(smem + 202984, 1); - // score_empty: 2 barriers, init_count=128 - mbarrier_init(smem + 202992, 128); - mbarrier_init(smem + 203000, 128); - // --- pipeline 'stats_pipe' --- - // stats_full: 2 barriers, init_count=128 - mbarrier_init(smem + 203008, 128); - mbarrier_init(smem + 203016, 128); - // stats_empty: 2 barriers, init_count=128 - mbarrier_init(smem + 203024, 128); - mbarrier_init(smem + 203032, 128); - // o_full: 1 barriers, init_count=1 - mbarrier_init(smem + 203040, 1); - // o_empty: 1 barriers, init_count=128 - mbarrier_init(smem + 203048, 128); - // --- pipeline 'work_pipe' --- - // work_full: 2 barriers, init_count=1 - mbarrier_init(smem + 203056, 1); - mbarrier_init(smem + 203064, 1); - // work_empty: 2 barriers, init_count=512 - mbarrier_init(smem + 203072, 512); - mbarrier_init(smem + 203080, 512); - // --- pipeline 'throttle_pipe' --- - // throttle_full: 2 barriers, init_count=128 - mbarrier_init(smem + 203088, 128); - mbarrier_init(smem + 203096, 128); - // throttle_empty: 2 barriers, init_count=32 - mbarrier_init(smem + 203104, 32); - mbarrier_init(smem + 203112, 32); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - - __syncwarp(); - - // TMEM alloc (512 columns, 512 used) - volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 203120); - if (warp == 0) { - int _tmem_hold = smem + 203120; - asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); - __syncwarp(); - asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); - } - - __syncthreads(); - asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); - - const int taddr = tmem_addr_storage[0]; - - // Kernel post-init ops - const int tmem_tmem = taddr; - - // ---- Ordered hardware-WG register redistribution ---- - // Inc phase consumes the registers released above. - if (warp >= 8 && warp <= 11) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 136;"); - } - - // ---- Role: softmax ---- - if (warp <= 3) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 152;"); - { // softmax_main - const int warp_rank = warp; - const int logical_row_origin = warp_rank * 16; - const int my_row = logical_row_origin + lane % 16; - const int col_half = lane / 16; - int score_cursor = 0; - int stats_cursor = 0; - unsigned int softmax_work_stage = 0; - unsigned int softmax_q = blockIdx.x; - unsigned int softmax_head = blockIdx.y; - unsigned int softmax_batch = blockIdx.z; - unsigned int _phase_work_full = 0; - #pragma unroll 1 - for (unsigned int _work = 0; _work < total_work_items; _work++) { - int q_begin = cum_seq_lens_q[(int)softmax_batch]; - int q_end = cum_seq_lens_q[(int)(softmax_batch + 1)]; - int q_local = (int)softmax_q; - int work_valid = q_local < q_end - q_begin; - asm volatile("griddepcontrol.wait;" ::: "memory"); - float softmax_rescale_log2 = bmm1_scale[0] * 1.4426950408889634f; - asm volatile("griddepcontrol.wait;" ::: "memory"); - float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; - if (work_valid != 0) { - int query_idx = q_begin + q_local; - int _max_4 = ((sparse_topk_lens[query_idx] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx] + sparse_topk_lens_offset) : (0)); - int _min_3 = ((_max_4) < (sparse_topk) ? (_max_4) : (sparse_topk)); - int active_topk = _min_3; - int num_tiles = (active_topk + 128 - 1) / 128; - float row_max = -CAKE_INF; - float row_sum = 0.0f; - #pragma unroll 1 - for (int tile = 0; tile < num_tiles; tile++) { - int pipeline_tile = score_cursor + tile; - int score_stage = pipeline_tile & 1; - int score_phase = pipeline_tile >> 1 & 1; - mbarrier_wait(score_full_addr + (score_stage) * 8, score_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - int score_col = ((score_stage != 0) ? 128 : 0); - float _tmem_load_0[64]; - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_0[0])), "=r"(*reinterpret_cast(&_tmem_load_0[1])), "=r"(*reinterpret_cast(&_tmem_load_0[2])), "=r"(*reinterpret_cast(&_tmem_load_0[3])), "=r"(*reinterpret_cast(&_tmem_load_0[4])), "=r"(*reinterpret_cast(&_tmem_load_0[5])), "=r"(*reinterpret_cast(&_tmem_load_0[6])), "=r"(*reinterpret_cast(&_tmem_load_0[7])), "=r"(*reinterpret_cast(&_tmem_load_0[8])), "=r"(*reinterpret_cast(&_tmem_load_0[9])), "=r"(*reinterpret_cast(&_tmem_load_0[10])), "=r"(*reinterpret_cast(&_tmem_load_0[11])), "=r"(*reinterpret_cast(&_tmem_load_0[12])), "=r"(*reinterpret_cast(&_tmem_load_0[13])), "=r"(*reinterpret_cast(&_tmem_load_0[14])), "=r"(*reinterpret_cast(&_tmem_load_0[15])), "=r"(*reinterpret_cast(&_tmem_load_0[16])), "=r"(*reinterpret_cast(&_tmem_load_0[17])), "=r"(*reinterpret_cast(&_tmem_load_0[18])), "=r"(*reinterpret_cast(&_tmem_load_0[19])), "=r"(*reinterpret_cast(&_tmem_load_0[20])), "=r"(*reinterpret_cast(&_tmem_load_0[21])), "=r"(*reinterpret_cast(&_tmem_load_0[22])), "=r"(*reinterpret_cast(&_tmem_load_0[23])), "=r"(*reinterpret_cast(&_tmem_load_0[24])), "=r"(*reinterpret_cast(&_tmem_load_0[25])), "=r"(*reinterpret_cast(&_tmem_load_0[26])), "=r"(*reinterpret_cast(&_tmem_load_0[27])), "=r"(*reinterpret_cast(&_tmem_load_0[28])), "=r"(*reinterpret_cast(&_tmem_load_0[29])), "=r"(*reinterpret_cast(&_tmem_load_0[30])), "=r"(*reinterpret_cast(&_tmem_load_0[31])) - : "r"(taddr + (unsigned int)score_col)); - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_0[32])), "=r"(*reinterpret_cast(&_tmem_load_0[33])), "=r"(*reinterpret_cast(&_tmem_load_0[34])), "=r"(*reinterpret_cast(&_tmem_load_0[35])), "=r"(*reinterpret_cast(&_tmem_load_0[36])), "=r"(*reinterpret_cast(&_tmem_load_0[37])), "=r"(*reinterpret_cast(&_tmem_load_0[38])), "=r"(*reinterpret_cast(&_tmem_load_0[39])), "=r"(*reinterpret_cast(&_tmem_load_0[40])), "=r"(*reinterpret_cast(&_tmem_load_0[41])), "=r"(*reinterpret_cast(&_tmem_load_0[42])), "=r"(*reinterpret_cast(&_tmem_load_0[43])), "=r"(*reinterpret_cast(&_tmem_load_0[44])), "=r"(*reinterpret_cast(&_tmem_load_0[45])), "=r"(*reinterpret_cast(&_tmem_load_0[46])), "=r"(*reinterpret_cast(&_tmem_load_0[47])), "=r"(*reinterpret_cast(&_tmem_load_0[48])), "=r"(*reinterpret_cast(&_tmem_load_0[49])), "=r"(*reinterpret_cast(&_tmem_load_0[50])), "=r"(*reinterpret_cast(&_tmem_load_0[51])), "=r"(*reinterpret_cast(&_tmem_load_0[52])), "=r"(*reinterpret_cast(&_tmem_load_0[53])), "=r"(*reinterpret_cast(&_tmem_load_0[54])), "=r"(*reinterpret_cast(&_tmem_load_0[55])), "=r"(*reinterpret_cast(&_tmem_load_0[56])), "=r"(*reinterpret_cast(&_tmem_load_0[57])), "=r"(*reinterpret_cast(&_tmem_load_0[58])), "=r"(*reinterpret_cast(&_tmem_load_0[59])), "=r"(*reinterpret_cast(&_tmem_load_0[60])), "=r"(*reinterpret_cast(&_tmem_load_0[61])), "=r"(*reinterpret_cast(&_tmem_load_0[62])), "=r"(*reinterpret_cast(&_tmem_load_0[63])) - : "r"(taddr + (unsigned int)score_col + 32)); - int valid_cols = active_topk - tile * 128 - col_half * 64; - if (valid_cols < 0) { - valid_cols = 0; - } - if (valid_cols > 64) { - valid_cols = 64; - } - uint32_t _slice_lo_mask_0; - { - int _lim_0 = valid_cols; - if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } - else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); - } - } - if (!(_slice_lo_mask_0 & (1u << 0))) _tmem_load_0[0] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 1))) _tmem_load_0[1] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 2))) _tmem_load_0[2] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 3))) _tmem_load_0[3] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 4))) _tmem_load_0[4] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 5))) _tmem_load_0[5] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 6))) _tmem_load_0[6] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 7))) _tmem_load_0[7] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 8))) _tmem_load_0[8] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 9))) _tmem_load_0[9] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 10))) _tmem_load_0[10] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 11))) _tmem_load_0[11] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 12))) _tmem_load_0[12] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 13))) _tmem_load_0[13] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 14))) _tmem_load_0[14] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 15))) _tmem_load_0[15] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 16))) _tmem_load_0[16] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 17))) _tmem_load_0[17] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 18))) _tmem_load_0[18] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 19))) _tmem_load_0[19] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 20))) _tmem_load_0[20] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 21))) _tmem_load_0[21] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 22))) _tmem_load_0[22] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 23))) _tmem_load_0[23] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 24))) _tmem_load_0[24] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 25))) _tmem_load_0[25] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 26))) _tmem_load_0[26] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 27))) _tmem_load_0[27] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 28))) _tmem_load_0[28] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 29))) _tmem_load_0[29] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 30))) _tmem_load_0[30] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 31))) _tmem_load_0[31] = -CAKE_INF; - uint32_t _slice_lo_mask_1; - { - int _lim_1 = valid_cols - 32; - if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } - else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); - } - } - if (!(_slice_lo_mask_1 & (1u << 0))) _tmem_load_0[32] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 1))) _tmem_load_0[33] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 2))) _tmem_load_0[34] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 3))) _tmem_load_0[35] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 4))) _tmem_load_0[36] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 5))) _tmem_load_0[37] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 6))) _tmem_load_0[38] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 7))) _tmem_load_0[39] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 8))) _tmem_load_0[40] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 9))) _tmem_load_0[41] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 10))) _tmem_load_0[42] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 11))) _tmem_load_0[43] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 12))) _tmem_load_0[44] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 13))) _tmem_load_0[45] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 14))) _tmem_load_0[46] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 15))) _tmem_load_0[47] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 16))) _tmem_load_0[48] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 17))) _tmem_load_0[49] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 18))) _tmem_load_0[50] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 19))) _tmem_load_0[51] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 20))) _tmem_load_0[52] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 21))) _tmem_load_0[53] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 22))) _tmem_load_0[54] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 23))) _tmem_load_0[55] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 24))) _tmem_load_0[56] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 25))) _tmem_load_0[57] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 26))) _tmem_load_0[58] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 27))) _tmem_load_0[59] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 28))) _tmem_load_0[60] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 29))) _tmem_load_0[61] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 30))) _tmem_load_0[62] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 31))) _tmem_load_0[63] = -CAKE_INF; - float2 _reg_reduce_max2_2 = {-CAKE_INF, -CAKE_INF}; - row_max_x32_accum(&_tmem_load_0[0], _reg_reduce_max2_2); - row_max_x32_accum(&_tmem_load_0[32], _reg_reduce_max2_2); - float _tmem_load_0_max = row_max_reduce(_reg_reduce_max2_2); - float _max_5 = max_noftz(row_max, _tmem_load_0_max); - float new_max = _max_5; - float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, new_max, 16); - float _max_6 = max_noftz(new_max, _shfl_xor_0); - new_max = _max_6; - float _fma_0 = __fmaf_rn(row_max, softmax_rescale_log2, (-new_max) * softmax_rescale_log2); - float delta = _fma_0; - float _exp2_0 = approx_exp2(delta); - float acc_scale = ((row_max > -CAKE_INF) ? _exp2_0 : 1.0f); - int stats_tile = stats_cursor + tile; - int stats_stage = stats_tile & 1; - int stats_phase = stats_tile >> 1 & 1; - mbarrier_wait(stats_empty_addr + (stats_stage) * 8, 1 - stats_phase); - float stats_pair[2]; - stats_pair[0] = row_max; - stats_pair[1] = new_max; - asm volatile( - "tcgen05.st.sync.aligned.32x32b.x2.b32" - " [%0], {%1, %2};" - :: "r"(taddr + (unsigned int)(stats_stage * 128)), "f"(stats_pair[0]), "f"(stats_pair[1])); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(stats_full_addr + (stats_stage) * 8); - float safe_max = ((new_max == -CAKE_INF) ? 0.0f : new_max); - float max_scaled = safe_max * softmax_scale_log2; - const float2 _fma_b2_3 = {softmax_scale_log2, softmax_scale_log2}; - const float2 _fma_c2_4 = {-max_scaled + 8.8073549f, -max_scaled + 8.8073549f}; - #pragma unroll - for (int _lf = 0; _lf < 32; _lf++) - fma_f32x2_inplace(&reinterpret_cast(_tmem_load_0)[_lf], _fma_b2_3, _fma_c2_4); - #pragma unroll - for (int _le = 0; _le < 64; _le++) { - _tmem_load_0[_le] = approx_exp2(_tmem_load_0[_le]); - } - float2 _reg_reduce_sum2_5 = make_float2(0.0f, 0.0f); - softmax_block_sum(&_tmem_load_0[0], &_reg_reduce_sum2_5); - softmax_block_sum(&_tmem_load_0[32], &_reg_reduce_sum2_5); - float _tmem_load_0_sum = _reg_reduce_sum2_5.x + _reg_reduce_sum2_5.y; - float block_sum = _tmem_load_0_sum; - float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, block_sum, 16); - block_sum = block_sum + _shfl_xor_1; - float _fma_1 = __fmaf_rn(row_sum, acc_scale, block_sum); - row_sum = _fma_1; - uint32_t _fp8_0[16]; - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_0[0]), "f"(_tmem_load_0[1]), - "f"(_tmem_load_0[2]), "f"(_tmem_load_0[3])); - _fp8_0[0] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_0[4]), "f"(_tmem_load_0[5]), - "f"(_tmem_load_0[6]), "f"(_tmem_load_0[7])); - _fp8_0[1] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_0[8]), "f"(_tmem_load_0[9]), - "f"(_tmem_load_0[10]), "f"(_tmem_load_0[11])); - _fp8_0[2] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_0[12]), "f"(_tmem_load_0[13]), - "f"(_tmem_load_0[14]), "f"(_tmem_load_0[15])); - _fp8_0[3] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_0[16]), "f"(_tmem_load_0[17]), - "f"(_tmem_load_0[18]), "f"(_tmem_load_0[19])); - _fp8_0[4] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_0[20]), "f"(_tmem_load_0[21]), - "f"(_tmem_load_0[22]), "f"(_tmem_load_0[23])); - _fp8_0[5] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_0[24]), "f"(_tmem_load_0[25]), - "f"(_tmem_load_0[26]), "f"(_tmem_load_0[27])); - _fp8_0[6] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_0[28]), "f"(_tmem_load_0[29]), - "f"(_tmem_load_0[30]), "f"(_tmem_load_0[31])); - _fp8_0[7] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_0[32]), "f"(_tmem_load_0[33]), - "f"(_tmem_load_0[34]), "f"(_tmem_load_0[35])); - _fp8_0[8] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_0[36]), "f"(_tmem_load_0[37]), - "f"(_tmem_load_0[38]), "f"(_tmem_load_0[39])); - _fp8_0[9] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_0[40]), "f"(_tmem_load_0[41]), - "f"(_tmem_load_0[42]), "f"(_tmem_load_0[43])); - _fp8_0[10] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_0[44]), "f"(_tmem_load_0[45]), - "f"(_tmem_load_0[46]), "f"(_tmem_load_0[47])); - _fp8_0[11] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_0[48]), "f"(_tmem_load_0[49]), - "f"(_tmem_load_0[50]), "f"(_tmem_load_0[51])); - _fp8_0[12] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_0[52]), "f"(_tmem_load_0[53]), - "f"(_tmem_load_0[54]), "f"(_tmem_load_0[55])); - _fp8_0[13] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_0[56]), "f"(_tmem_load_0[57]), - "f"(_tmem_load_0[58]), "f"(_tmem_load_0[59])); - _fp8_0[14] = _packed; - } - { - uint32_t _packed; - asm volatile("{\n\t" - ".reg .b16 _lo;\n\t" - ".reg .b16 _hi;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" - "mov.b32 %0, {_lo, _hi};\n\t" - "}" - : "=r"(_packed) : "f"(_tmem_load_0[60]), "f"(_tmem_load_0[61]), - "f"(_tmem_load_0[62]), "f"(_tmem_load_0[63])); - _fp8_0[15] = _packed; - } - int p_col = ((score_stage != 0) ? 160 : 32); - asm volatile( - "tcgen05.st.sync.aligned.16x32bx2.x16.b32" - " [%0], 16, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" - :: "r"(taddr + (unsigned int)p_col), "r"(*reinterpret_cast(&_fp8_0[0])), "r"(*reinterpret_cast(&_fp8_0[1])), "r"(*reinterpret_cast(&_fp8_0[2])), "r"(*reinterpret_cast(&_fp8_0[3])), "r"(*reinterpret_cast(&_fp8_0[4])), "r"(*reinterpret_cast(&_fp8_0[5])), "r"(*reinterpret_cast(&_fp8_0[6])), "r"(*reinterpret_cast(&_fp8_0[7])), "r"(*reinterpret_cast(&_fp8_0[8])), "r"(*reinterpret_cast(&_fp8_0[9])), "r"(*reinterpret_cast(&_fp8_0[10])), "r"(*reinterpret_cast(&_fp8_0[11])), "r"(*reinterpret_cast(&_fp8_0[12])), "r"(*reinterpret_cast(&_fp8_0[13])), "r"(*reinterpret_cast(&_fp8_0[14])), "r"(*reinterpret_cast(&_fp8_0[15]))); - asm volatile( - "tcgen05.st.sync.aligned.16x32bx2.x16.b32" - " [%0], 16, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" - :: "r"(taddr + 1048576 + (unsigned int)p_col), "r"(*reinterpret_cast(&_fp8_0[0])), "r"(*reinterpret_cast(&_fp8_0[1])), "r"(*reinterpret_cast(&_fp8_0[2])), "r"(*reinterpret_cast(&_fp8_0[3])), "r"(*reinterpret_cast(&_fp8_0[4])), "r"(*reinterpret_cast(&_fp8_0[5])), "r"(*reinterpret_cast(&_fp8_0[6])), "r"(*reinterpret_cast(&_fp8_0[7])), "r"(*reinterpret_cast(&_fp8_0[8])), "r"(*reinterpret_cast(&_fp8_0[9])), "r"(*reinterpret_cast(&_fp8_0[10])), "r"(*reinterpret_cast(&_fp8_0[11])), "r"(*reinterpret_cast(&_fp8_0[12])), "r"(*reinterpret_cast(&_fp8_0[13])), "r"(*reinterpret_cast(&_fp8_0[14])), "r"(*reinterpret_cast(&_fp8_0[15]))); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - if (num_tiles > tile + 1) { - mbarrier_arrive(score_empty_addr + (score_stage) * 8); - } - row_max = new_max; - } - int final_stats_tile = stats_cursor + num_tiles; - int final_stats_stage = final_stats_tile & 1; - int final_stats_phase = final_stats_tile >> 1 & 1; - mbarrier_wait(stats_empty_addr + (final_stats_stage) * 8, 1 - final_stats_phase); - float final_stats[2]; - final_stats[0] = row_sum; - final_stats[1] = row_max; - asm volatile( - "tcgen05.st.sync.aligned.32x32b.x2.b32" - " [%0], {%1, %2};" - :: "r"(taddr + (unsigned int)(final_stats_stage * 128)), "f"(final_stats[0]), "f"(final_stats[1])); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(stats_full_addr + (final_stats_stage) * 8); - int tail_stats_tile = final_stats_tile + 1; - int tail_stats_stage = tail_stats_tile & 1; - int tail_stats_phase = tail_stats_tile >> 1 & 1; - mbarrier_wait(stats_empty_addr + (tail_stats_stage) * 8, 1 - tail_stats_phase); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(stats_full_addr + (tail_stats_stage) * 8); - int last_score_tile = score_cursor + num_tiles - 1; - mbarrier_arrive(score_empty_addr + (last_score_tile & 1) * 8); - stats_cursor = stats_cursor + num_tiles + 2; - score_cursor = score_cursor + num_tiles; - } - mbarrier_wait(work_full_addr + (softmax_work_stage) * 8, _phase_work_full); - uint32_t _clc_valid_3 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_3) - : "r"(work_response_addr + softmax_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_9 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_9) - : "r"(work_response_addr + softmax_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_10 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_10) - : "r"(work_response_addr + softmax_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_11 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_11) - : "r"(work_response_addr + softmax_work_stage * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(work_empty_addr + (softmax_work_stage) * 8); - softmax_work_stage += 1; - if (softmax_work_stage == 2) { softmax_work_stage = 0; _phase_work_full ^= 1; } - if (_clc_valid_3 == 0) { - break; - } - softmax_q = _clc_ctaid_9; - softmax_head = _clc_ctaid_10; - softmax_batch = _clc_ctaid_11; - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - } - } - // ---- Role: correction ---- - if (warp >= 4 && warp <= 7) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 144;"); - { // correction_main - const int warp_rank_1 = warp - 4; - const int logical_row_origin_1 = warp_rank_1 * 16; - const int my_row_1 = logical_row_origin_1 + lane % 16; - const int col_half_1 = lane / 16; - int score_cursor_1 = 0; - int stats_cursor_1 = 0; - unsigned int correction_work_stage = 0; - unsigned int correction_q = blockIdx.x; - unsigned int correction_head = blockIdx.y; - unsigned int correction_batch = blockIdx.z; - unsigned int _phase_work_full_1 = 0; - #pragma unroll 1 - for (unsigned int _work_1 = 0; _work_1 < total_work_items; _work_1++) { - int q_begin_1 = cum_seq_lens_q[(int)correction_batch]; - int q_end_1 = cum_seq_lens_q[(int)(correction_batch + 1)]; - int q_local_1 = (int)correction_q; - int work_valid_1 = q_local_1 < q_end_1 - q_begin_1; - asm volatile("griddepcontrol.wait;" ::: "memory"); - float softmax_scale_log2_1 = bmm1_scale[0] * 1.4426950408889634f; - float output_scale = bmm2_scale[0]; - if (work_valid_1 != 0) { - int query_idx_1 = q_begin_1 + q_local_1; - int head_base = (int)correction_head * 64; - int head_idx = head_base + my_row_1; - int _max_7 = ((sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_1] + sparse_topk_lens_offset) : (0)); - int _min_4 = ((_max_7) < (sparse_topk) ? (_max_7) : (sparse_topk)); - int active_topk_1 = _min_4; - int num_tiles_1 = (active_topk_1 + 128 - 1) / 128; - int first_stats_tile = stats_cursor_1; - int first_stats_stage = first_stats_tile & 1; - int first_stats_phase = first_stats_tile >> 1 & 1; - mbarrier_wait(stats_full_addr + (first_stats_stage) * 8, first_stats_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - mbarrier_arrive(stats_empty_addr + (first_stats_stage) * 8); - #pragma unroll 1 - for (int tile_1 = 1; tile_1 < num_tiles_1; tile_1++) { - int stats_tile_1 = stats_cursor_1 + tile_1; - int stats_stage_1 = stats_tile_1 & 1; - int stats_phase_1 = stats_tile_1 >> 1 & 1; - mbarrier_wait(stats_full_addr + (stats_stage_1) * 8, stats_phase_1); - asm volatile("tcgen05.fence::after_thread_sync;"); - float _tmem_load_1[2]; - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x2.b32" - " {%0, %1}, [%2];" - : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]) - : "r"(taddr + (unsigned int)(stats_stage_1 * 128))); - asm volatile("tcgen05.wait::ld.sync.aligned;"); - mbarrier_arrive(stats_empty_addr + (stats_stage_1) * 8); - int output_tile = score_cursor_1 + tile_1 - 1; - mbarrier_wait(o_full_addr, output_tile & 1); - asm volatile("tcgen05.fence::after_thread_sync;"); - float max_delta = _tmem_load_1[0] - _tmem_load_1[1]; - float _exp2_1 = approx_exp2(max_delta * softmax_scale_log2_1); - float acc_scale_1 = ((max_delta != 0.0f) ? _exp2_1 : 1.0f); - int _vote_0 = __any_sync(0xFFFFFFFF, acc_scale_1 < 1.0f); - int any_rescale = _vote_0; - if (any_rescale != 0) { - #pragma unroll - for (int output_half = 0; output_half < 2; output_half++) { - int row_bank = ((output_half != 0) ? 1048576 : 0); - #pragma unroll - for (int col = 0; col < 128; col += 64) { - float _tmem_load_2[64]; - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 128;" - : "=r"(*reinterpret_cast(&_tmem_load_2[0])), "=r"(*reinterpret_cast(&_tmem_load_2[1])), "=r"(*reinterpret_cast(&_tmem_load_2[2])), "=r"(*reinterpret_cast(&_tmem_load_2[3])), "=r"(*reinterpret_cast(&_tmem_load_2[4])), "=r"(*reinterpret_cast(&_tmem_load_2[5])), "=r"(*reinterpret_cast(&_tmem_load_2[6])), "=r"(*reinterpret_cast(&_tmem_load_2[7])), "=r"(*reinterpret_cast(&_tmem_load_2[8])), "=r"(*reinterpret_cast(&_tmem_load_2[9])), "=r"(*reinterpret_cast(&_tmem_load_2[10])), "=r"(*reinterpret_cast(&_tmem_load_2[11])), "=r"(*reinterpret_cast(&_tmem_load_2[12])), "=r"(*reinterpret_cast(&_tmem_load_2[13])), "=r"(*reinterpret_cast(&_tmem_load_2[14])), "=r"(*reinterpret_cast(&_tmem_load_2[15])), "=r"(*reinterpret_cast(&_tmem_load_2[16])), "=r"(*reinterpret_cast(&_tmem_load_2[17])), "=r"(*reinterpret_cast(&_tmem_load_2[18])), "=r"(*reinterpret_cast(&_tmem_load_2[19])), "=r"(*reinterpret_cast(&_tmem_load_2[20])), "=r"(*reinterpret_cast(&_tmem_load_2[21])), "=r"(*reinterpret_cast(&_tmem_load_2[22])), "=r"(*reinterpret_cast(&_tmem_load_2[23])), "=r"(*reinterpret_cast(&_tmem_load_2[24])), "=r"(*reinterpret_cast(&_tmem_load_2[25])), "=r"(*reinterpret_cast(&_tmem_load_2[26])), "=r"(*reinterpret_cast(&_tmem_load_2[27])), "=r"(*reinterpret_cast(&_tmem_load_2[28])), "=r"(*reinterpret_cast(&_tmem_load_2[29])), "=r"(*reinterpret_cast(&_tmem_load_2[30])), "=r"(*reinterpret_cast(&_tmem_load_2[31])) - : "r"(taddr + 256 + (unsigned int)row_bank + (unsigned int)col)); - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 128;" - : "=r"(*reinterpret_cast(&_tmem_load_2[32])), "=r"(*reinterpret_cast(&_tmem_load_2[33])), "=r"(*reinterpret_cast(&_tmem_load_2[34])), "=r"(*reinterpret_cast(&_tmem_load_2[35])), "=r"(*reinterpret_cast(&_tmem_load_2[36])), "=r"(*reinterpret_cast(&_tmem_load_2[37])), "=r"(*reinterpret_cast(&_tmem_load_2[38])), "=r"(*reinterpret_cast(&_tmem_load_2[39])), "=r"(*reinterpret_cast(&_tmem_load_2[40])), "=r"(*reinterpret_cast(&_tmem_load_2[41])), "=r"(*reinterpret_cast(&_tmem_load_2[42])), "=r"(*reinterpret_cast(&_tmem_load_2[43])), "=r"(*reinterpret_cast(&_tmem_load_2[44])), "=r"(*reinterpret_cast(&_tmem_load_2[45])), "=r"(*reinterpret_cast(&_tmem_load_2[46])), "=r"(*reinterpret_cast(&_tmem_load_2[47])), "=r"(*reinterpret_cast(&_tmem_load_2[48])), "=r"(*reinterpret_cast(&_tmem_load_2[49])), "=r"(*reinterpret_cast(&_tmem_load_2[50])), "=r"(*reinterpret_cast(&_tmem_load_2[51])), "=r"(*reinterpret_cast(&_tmem_load_2[52])), "=r"(*reinterpret_cast(&_tmem_load_2[53])), "=r"(*reinterpret_cast(&_tmem_load_2[54])), "=r"(*reinterpret_cast(&_tmem_load_2[55])), "=r"(*reinterpret_cast(&_tmem_load_2[56])), "=r"(*reinterpret_cast(&_tmem_load_2[57])), "=r"(*reinterpret_cast(&_tmem_load_2[58])), "=r"(*reinterpret_cast(&_tmem_load_2[59])), "=r"(*reinterpret_cast(&_tmem_load_2[60])), "=r"(*reinterpret_cast(&_tmem_load_2[61])), "=r"(*reinterpret_cast(&_tmem_load_2[62])), "=r"(*reinterpret_cast(&_tmem_load_2[63])) - : "r"(taddr + 256 + (unsigned int)row_bank + (unsigned int)col + 32)); - const float2 _scale2_0 = {acc_scale_1, acc_scale_1}; - #pragma unroll - for (int _ls = 0; _ls < 32; _ls++) - mul_f32x2_inplace(&reinterpret_cast(_tmem_load_2)[_ls], _scale2_0); - asm volatile( - "tcgen05.st.sync.aligned.16x32bx2.x64.b32" - " [%0], 128, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64};" - :: "r"(taddr + 256 + (unsigned int)row_bank + (unsigned int)col), "r"(*reinterpret_cast(&_tmem_load_2[0])), "r"(*reinterpret_cast(&_tmem_load_2[1])), "r"(*reinterpret_cast(&_tmem_load_2[2])), "r"(*reinterpret_cast(&_tmem_load_2[3])), "r"(*reinterpret_cast(&_tmem_load_2[4])), "r"(*reinterpret_cast(&_tmem_load_2[5])), "r"(*reinterpret_cast(&_tmem_load_2[6])), "r"(*reinterpret_cast(&_tmem_load_2[7])), "r"(*reinterpret_cast(&_tmem_load_2[8])), "r"(*reinterpret_cast(&_tmem_load_2[9])), "r"(*reinterpret_cast(&_tmem_load_2[10])), "r"(*reinterpret_cast(&_tmem_load_2[11])), "r"(*reinterpret_cast(&_tmem_load_2[12])), "r"(*reinterpret_cast(&_tmem_load_2[13])), "r"(*reinterpret_cast(&_tmem_load_2[14])), "r"(*reinterpret_cast(&_tmem_load_2[15])), "r"(*reinterpret_cast(&_tmem_load_2[16])), "r"(*reinterpret_cast(&_tmem_load_2[17])), "r"(*reinterpret_cast(&_tmem_load_2[18])), "r"(*reinterpret_cast(&_tmem_load_2[19])), "r"(*reinterpret_cast(&_tmem_load_2[20])), "r"(*reinterpret_cast(&_tmem_load_2[21])), "r"(*reinterpret_cast(&_tmem_load_2[22])), "r"(*reinterpret_cast(&_tmem_load_2[23])), "r"(*reinterpret_cast(&_tmem_load_2[24])), "r"(*reinterpret_cast(&_tmem_load_2[25])), "r"(*reinterpret_cast(&_tmem_load_2[26])), "r"(*reinterpret_cast(&_tmem_load_2[27])), "r"(*reinterpret_cast(&_tmem_load_2[28])), "r"(*reinterpret_cast(&_tmem_load_2[29])), "r"(*reinterpret_cast(&_tmem_load_2[30])), "r"(*reinterpret_cast(&_tmem_load_2[31])), "r"(*reinterpret_cast(&_tmem_load_2[32])), "r"(*reinterpret_cast(&_tmem_load_2[33])), "r"(*reinterpret_cast(&_tmem_load_2[34])), "r"(*reinterpret_cast(&_tmem_load_2[35])), "r"(*reinterpret_cast(&_tmem_load_2[36])), "r"(*reinterpret_cast(&_tmem_load_2[37])), "r"(*reinterpret_cast(&_tmem_load_2[38])), "r"(*reinterpret_cast(&_tmem_load_2[39])), "r"(*reinterpret_cast(&_tmem_load_2[40])), "r"(*reinterpret_cast(&_tmem_load_2[41])), "r"(*reinterpret_cast(&_tmem_load_2[42])), "r"(*reinterpret_cast(&_tmem_load_2[43])), "r"(*reinterpret_cast(&_tmem_load_2[44])), "r"(*reinterpret_cast(&_tmem_load_2[45])), "r"(*reinterpret_cast(&_tmem_load_2[46])), "r"(*reinterpret_cast(&_tmem_load_2[47])), "r"(*reinterpret_cast(&_tmem_load_2[48])), "r"(*reinterpret_cast(&_tmem_load_2[49])), "r"(*reinterpret_cast(&_tmem_load_2[50])), "r"(*reinterpret_cast(&_tmem_load_2[51])), "r"(*reinterpret_cast(&_tmem_load_2[52])), "r"(*reinterpret_cast(&_tmem_load_2[53])), "r"(*reinterpret_cast(&_tmem_load_2[54])), "r"(*reinterpret_cast(&_tmem_load_2[55])), "r"(*reinterpret_cast(&_tmem_load_2[56])), "r"(*reinterpret_cast(&_tmem_load_2[57])), "r"(*reinterpret_cast(&_tmem_load_2[58])), "r"(*reinterpret_cast(&_tmem_load_2[59])), "r"(*reinterpret_cast(&_tmem_load_2[60])), "r"(*reinterpret_cast(&_tmem_load_2[61])), "r"(*reinterpret_cast(&_tmem_load_2[62])), "r"(*reinterpret_cast(&_tmem_load_2[63]))); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - } - } - } - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(o_empty_addr); - } - int final_stats_tile_1 = stats_cursor_1 + num_tiles_1; - int final_stats_stage_1 = final_stats_tile_1 & 1; - int final_stats_phase_1 = final_stats_tile_1 >> 1 & 1; - mbarrier_wait(stats_full_addr + (final_stats_stage_1) * 8, final_stats_phase_1); - asm volatile("tcgen05.fence::after_thread_sync;"); - float _tmem_load_3[2]; - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x2.b32" - " {%0, %1}, [%2];" - : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]) - : "r"(taddr + (unsigned int)(final_stats_stage_1 * 128))); - asm volatile("tcgen05.wait::ld.sync.aligned;"); - mbarrier_arrive(stats_empty_addr + (final_stats_stage_1) * 8); - int tail_stats_tile_1 = final_stats_tile_1 + 1; - int tail_stats_stage_1 = tail_stats_tile_1 & 1; - int tail_stats_phase_1 = tail_stats_tile_1 >> 1 & 1; - mbarrier_wait(stats_full_addr + (tail_stats_stage_1) * 8, tail_stats_phase_1); - asm volatile("tcgen05.fence::after_thread_sync;"); - mbarrier_arrive(stats_empty_addr + (tail_stats_stage_1) * 8); - int last_output_tile = score_cursor_1 + num_tiles_1 - 1; - mbarrier_wait(o_full_addr, last_output_tile & 1); - asm volatile("tcgen05.fence::after_thread_sync;"); - float total_sum = _tmem_load_3[0]; - float final_max = _tmem_load_3[1]; - int output_in_bounds = head_idx < num_heads; - if (((int)(has_sinks != 0) & output_in_bounds) != 0) { - float _exp2_2 = approx_exp2(sinks[head_idx] * 1.4426950408889634f - final_max * softmax_scale_log2_1); - total_sum = total_sum + _exp2_2 * 448.0f; - } - float _rcp_0 = approx_rcp(total_sum); - float inv_sum = ((total_sum > 0.0f) ? output_scale * _rcp_0 : 0.0f); - int output_base = (query_idx_1 * num_heads + head_idx) * 512; - #pragma unroll - for (int output_half_1 = 0; output_half_1 < 2; output_half_1++) { - int row_bank_1 = ((output_half_1 != 0) ? 1048576 : 0); - int col_base = col_half_1 * 128; - #pragma unroll 2 - for (int col_1 = 0; col_1 < 128; col_1 += 8) { - float _tmem_load_4[8]; - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x8.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7}, [%8], 128;" - : "=r"(*reinterpret_cast(&_tmem_load_4[0])), "=r"(*reinterpret_cast(&_tmem_load_4[1])), "=r"(*reinterpret_cast(&_tmem_load_4[2])), "=r"(*reinterpret_cast(&_tmem_load_4[3])), "=r"(*reinterpret_cast(&_tmem_load_4[4])), "=r"(*reinterpret_cast(&_tmem_load_4[5])), "=r"(*reinterpret_cast(&_tmem_load_4[6])), "=r"(*reinterpret_cast(&_tmem_load_4[7])) - : "r"(taddr + 256 + (unsigned int)row_bank_1 + (unsigned int)col_1)); - const float2 _scale2_1 = {inv_sum, inv_sum}; - #pragma unroll - for (int _ls = 0; _ls < 4; _ls++) - mul_f32x2_inplace(&reinterpret_cast(_tmem_load_4)[_ls], _scale2_1); - uint32_t _tmem_load_4_bf16[4]; - #pragma unroll - for (int _lp = 0; _lp < 4; _lp++) { - __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_4[_lp*2 + 0], _tmem_load_4[_lp*2+1 + 0])); - _tmem_load_4_bf16[_lp] = *(uint32_t*)&_bf2; - } - if (output_in_bounds != 0) { - reinterpret_cast(O + (output_base + output_half_1 * 256 + col_base + col_1))[0] = reinterpret_cast(_tmem_load_4_bf16)[0]; - } - } - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(o_empty_addr); - stats_cursor_1 = stats_cursor_1 + num_tiles_1 + 2; - score_cursor_1 = score_cursor_1 + num_tiles_1; - } - mbarrier_wait(work_full_addr + (correction_work_stage) * 8, _phase_work_full_1); - uint32_t _clc_valid_4 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_4) - : "r"(work_response_addr + correction_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_12 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_12) - : "r"(work_response_addr + correction_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_13 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_13) - : "r"(work_response_addr + correction_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_14 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_14) - : "r"(work_response_addr + correction_work_stage * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(work_empty_addr + (correction_work_stage) * 8); - correction_work_stage += 1; - if (correction_work_stage == 2) { correction_work_stage = 0; _phase_work_full_1 ^= 1; } - if (_clc_valid_4 == 0) { - break; - } - correction_q = _clc_ctaid_12; - correction_head = _clc_ctaid_13; - correction_batch = _clc_ctaid_14; - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile("barrier.sync 7, 128;" ::: "memory"); - if (warp == 4) { - int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); - asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); - } - } - } - // ---- Role: mma ---- - if (warp == 8) { - { // mma_main - CakePipelineState<4> _kv_cons_state_state{0u, 0u, 0u}; - unsigned int q_cons_stage = 0; - int score_prod_stage = 0; - int score_prod_phase = 0; - int p_cons_stage = 0; - int p_cons_phase = 0; - int o_prod_phase = 0; - unsigned int mma_work_stage = 0; - unsigned int mma_q = blockIdx.x; - unsigned int mma_head = blockIdx.y; - unsigned int mma_batch = blockIdx.z; - #pragma unroll - for (int score_lookahead = 0; score_lookahead < 2; score_lookahead++) { - int lookahead_stage = score_lookahead & 1; - int lookahead_phase = score_lookahead >> 1 & 1; - mbarrier_wait(score_empty_addr + (lookahead_stage) * 8, 1 - lookahead_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - } - unsigned int _phase_q_full = 0; - unsigned int _phase_work_full_2 = 0; - #pragma unroll 1 - for (unsigned int _work_2 = 0; _work_2 < total_work_items; _work_2++) { - int q_begin_2 = cum_seq_lens_q[(int)mma_batch]; - int q_end_2 = cum_seq_lens_q[(int)(mma_batch + 1)]; - int q_local_2 = (int)mma_q; - int work_valid_2 = q_local_2 < q_end_2 - q_begin_2; - if (work_valid_2 != 0) { - int query_idx_2 = q_begin_2 + q_local_2; - int _max_3 = ((sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_2] + sparse_topk_lens_offset) : (0)); - int _min_2 = ((_max_3) < (sparse_topk) ? (_max_3) : (sparse_topk)); - int active_topk_2 = _min_2; - int num_tiles_2 = (active_topk_2 + 128 - 1) / 128; - if (num_tiles_2 > 0) { - mbarrier_wait(q_full_addr + (q_cons_stage) * 8, _phase_q_full); - asm volatile("tcgen05.fence::after_thread_sync;"); - int score_col_1 = ((score_prod_stage != 0) ? 128 : 0); - mbarrier_wait(kv_full_addr + (_kv_cons_state_state.index()) * 8, _kv_cons_state_state.phase()); - asm volatile("tcgen05.fence::after_thread_sync;"); - int _mma_a_lo_0 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (q_cons_stage) * 2048); - int _mma_b_lo_0 = make_warp_uniform((((smem_k_addr) >> 4) & 0x3FFF) + (_kv_cons_state_state.index()) * 2048); - { - uint64_t _mma_ss_a_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_0); - uint64_t _mma_ss_b_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_0); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69206032, 0); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 506U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 1018U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69206032, 1); - } - } - elect_commit(kv_empty_addr + (_kv_cons_state_state.index()) * 8); - ++_kv_cons_state_state; - mbarrier_wait(kv_full_addr + (_kv_cons_state_state.index()) * 8, _kv_cons_state_state.phase()); - asm volatile("tcgen05.fence::after_thread_sync;"); - int _mma_a_lo_1 = make_warp_uniform((((smem_q_addr + 16384) >> 4) & 0x3FFF) + (q_cons_stage) * 2048); - int _mma_b_lo_1 = make_warp_uniform((((smem_k_addr) >> 4) & 0x3FFF) + (_kv_cons_state_state.index()) * 2048); - { - uint64_t _mma_ss_a_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_1); - uint64_t _mma_ss_b_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_1); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 506U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 1018U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69206032, 1); - } - } - elect_commit(score_full_addr + (score_prod_stage) * 8); - elect_commit(kv_empty_addr + (_kv_cons_state_state.index()) * 8); - ++_kv_cons_state_state; - score_prod_stage ^= 1; - if (score_prod_stage == 0) { - score_prod_phase ^= 1; - } - int first_pv = 1; - #pragma unroll 1 - for (int _steady = 0; _steady < num_tiles_2 - 1; _steady++) { - score_col_1 = ((score_prod_stage != 0) ? 128 : 0); - mbarrier_wait(kv_full_addr + (_kv_cons_state_state.index()) * 8, _kv_cons_state_state.phase()); - asm volatile("tcgen05.fence::after_thread_sync;"); - int _mma_a_lo_2 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (q_cons_stage) * 2048); - int _mma_b_lo_2 = make_warp_uniform((((smem_k_addr) >> 4) & 0x3FFF) + (_kv_cons_state_state.index()) * 2048); - { - uint64_t _mma_ss_a_desc_2 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_2); - uint64_t _mma_ss_b_desc_2 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_2); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 69206032, 0); - } - incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_2, 506U); - incr_smem_desc_lo(_mma_ss_b_desc_2, 1018U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_2, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_2, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_2, _mma_ss_b_desc_2, 69206032, 1); - } - } - elect_commit(kv_empty_addr + (_kv_cons_state_state.index()) * 8); - ++_kv_cons_state_state; - mbarrier_wait(kv_full_addr + (_kv_cons_state_state.index()) * 8, _kv_cons_state_state.phase()); - asm volatile("tcgen05.fence::after_thread_sync;"); - int _mma_a_lo_3 = make_warp_uniform((((smem_q_addr + 16384) >> 4) & 0x3FFF) + (q_cons_stage) * 2048); - int _mma_b_lo_3 = make_warp_uniform((((smem_k_addr) >> 4) & 0x3FFF) + (_kv_cons_state_state.index()) * 2048); - { - uint64_t _mma_ss_a_desc_3 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_3); - uint64_t _mma_ss_b_desc_3 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_3); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_3, 506U); - incr_smem_desc_lo(_mma_ss_b_desc_3, 1018U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 69206032, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_3, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_3, 2U); - if (elect_sync()) { - tcgen05_mma_f8f6f4((tmem_tmem + (score_col_1)), _mma_ss_a_desc_3, _mma_ss_b_desc_3, 69206032, 1); - } - } - elect_commit(score_full_addr + (score_prod_stage) * 8); - elect_commit(kv_empty_addr + (_kv_cons_state_state.index()) * 8); - ++_kv_cons_state_state; - score_prod_stage ^= 1; - if (score_prod_stage == 0) { - score_prod_phase ^= 1; - } - mbarrier_wait(score_empty_addr + (score_prod_stage) * 8, 1 - score_prod_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - mbarrier_wait(o_empty_addr, 1 - o_prod_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - #pragma unroll - for (int v_half = 0; v_half < 2; v_half++) { - mbarrier_wait(kv_full_addr + (_kv_cons_state_state.index()) * 8, _kv_cons_state_state.phase()); - asm volatile("tcgen05.fence::after_thread_sync;"); - int row_bank_2 = ((v_half != 0) ? 1048576 : 0); - int p_col_1 = ((p_cons_stage != 0) ? 160 : 32) + row_bank_2; - int _mma_b_lo_4 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (_kv_cons_state_state.index()) * 2048); - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 dhi, blo, ta, id;\n\t" - ".reg .b64 db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "" - "mov.b32 dhi, 0x40004040;\n\t" - "mov.b32 id, 71368720;\n\t" - "mov.b32 ta, %2;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p0;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" - "}\n" - :: "r"((tmem_tmem + (256 + row_bank_2))), "r"(_mma_b_lo_4), "r"(tmem_tmem + p_col_1), "r"(((first_pv) ? 0 : 1))); - elect_commit(kv_empty_addr + (_kv_cons_state_state.index()) * 8); - ++_kv_cons_state_state; - } - first_pv = 0; - elect_commit(o_full_addr); - p_cons_stage ^= 1; - if (p_cons_stage == 0) { - p_cons_phase ^= 1; - } - o_prod_phase ^= 1; - } - elect_commit(q_empty_addr + (q_cons_stage) * 8); - q_cons_stage += 1; - if (q_cons_stage == 2) { q_cons_stage = 0; _phase_q_full ^= 1; } - int tail_score_stage = score_prod_stage ^ 1; - int tail_score_phase = score_prod_phase ^ score_prod_stage; - mbarrier_wait(score_empty_addr + (tail_score_stage) * 8, 1 - tail_score_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - mbarrier_wait(o_empty_addr, 1 - o_prod_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - #pragma unroll - for (int v_half_1 = 0; v_half_1 < 2; v_half_1++) { - mbarrier_wait(kv_full_addr + (_kv_cons_state_state.index()) * 8, _kv_cons_state_state.phase()); - asm volatile("tcgen05.fence::after_thread_sync;"); - int row_bank_3 = ((v_half_1 != 0) ? 1048576 : 0); - int p_col_2 = ((p_cons_stage != 0) ? 160 : 32) + row_bank_3; - int _mma_b_lo_5 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (_kv_cons_state_state.index()) * 2048); - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 dhi, blo, ta, id;\n\t" - ".reg .b64 db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "" - "mov.b32 dhi, 0x40004040;\n\t" - "mov.b32 id, 71368720;\n\t" - "mov.b32 ta, %2;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p0;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 256;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], [ta], db, id, p1;\n\t" - "}\n" - :: "r"((tmem_tmem + (256 + row_bank_3))), "r"(_mma_b_lo_5), "r"(tmem_tmem + p_col_2), "r"(((first_pv) ? 0 : 1))); - elect_commit(kv_empty_addr + (_kv_cons_state_state.index()) * 8); - ++_kv_cons_state_state; - } - elect_commit(o_full_addr); - p_cons_stage ^= 1; - if (p_cons_stage == 0) { - p_cons_phase ^= 1; - } - o_prod_phase ^= 1; - } - } - mbarrier_wait(work_full_addr + (mma_work_stage) * 8, _phase_work_full_2); - uint32_t _clc_valid_2 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_2) - : "r"(work_response_addr + mma_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_6 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_6) - : "r"(work_response_addr + mma_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_7 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_7) - : "r"(work_response_addr + mma_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_8 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_8) - : "r"(work_response_addr + mma_work_stage * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(work_empty_addr + (mma_work_stage) * 8); - mma_work_stage += 1; - if (mma_work_stage == 2) { mma_work_stage = 0; _phase_work_full_2 ^= 1; } - if (_clc_valid_2 == 0) { - break; - } - mma_q = _clc_ctaid_6; - mma_head = _clc_ctaid_7; - mma_batch = _clc_ctaid_8; - } - } - } - // ---- Role: page_loader ---- - if (warp == 9) { - { // page_loader_main - unsigned int page_prod_stage = 0; - unsigned int page_prod_phase = 1; - unsigned int page_work_stage = 0; - unsigned int page_q = blockIdx.x; - unsigned int page_head = blockIdx.y; - unsigned int page_batch = blockIdx.z; - unsigned int _phase_work_full_3 = 0; - #pragma unroll 1 - for (unsigned int _work_3 = 0; _work_3 < total_work_items; _work_3++) { - int q_begin_3 = cum_seq_lens_q[(int)page_batch]; - int q_end_3 = cum_seq_lens_q[(int)(page_batch + 1)]; - int q_local_3 = (int)page_q; - int work_valid_3 = q_local_3 < q_end_3 - q_begin_3; - if (work_valid_3 != 0) { - int query_idx_3 = q_begin_3 + q_local_3; - int _max_0 = ((sparse_topk_lens[query_idx_3] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_3] + sparse_topk_lens_offset) : (0)); - int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); - int active_topk_3 = _min_0; - int num_tiles_3 = (active_topk_3 + 128 - 1) / 128; - int num_page_passes = (num_tiles_3 + 1) / 2; - int* swa_row = swa_indices + (query_idx_3 * swa_index_stride); - int* compressed_row = compressed_indices + (query_idx_3 * compressed_index_stride); - int _max_1 = (((active_topk_3 - 1) / 4 * 4) > (0) ? ((active_topk_3 - 1) / 4 * 4) : (0)); - int last_vec_base = _max_1; - #pragma unroll 1 - for (int page_pass = 0; page_pass < num_page_passes; page_pass++) { - #pragma unroll - for (int _duplicate = 0; _duplicate < 2; _duplicate++) { - mbarrier_wait(page_empty_addr + (page_prod_stage) * 8, page_prod_phase); - int page_dst = smem_page_offsets_addr + page_prod_stage * 1024; - #pragma unroll - for (int page_half = 0; page_half < 2; page_half++) { - int page_offset = page_pass * 256 + page_half * 128 + lane * 4; - int clamped_offset = ((page_offset < active_topk_3) ? page_offset : last_vec_base); - int* page_src = ((clamped_offset < 128) ? swa_row + clamped_offset : compressed_row + (clamped_offset - 128)); - asm volatile("cp.async.cg.shared::cta.global.L2::128B [%0], [%1], 16;" - :: "r"(page_dst + (page_half * 128 + lane * 4) * 4), "l"(page_src)); - } - asm volatile("cp.async.commit_group;"); - asm volatile("cp.async.wait_group 0;"); - mbarrier_arrive(page_full_addr + (page_prod_stage) * 8); - page_prod_stage += 1; - if (page_prod_stage == 6) { page_prod_stage = 0; page_prod_phase ^= 1; } - } - } - } - mbarrier_wait(work_full_addr + (page_work_stage) * 8, _phase_work_full_3); - uint32_t _clc_valid_0 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_0) - : "r"(work_response_addr + page_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_0 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_0) - : "r"(work_response_addr + page_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_1 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_1) - : "r"(work_response_addr + page_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_2 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_2) - : "r"(work_response_addr + page_work_stage * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(work_empty_addr + (page_work_stage) * 8); - page_work_stage += 1; - if (page_work_stage == 2) { page_work_stage = 0; _phase_work_full_3 ^= 1; } - if (_clc_valid_0 == 0) { - break; - } - page_q = _clc_ctaid_0; - page_head = _clc_ctaid_1; - page_batch = _clc_ctaid_2; - } - } - } - // ---- Role: scheduler ---- - if (warp == 10) { - { // scheduler_main - unsigned int scheduler_work_stage = 0; - unsigned int scheduler_throttle_stage = 0; - unsigned int _phase_throttle_full = 0; - unsigned int _phase_work_empty = 1; - unsigned int _phase_work_full_4 = 0; - #pragma unroll 1 - for (unsigned int _work_4 = 0; _work_4 < total_work_items; _work_4++) { - mbarrier_wait(throttle_full_addr + (scheduler_throttle_stage) * 8, _phase_throttle_full); - mbarrier_arrive(throttle_empty_addr + (scheduler_throttle_stage) * 8); - scheduler_throttle_stage += 1; - if (scheduler_throttle_stage == 2) { scheduler_throttle_stage = 0; _phase_throttle_full ^= 1; } - if (elect_sync()) { - mbarrier_wait(work_empty_addr + (scheduler_work_stage) * 8, _phase_work_empty); - mbarrier_arrive_expect_tx(work_full_addr + (scheduler_work_stage) * 8, 16); - asm volatile( - "fence.proxy.async.shared::cta;\n\t" - "clusterlaunchcontrol.try_cancel.async.shared::cta" - ".mbarrier::complete_tx::bytes.b128" - " [%0], [%1];" - :: "r"(work_response_addr + scheduler_work_stage * 16 + 0 * 16), "r"(work_full_addr + scheduler_work_stage * 8) - : "memory"); - } - mbarrier_wait(work_full_addr + (scheduler_work_stage) * 8, _phase_work_full_4); - uint32_t _clc_valid_5 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_5) - : "r"(work_response_addr + scheduler_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_15 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_15) - : "r"(work_response_addr + scheduler_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_16 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_16) - : "r"(work_response_addr + scheduler_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_17 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_17) - : "r"(work_response_addr + scheduler_work_stage * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(work_empty_addr + (scheduler_work_stage) * 8); - scheduler_work_stage += 1; - if (scheduler_work_stage == 2) { scheduler_work_stage = 0; _phase_work_empty ^= 1; _phase_work_full_4 ^= 1; } - if (_clc_valid_5 == 0) { - break; - } - } - } - } - // ---- Role: padding ---- - if (warp == 11) { - { // padding_main - unsigned int padding_work_stage = 0; - unsigned int _phase_work_full_5 = 0; - #pragma unroll 1 - for (unsigned int _work_5 = 0; _work_5 < total_work_items; _work_5++) { - mbarrier_wait(work_full_addr + (padding_work_stage) * 8, _phase_work_full_5); - uint32_t _clc_valid_6 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_6) - : "r"(work_response_addr + padding_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_18 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_18) - : "r"(work_response_addr + padding_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_19 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_19) - : "r"(work_response_addr + padding_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_20 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_20) - : "r"(work_response_addr + padding_work_stage * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(work_empty_addr + (padding_work_stage) * 8); - padding_work_stage += 1; - if (padding_work_stage == 2) { padding_work_stage = 0; _phase_work_full_5 ^= 1; } - if (_clc_valid_6 == 0) { - break; - } - } - } - } - // ---- Role: load ---- - if (warp >= 12 && warp <= 15) { - asm volatile("setmaxnreg.dec.sync.aligned.u32 72;"); - { // load_main - const int load_rank = warp - 12; - unsigned int q_prod_stage = 0; - unsigned int kv_prod_stage = 0; - unsigned int page_k_stage = 0; - unsigned int page_k_phase = 0; - unsigned int page_v_stage = 1; - unsigned int page_v_phase = 0; - unsigned int load_work_stage = 0; - unsigned int throttle_stage = 0; - unsigned int load_q = blockIdx.x; - unsigned int load_head = blockIdx.y; - unsigned int load_batch = blockIdx.z; - unsigned int _phase_throttle_empty = 1; - unsigned int _phase_q_empty = 1; - unsigned int _phase_kv_empty = 1; - unsigned int _phase_work_full_6 = 0; - #pragma unroll 1 - for (unsigned int _work_6 = 0; _work_6 < total_work_items; _work_6++) { - mbarrier_wait(throttle_empty_addr + (throttle_stage) * 8, _phase_throttle_empty); - mbarrier_arrive(throttle_full_addr + (throttle_stage) * 8); - throttle_stage += 1; - if (throttle_stage == 2) { throttle_stage = 0; _phase_throttle_empty ^= 1; } - int q_begin_4 = cum_seq_lens_q[(int)load_batch]; - int q_end_4 = cum_seq_lens_q[(int)(load_batch + 1)]; - int q_local_4 = (int)load_q; - int work_valid_4 = q_local_4 < q_end_4 - q_begin_4; - if (work_valid_4 != 0) { - int query_idx_4 = q_begin_4 + q_local_4; - int _max_2 = ((sparse_topk_lens[query_idx_4] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx_4] + sparse_topk_lens_offset) : (0)); - int _min_1 = ((_max_2) < (sparse_topk) ? (_max_2) : (sparse_topk)); - int active_topk_4 = _min_1; - int num_tiles_4 = (active_topk_4 + 128 - 1) / 128; - mbarrier_wait(q_empty_addr + (q_prod_stage) * 8, _phase_q_empty); - if (load_rank == 0) { - if (elect_sync()) { - mbarrier_arrive_expect_tx(q_full_addr + (q_prod_stage) * 8, 32768); - } - } - asm volatile("griddepcontrol.wait;" ::: "memory"); - if (load_rank == 0) { - if (elect_sync()) { - #pragma unroll - for (int q_chunk = 0; q_chunk < 4; q_chunk++) { - asm volatile( - "cp.async.bulk.tensor.4d.shared::cluster.global.tile.mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5}], [%6];" - :: "r"(smem_q_addr + q_prod_stage * 32768 + (unsigned int)(q_chunk * 8192)), "l"((&tmap_q)), "r"(q_chunk * 128), "r"(0), "r"((int)load_head), "r"(query_idx_4), "r"(q_full_addr + (q_prod_stage) * 8) : "memory"); - } - } - } - q_prod_stage += 1; - if (q_prod_stage == 2) { q_prod_stage = 0; _phase_q_empty ^= 1; } - #pragma unroll 1 - for (int tile_2 = 0; tile_2 < num_tiles_4; tile_2++) { - if ((tile_2 & 1) == 0) { - mbarrier_wait(page_full_addr + (page_k_stage) * 8, page_k_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - } - int page_k_base = smem_page_offsets_addr + page_k_stage * 1024; - #pragma unroll - for (int k_half = 0; k_half < 2; k_half++) { - mbarrier_wait(kv_empty_addr + (kv_prod_stage) * 8, _phase_kv_empty); - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (kv_prod_stage) * 8, 8192); - if (tile_2 == 0) { - #pragma unroll - for (int local_group = 0; local_group < 8; local_group++) { - int group = load_rank * 8 + local_group; - int group_offset = (tile_2 & 1) * 128 + group * 4; - int rows[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&rows[0])), "=r"(*reinterpret_cast(&rows[(0) + 1])), "=r"(*reinterpret_cast(&rows[(0) + 2])), "=r"(*reinterpret_cast(&rows[(0) + 3])) - : "r"(page_k_base + group_offset * 4)); - int k_dst = smem_k_addr + kv_prod_stage * 32768; - tma_gather4_gmem2smem(k_dst + group * 512, (&tmap_swa_kv), k_half * 256, rows[0], rows[1], rows[2], rows[3], kv_full_addr + (kv_prod_stage) * 8); - tma_gather4_gmem2smem(k_dst + 16384 + group * 512, (&tmap_swa_kv), k_half * 256 + 128, rows[0], rows[1], rows[2], rows[3], kv_full_addr + (kv_prod_stage) * 8); - } - } else { - #pragma unroll - for (int local_group_1 = 0; local_group_1 < 8; local_group_1++) { - int group_1 = load_rank * 8 + local_group_1; - int group_offset_1 = (tile_2 & 1) * 128 + group_1 * 4; - int rows_1[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&rows_1[0])), "=r"(*reinterpret_cast(&rows_1[(0) + 1])), "=r"(*reinterpret_cast(&rows_1[(0) + 2])), "=r"(*reinterpret_cast(&rows_1[(0) + 3])) - : "r"(page_k_base + group_offset_1 * 4)); - int k_dst_1 = smem_k_addr + kv_prod_stage * 32768; - tma_gather4_gmem2smem(k_dst_1 + group_1 * 512, (&tmap_compressed_kv), k_half * 256, rows_1[0], rows_1[1], rows_1[2], rows_1[3], kv_full_addr + (kv_prod_stage) * 8); - tma_gather4_gmem2smem(k_dst_1 + 16384 + group_1 * 512, (&tmap_compressed_kv), k_half * 256 + 128, rows_1[0], rows_1[1], rows_1[2], rows_1[3], kv_full_addr + (kv_prod_stage) * 8); - } - } - } - kv_prod_stage += 1; - if (kv_prod_stage == 4) { kv_prod_stage = 0; _phase_kv_empty ^= 1; } - } - if ((tile_2 & 1) != 0 || tile_2 == num_tiles_4 - 1) { - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(page_empty_addr + (page_k_stage) * 8); - page_k_stage += 1; - if (page_k_stage == 6) { page_k_stage = 0; page_k_phase ^= 1; } - page_k_stage += 1; - if (page_k_stage == 6) { page_k_stage = 0; page_k_phase ^= 1; } - } - if (tile_2 > 0) { - int prev_tile = tile_2 - 1; - if ((prev_tile & 1) == 0) { - mbarrier_wait(page_full_addr + (page_v_stage) * 8, page_v_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - } - int page_v_base = smem_page_offsets_addr + page_v_stage * 1024; - #pragma unroll - for (int v_half_2 = 0; v_half_2 < 2; v_half_2++) { - mbarrier_wait(kv_empty_addr + (kv_prod_stage) * 8, _phase_kv_empty); - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (kv_prod_stage) * 8, 8192); - if (prev_tile == 0) { - #pragma unroll - for (int local_group_2 = 0; local_group_2 < 8; local_group_2++) { - int group_2 = load_rank * 8 + local_group_2; - int group_offset_2 = (prev_tile & 1) * 128 + group_2 * 4; - int rows_2[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&rows_2[0])), "=r"(*reinterpret_cast(&rows_2[(0) + 1])), "=r"(*reinterpret_cast(&rows_2[(0) + 2])), "=r"(*reinterpret_cast(&rows_2[(0) + 3])) - : "r"(page_v_base + group_offset_2 * 4)); - int v_dst = smem_v_addr + kv_prod_stage * 32768; - tma_gather4_gmem2smem(v_dst + group_2 * 512, (&tmap_swa_kv), v_half_2 * 256, rows_2[0], rows_2[1], rows_2[2], rows_2[3], kv_full_addr + (kv_prod_stage) * 8); - tma_gather4_gmem2smem(v_dst + 16384 + group_2 * 512, (&tmap_swa_kv), v_half_2 * 256 + 128, rows_2[0], rows_2[1], rows_2[2], rows_2[3], kv_full_addr + (kv_prod_stage) * 8); - } - } else { - #pragma unroll - for (int local_group_3 = 0; local_group_3 < 8; local_group_3++) { - int group_3 = load_rank * 8 + local_group_3; - int group_offset_3 = (prev_tile & 1) * 128 + group_3 * 4; - int rows_3[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&rows_3[0])), "=r"(*reinterpret_cast(&rows_3[(0) + 1])), "=r"(*reinterpret_cast(&rows_3[(0) + 2])), "=r"(*reinterpret_cast(&rows_3[(0) + 3])) - : "r"(page_v_base + group_offset_3 * 4)); - int v_dst_1 = smem_v_addr + kv_prod_stage * 32768; - tma_gather4_gmem2smem(v_dst_1 + group_3 * 512, (&tmap_compressed_kv), v_half_2 * 256, rows_3[0], rows_3[1], rows_3[2], rows_3[3], kv_full_addr + (kv_prod_stage) * 8); - tma_gather4_gmem2smem(v_dst_1 + 16384 + group_3 * 512, (&tmap_compressed_kv), v_half_2 * 256 + 128, rows_3[0], rows_3[1], rows_3[2], rows_3[3], kv_full_addr + (kv_prod_stage) * 8); - } - } - } - kv_prod_stage += 1; - if (kv_prod_stage == 4) { kv_prod_stage = 0; _phase_kv_empty ^= 1; } - } - if ((prev_tile & 1) != 0) { - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(page_empty_addr + (page_v_stage) * 8); - page_v_stage += 1; - if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } - page_v_stage += 1; - if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } - } - } - } - int last_tile = num_tiles_4 - 1; - mbarrier_wait(page_full_addr + (page_v_stage) * 8, page_v_phase); - asm volatile("tcgen05.fence::after_thread_sync;"); - int page_v_base_1 = smem_page_offsets_addr + page_v_stage * 1024; - #pragma unroll - for (int v_half_3 = 0; v_half_3 < 2; v_half_3++) { - mbarrier_wait(kv_empty_addr + (kv_prod_stage) * 8, _phase_kv_empty); - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (kv_prod_stage) * 8, 8192); - if (last_tile == 0) { - #pragma unroll - for (int local_group_4 = 0; local_group_4 < 8; local_group_4++) { - int group_4 = load_rank * 8 + local_group_4; - int group_offset_4 = (last_tile & 1) * 128 + group_4 * 4; - int rows_4[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&rows_4[0])), "=r"(*reinterpret_cast(&rows_4[(0) + 1])), "=r"(*reinterpret_cast(&rows_4[(0) + 2])), "=r"(*reinterpret_cast(&rows_4[(0) + 3])) - : "r"(page_v_base_1 + group_offset_4 * 4)); - int v_dst_2 = smem_v_addr + kv_prod_stage * 32768; - tma_gather4_gmem2smem(v_dst_2 + group_4 * 512, (&tmap_swa_kv), v_half_3 * 256, rows_4[0], rows_4[1], rows_4[2], rows_4[3], kv_full_addr + (kv_prod_stage) * 8); - tma_gather4_gmem2smem(v_dst_2 + 16384 + group_4 * 512, (&tmap_swa_kv), v_half_3 * 256 + 128, rows_4[0], rows_4[1], rows_4[2], rows_4[3], kv_full_addr + (kv_prod_stage) * 8); - } - } else { - #pragma unroll - for (int local_group_5 = 0; local_group_5 < 8; local_group_5++) { - int group_5 = load_rank * 8 + local_group_5; - int group_offset_5 = (last_tile & 1) * 128 + group_5 * 4; - int rows_5[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&rows_5[0])), "=r"(*reinterpret_cast(&rows_5[(0) + 1])), "=r"(*reinterpret_cast(&rows_5[(0) + 2])), "=r"(*reinterpret_cast(&rows_5[(0) + 3])) - : "r"(page_v_base_1 + group_offset_5 * 4)); - int v_dst_3 = smem_v_addr + kv_prod_stage * 32768; - tma_gather4_gmem2smem(v_dst_3 + group_5 * 512, (&tmap_compressed_kv), v_half_3 * 256, rows_5[0], rows_5[1], rows_5[2], rows_5[3], kv_full_addr + (kv_prod_stage) * 8); - tma_gather4_gmem2smem(v_dst_3 + 16384 + group_5 * 512, (&tmap_compressed_kv), v_half_3 * 256 + 128, rows_5[0], rows_5[1], rows_5[2], rows_5[3], kv_full_addr + (kv_prod_stage) * 8); - } - } - } - kv_prod_stage += 1; - if (kv_prod_stage == 4) { kv_prod_stage = 0; _phase_kv_empty ^= 1; } - } - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(page_empty_addr + (page_v_stage) * 8); - page_v_stage += 1; - if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } - page_v_stage += 1; - if (page_v_stage == 6) { page_v_stage = 0; page_v_phase ^= 1; } - } - mbarrier_wait(work_full_addr + (load_work_stage) * 8, _phase_work_full_6); - uint32_t _clc_valid_1 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "selp.u32 %0, 1, 0, p1;\n\t" - "}\n" - : "=r"(_clc_valid_1) - : "r"(work_response_addr + load_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_3 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_3) - : "r"(work_response_addr + load_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_4 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::y.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_4) - : "r"(work_response_addr + load_work_stage * 16 + 0 * 16) - : "memory"); - uint32_t _clc_ctaid_5 = 0; - asm volatile( - "{\n\t" - ".reg .pred p1;\n\t" - ".reg .b128 clc_r;\n\t" - "ld.shared.b128 clc_r, [%1];\n\t" - "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_r;\n\t" - "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid::z.b32.b128 %0, clc_r;\n\t" - "}\n" - : "=r"(_clc_ctaid_5) - : "r"(work_response_addr + load_work_stage * 16 + 0 * 16) - : "memory"); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(work_empty_addr + (load_work_stage) * 8); - load_work_stage += 1; - if (load_work_stage == 2) { load_work_stage = 0; _phase_work_full_6 ^= 1; } - if (_clc_valid_1 == 0) { - break; - } - load_q = _clc_ctaid_3; - load_head = _clc_ctaid_4; - load_batch = _clc_ctaid_5; - } - } - } - - // Cleanup -} - -} // extern "C" diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_f89bc6561b38462e9538_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_f89bc6561b38462e9538_binding.cu deleted file mode 100644 index 02f9f0f2ff8..00000000000 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_f89bc6561b38462e9538_binding.cu +++ /dev/null @@ -1,389 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -// Generated by the Cake source exporter — do not edit. -// tvm-ffi direct-source launcher for Cake kernel 'kernel_cake_dsv4_f89bc6561b38462e9538'. -#include -#include -#include -#include - -#include "tvm_ffi_utils.h" - -#include -#include -#include -#include -#include -#include -#include - -extern "C" __global__ void kernel_cake_dsv4_f89bc6561b38462e9538(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int num_head_tiles, int has_sinks); - - -namespace cake_host_shim_13f3ad5a6154b3a4 { - -using tvm::ffi::Optional; -using tvm::ffi::TensorView; - -class ScopedCudaDevice { - public: - explicit ScopedCudaDevice(int device_id) { - cudaError_t error = cudaGetDevice(&previous_device_); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed before host-shim launch: cudaError=" - << static_cast(error); - if (previous_device_ != device_id) { - error = cudaSetDevice(device_id); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaSetDevice failed before host-shim launch for cuda:" - << device_id << ": cudaError=" << static_cast(error); - restore_ = true; - } - } - - ScopedCudaDevice(const ScopedCudaDevice&) = delete; - ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; - - ~ScopedCudaDevice() noexcept { - if (restore_) { - (void)cudaSetDevice(previous_device_); - } - } - - private: - int previous_device_ = -1; - bool restore_ = false; -}; - -inline int64_t CakeDeviceMultiprocessorCount(int device_id) { - constexpr int kMaxCachedCudaDevices = 64; - TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) - << "physical-SM-count cache does not cover cuda:" << device_id; - static std::atomic count_by_device[kMaxCachedCudaDevices]{}; - int cached = count_by_device[device_id].load(std::memory_order_acquire); - if (cached > 0) return cached; - - int count = 0; - cudaError_t error = cudaDeviceGetAttribute( - &count, cudaDevAttrMultiProcessorCount, device_id); - TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) - << "querying multiProcessorCount failed for cuda:" << device_id - << ": cudaError=" << static_cast(error) << ", count=" << count; - // Concurrent first launches may repeat the immutable device query, but all - // publication is atomic and every later hot-path lookup is one acquire load. - count_by_device[device_id].store(count, std::memory_order_release); - return count; -} - -inline void CheckCudaTensor(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) - << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; -} - -inline void CheckSameCudaDevice( - const TensorView& t, - const TensorView& reference, - const char* name, - const char* reference_name) { - TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) - << name << " must be on the same CUDA device as " << reference_name - << ": got cuda:" << t.device().device_id - << " versus cuda:" << reference.device().device_id; -} - -inline void CheckCurrentCudaDevice( - const TensorView& reference, - const char* reference_name) { - int current_device = -1; - cudaError_t error = cudaGetDevice(¤t_device); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed while validating " << reference_name - << ": cudaError=" << static_cast(error); - TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) - << "current CUDA device must match " << reference_name - << ": current=cuda:" << current_device - << ", tensor=cuda:" << reference.device().device_id; -} - -inline void CheckContiguous(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; -} - -inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { - DLDataType d = t.dtype(); - TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) - << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits - << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits - << ", lanes=" << (int)d.lanes << ")"; -} - -// A logical axis.outer(trailing) folds every source dim above the trailing -// dimensions. Shape products are independent of physical strides, so verify -// the leading dimensions form one dense row-major chain instead of inventing -// a "folded stride". The descriptor reads its exact adjacent physical step -// separately through stride[-(trailing + 1)]. -inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { - TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) - << name << " cannot fold leading dimensions above " << trailing - << " trailing dims from ndim=" << t.ndim(); - int outer_last = t.ndim() - trailing - 1; - if (outer_last <= 0) { - return; - } - int64_t step = t.stride(outer_last); - TVM_FFI_CHECK(step > 0, ValueError) - << name << " physical strides must be positive"; - int64_t expected = step; - for (int axis = outer_last - 1; axis >= 0; --axis) { - expected *= t.size(axis + 1); - if (t.size(axis) > 1) { - TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) - << name << " leading dims are not physically foldable above " << trailing - << " trailing dims: stride(" << axis << ")=" << t.stride(axis) - << ", expected " << expected; - } - } -} - -// 4D TMA descriptor for buffer 'tmap_q' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 3, ValueError) - << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - int64_t d2 = t.size(t.ndim() - 2); - int64_t d3 = t.size(t.ndim() - 3); - TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) - << "TMA source 'tmap_q' trailing dims must be positive"; - TVM_FFI_CHECK(d1 % 64 == 0, ValueError) - << "TMA source 'tmap_q' extent " << d1 - << " must divide exactly by " << 64; - uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(d3)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) - << "TMA box (64, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; - int64_t carrier_stride_0 = d1; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; - int64_t carrier_stride_1 = 64; - TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; - TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; - TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; - int64_t carrier_stride_2 = (d2 * d1); - TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; - TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; - TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; - uint64_t global_strides[3] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - (uint64_t)((carrier_stride_1 * 16) / 8), - (uint64_t)((carrier_stride_2 * 16) / 8), - }; - uint32_t box_dim[4] = {64u, 64u, 1u, 1u}; - uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_swa_kv' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_swa_kv(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_swa_kv' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_swa_kv' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_swa_kv' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_swa_kv"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_swa_kv' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_kv'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_swa_kv' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_swa_kv') failed: CUresult=" << (int)r; - return tm; -} - -void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_kv, TensorView arg_O, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_num_head_tiles, int64_t arg_has_sinks, int64_t grid_x, int64_t grid_y, int64_t grid_z) { - DLDevice dev = arg_tmap_q.device(); - ScopedCudaDevice device_guard(dev.device_id); - CheckCudaTensor(arg_tmap_q, "tmap_q"); - CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); - CheckContiguous(arg_tmap_q, "tmap_q"); - CheckCudaTensor(arg_tmap_swa_kv, "tmap_swa_kv"); - CheckDtype(arg_tmap_swa_kv, "tmap_swa_kv", 4, 16, 1); - CheckCudaTensor(arg_O, "O"); - CheckDtype(arg_O, "O", 4, 16, 1); - CheckContiguous(arg_O, "O"); - CheckCudaTensor(arg_swa_indices, "swa_indices"); - CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); - CheckContiguous(arg_swa_indices, "swa_indices"); - CheckCudaTensor(arg_compressed_indices, "compressed_indices"); - CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); - CheckContiguous(arg_compressed_indices, "compressed_indices"); - CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); - CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckCudaTensor(arg_sinks, "sinks"); - CheckDtype(arg_sinks, "sinks", 2, 32, 1); - CheckContiguous(arg_sinks, "sinks"); - CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); - CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); - CheckContiguous(arg_bmm1_scale, "bmm1_scale"); - CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); - CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); - CheckContiguous(arg_bmm2_scale, "bmm2_scale"); - TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) - << "scalar 'num_heads' value " << arg_num_heads - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) - << "scalar 'swa_index_stride' value " << arg_swa_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) - << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) - << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) - << "scalar 'sparse_topk' value " << arg_sparse_topk - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_num_head_tiles >= -2147483648LL && arg_num_head_tiles <= 2147483647LL, ValueError) - << "scalar 'num_head_tiles' value " << arg_num_head_tiles - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) - << "scalar 'has_sinks' value " << arg_has_sinks - << " is outside i32 range [-2147483648, 2147483647]"; - CheckSameCudaDevice(arg_tmap_swa_kv, arg_tmap_q, "tmap_swa_kv", "tmap_q"); - CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); - CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); - CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); - CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); - CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); - CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); - CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); - TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) - << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y - << ", " << grid_z << ")"; - - cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); - CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); - CUtensorMap p_tmap_swa_kv = EncodeTma_tmap_swa_kv(arg_tmap_swa_kv); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); - int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); - int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); - int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); - float* p_sinks = static_cast(arg_sinks.data_ptr()); - float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); - float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); - int32_t v_num_heads = (int32_t)arg_num_heads; - int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; - int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; - int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; - int32_t v_sparse_topk = (int32_t)arg_sparse_topk; - int32_t v_num_head_tiles = (int32_t)arg_num_head_tiles; - int32_t v_has_sinks = (int32_t)arg_has_sinks; - void* kargs[] = {&p_tmap_q, &p_tmap_swa_kv, &p_O, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_num_head_tiles, &v_has_sinks}; - - static std::mutex smem_mu; - static auto* smem_status_by_device = - new std::unordered_map(); - cudaError_t smem_status = cudaSuccess; - { - std::lock_guard lock(smem_mu); - auto it = smem_status_by_device->find(dev.device_id); - if (it == smem_status_by_device->end()) { - int current_device = -1; - smem_status = cudaGetDevice(¤t_device); - if (smem_status == cudaSuccess && current_device != dev.device_id) { - smem_status = cudaErrorInvalidDevice; - } - if (smem_status == cudaSuccess) { - smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_f89bc6561b38462e9538), - cudaFuncAttributeMaxDynamicSharedMemorySize, - 198528); - } - smem_status_by_device->emplace(dev.device_id, smem_status); - } else { - smem_status = it->second; - } - } - TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_f89bc6561b38462e9538 failed: " - << cudaGetErrorString(smem_status); - dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); - dim3 block(512u, 1u, 1u); - - cudaLaunchConfig_t config{}; - config.gridDim = grid; - config.blockDim = block; - config.dynamicSmemBytes = 198528u; - config.stream = stream; - cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_f89bc6561b38462e9538), kargs); - TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_f89bc6561b38462e9538 failed: " - << cudaGetErrorString(launch_status); -} - -} // namespace cake_host_shim_13f3ad5a6154b3a4 - -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_13f3ad5a6154b3a4::Run); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_f89bc6561b38462e9538_kernel.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_f89bc6561b38462e9538_kernel.cu deleted file mode 100644 index d9d9646824e..00000000000 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_f89bc6561b38462e9538_kernel.cu +++ /dev/null @@ -1,1642 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -typedef signed char int8_t; -typedef unsigned char uint8_t; -typedef unsigned short uint16_t; -typedef unsigned int uint32_t; -#if defined(__CUDACC_RTC__) -typedef unsigned long long uint64_t; -#else -typedef unsigned long uint64_t; -#endif -static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); -typedef signed int int32_t; -typedef short int int16_t; -struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; -struct __align__(64) CakeTensorMap64 { uint64_t opaque[16]; }; -static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); -static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); - -#if defined(__CUDACC_RTC__) -typedef struct __align__(128) { uint64_t opaque[16]; } CUtensorMap; -#else -#include -#endif - -static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); -static_assert(alignof(CakeTensorMap) >= alignof(CUtensorMap), "CakeTensorMap alignment must cover the CUtensorMap CUDA ABI"); -#include -#include - -#define CAKE_INF CUDART_INF_F -#define TMEM_NCOLS 384 -#define TMEM_TMEM_OFFSET 0 -#define NUM_KV_PIPE_STAGES 4 -#define SMEM_SMEM_Q_OFF 1024 -#define SMEM_SMEM_Q_STAGE_BYTES 8192 -#define SMEM_SMEM_Q_STRIDE 8192 -#define SMEM_SMEM_KV_OFF 66560 -#define SMEM_SMEM_KV_STAGE_BYTES 32768 -#define SMEM_SMEM_KV_STRIDE 32768 -#define SMEM_SMEM_V_OFF 66560 -#define SMEM_SMEM_V_STAGE_BYTES 32768 -#define SMEM_SMEM_V_STRIDE 32768 -#define SMEM_SMEM_SUM_OFF 197632 -#define SMEM_SMEM_SUM_STAGE_BYTES 256 -#define SMEM_SMEM_SUM_STRIDE 256 -#define SMEM_SMEM_INDEX_FLAGS_OFF 198400 -#define SMEM_SMEM_INDEX_FLAGS_STAGE_BYTES 16 -#define SMEM_SMEM_INDEX_FLAGS_STRIDE 16 -#define SMEM_SMEM_INDICES_OFF 197888 -#define SMEM_SMEM_INDICES_STAGE_BYTES 512 -#define SMEM_SMEM_INDICES_STRIDE 512 -#define SMEM_TOTAL 198528 -#define THREADS 512 - -#include - -__device__ __forceinline__ uint32_t elect_sync() { - uint32_t pred = 0; - asm volatile( - "{\n\t" - ".reg .pred %%px;\n\t" - "elect.sync _|%%px, %1;\n\t" - "@%%px mov.s32 %0, 1;\n\t" - "}\n" - : "+r"(pred) - : "r"(0xFFFFFFFF)); - return pred; -} - - -__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { - asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" - :: "r"(mbar_addr), "r"(count) : "memory"); -} - -__device__ __forceinline__ void mbarrier_init_generic(void* mbar_addr, int count) { - asm volatile("mbarrier.init.b64 [%0], %1;" - :: "l"(mbar_addr), "r"(count) : "memory"); -} - - -__device__ __forceinline__ uint32_t mbarrier_try_wait_plain(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - -__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { - uint32_t token; - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%1], %2;\n\t" - "selp.u32 %0, 1, 0, P1;\n\t" - "}\n" - : "=r"(token) - : "r"(mbar_addr), "r"(phase) : "memory"); - return token; -} - - -// CTA-local pipelines have short, resident producer/consumer edges. Omitting -// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the -// explicit loop still makes this helper blocking until acquire succeeds. -__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE;\n\t" - "bra.uni LAB_WAIT;\n\t" - "DONE:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Source-faithful relaxed CTA wait used only by a typed protocol that does -// not attach the PTX acquire qualifier, such as FA4's interior P-ready edge. -__device__ __forceinline__ void mbarrier_wait_relaxed(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, 10000000;\n\t" - "@P1 bra.uni DONE_RELAXED;\n\t" - "bra.uni LAB_WAIT_RELAXED;\n\t" - "DONE_RELAXED:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -// Exact source ports may request the PTX suspendTimeHint operand explicitly. -// The hint is expressed in nanoseconds and is kept separate from the canonical -// no-hint CTA helper so unrelated schedules retain their existing retry path. -__device__ __forceinline__ void mbarrier_wait_suspend( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_SUSPEND:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_SUSPEND;\n\t" - "bra.uni LAB_WAIT_SUSPEND;\n\t" - "DONE_SUSPEND:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1;\n\t" - "@P1 bra.uni DONE_CLUSTER;\n\t" - "bra.uni LAB_WAIT_CLUSTER;\n\t" - "DONE_CLUSTER:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - ".reg .u32 WAIT_ADDR;\n\t" - "mov.u32 WAIT_ADDR, %0;\n\t" - "LAB_WAIT_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" - " P1, [WAIT_ADDR], %1, %2;\n\t" - "@P1 bra.uni DONE_HINT;\n\t" - "bra.uni LAB_WAIT_HINT;\n\t" - "DONE_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -// Exact unqualified CTA wait used by source schedules whose PTX intentionally -// omits the acquire qualifier while retaining a typed suspendTimeHint operand. -__device__ __forceinline__ void mbarrier_wait_relaxed_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_RELAXED_HINT:\n\t" - "mbarrier.try_wait.parity.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra DONE_RELAXED_HINT;\n\t" - "bra LAB_WAIT_RELAXED_HINT;\n\t" - "DONE_RELAXED_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint)); -} - -__device__ __forceinline__ void mbarrier_wait_cluster_hint( - int mbar_addr, int phase, uint32_t suspend_time_hint) { - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "LAB_WAIT_CLUSTER_HINT:\n\t" - "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" - " P1, [%0], %1, %2;\n\t" - "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" - "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" - "DONE_CLUSTER_HINT:\n\t" - "}\n" - :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); -} - -__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_suspend( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_suspend(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { - if (token == 0) { - mbarrier_wait_cluster(mbar_addr, phase); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_hint(mbar_addr, phase, suspend_time_hint); - } -} - -__device__ __forceinline__ void mbarrier_wait_token_cluster_hint( - int mbar_addr, int phase, uint32_t token, uint32_t suspend_time_hint) { - if (token == 0) { - mbarrier_wait_cluster_hint(mbar_addr, phase, suspend_time_hint); - } -} - - -__device__ __forceinline__ void tcgen05_mma_f16( - int taddr, uint64_t a_desc, uint64_t b_desc, - uint32_t i_desc, int enable_input_d) { - asm volatile( - "{\n\t" - ".reg .pred p;\n\t" - "setp.ne.b32 p, %4, 0;\n\t" - "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" - "}\n" - :: "r"(taddr), "l"(a_desc), "l"(b_desc), - "r"(i_desc), "r"(enable_input_d) - : "memory"); -} - - -__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { - return (x & 0x3FFFFULL) >> 4ULL; -} - - -union MmaSmemDesc { - uint64_t u64; - uint32_t u32[2]; -}; - -__device__ __forceinline__ void incr_smem_desc_lo(uint64_t& smem_desc, uint32_t offset) { - MmaSmemDesc tmp; - tmp.u64 = smem_desc; - tmp.u32[0] += offset; - smem_desc = tmp.u64; -} - - -__device__ __forceinline__ void elect_commit(int mbar_addr) { - asm volatile( - "{\n\t" - ".reg .pred leader;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" - ".shared::cluster.b64 [%0];\n\t" - "}\n" - :: "r"(mbar_addr)); -} - - -__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { - asm volatile( - "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { - asm volatile( - "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" - :: "r"(mbar_addr), "r"(bytes) : "memory"); -} - - -__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { - asm volatile( - "tcgen05.ld.sync.aligned.32x32b.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7," - " %8, %9, %10, %11, %12, %13, %14, %15," - " %16, %17, %18, %19, %20, %21, %22, %23," - " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" - : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), - "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), - "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), - "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), - "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), - "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), - "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), - "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) - : "r"(tmem_addr)); -} - - -__device__ __forceinline__ float approx_exp2(float x) { - float y; - asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float approx_rcp(float x) { - float y; - asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} - - -__device__ __forceinline__ float max_noftz(float a, float b) { - float c; - asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); - return c; -} - - -__device__ __forceinline__ float warp_reduce_max(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val = max_noftz(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); - return val; -} - - -__device__ __forceinline__ float warp_reduce_sum(float val) { - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) - val += __shfl_xor_sync(0xFFFFFFFF, val, offset); - return val; -} - - -__device__ __forceinline__ float row_max_reduce(float2 acc) { - return max_noftz(acc.x, acc.y); -} - - -__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { - #pragma unroll - for (int j = 0; j < 16; j++) { - if (j % 2 == 0) - acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); - else - acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); - } -} - - -__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void fma_f32x2_noftz_inplace(float2* a, float2 b, float2 c) { - unsigned long long r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(r) - : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - *(unsigned long long*)a = r; -} - -__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { - asm("mul.rn.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { - asm("mul.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { - asm("add.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void add_f32x2_noftz_inplace(float2* a, float2 b) { - asm("add.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { - asm("sub.rn.ftz.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ void sub_f32x2_noftz_inplace(float2* a, float2 b) { - asm("sub.f32x2 %0, %0, %1;" - : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); -} - -__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("add.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { - float2 r; - asm("sub.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 sub_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("sub.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ void fma_scale_x32( - float* sv, const float2* scale2, const float2* neg_max2) -{ - float2* sv_2 = reinterpret_cast(sv); - #pragma unroll - for (int j = 0; j < 16; j++) - fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); -} - -__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), - "l"(*(unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_noftz(float2 a, float2 b) { - float2 r; - asm("mul.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); - return r; -} - - -// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) - -__device__ __forceinline__ float2 add_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("add.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("add.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("add.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("add.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("add.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("add.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("add.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 add_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("add.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rn_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rn.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rz_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rz.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rm_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rm.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_noftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 mul_f32x2_rp_ftz(float2 a, float2 b) { - float2 r; - asm("mul.rp.ftz.f32x2 %0, %1, %2;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rz.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rz_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rz.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rm.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rm_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rm.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_noftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm("fma.rp.ftz.f32x2 %0, %1, %2, %3;" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - -__device__ __forceinline__ float2 fma_sub_f32x2_rp_ftz(float2 a, float2 b, float2 c) { - float2 r; - asm volatile("{\n\t" - ".reg .f32 _c0, _c1;\n\t" - ".reg .b64 _neg_c;\n\t" - "mov.b64 {_c0, _c1}, %3;\n\t" - "neg.f32 _c0, _c0;\n\t" - "neg.f32 _c1, _c1;\n\t" - "mov.b64 _neg_c, {_c0, _c1};\n\t" - "fma.rp.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" - "}\n" - : "=l"(*(unsigned long long*)&r) - : "l"(*(const unsigned long long*)&a), - "l"(*(const unsigned long long*)&b), - "l"(*(const unsigned long long*)&c)); - return r; -} - - -__device__ __forceinline__ void fence_async_shared() { - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); -} - - -__device__ __forceinline__ uint64_t make_smem_desc(int addr) { - const int SBO = 1024; - return desc_encode(addr) - | (desc_encode(SBO) << 32ULL) - | (1ULL << 46ULL) - | (2ULL << 61ULL); -} - - -__device__ __forceinline__ void tma_2d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3}], [%4];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_4d_gmem2smem( - int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { - asm volatile( - "cp.async.bulk.tensor.4d.shared::cta.global" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5}], [%6];" - :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tma_gather4_gmem2smem( - int dst, const void *tmap_ptr, - int col_idx, int row0, int row1, int row2, int row3, - int mbar_addr) { - // Canonical .shared::cta form for non-multicast gather4, matching - // trtllm-gen / cuda_ptx and the PTX ISA qualifier order - // (dim.dst.src.load_mode.completion_mechanism). Per the PTX grammar, - // .shared::cluster is reserved for the multicast variant (ctaMask). - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" - ".mbarrier::complete_tx::bytes" - " [%0], [%1, {%2, %3, %4, %5, %6}], [%7];" - :: "r"(dst), "l"(tmap_ptr), "r"(col_idx), - "r"(row0), "r"(row1), "r"(row2), "r"(row3), - "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { - asm volatile( - "tcgen05.commit.cta_group::1.mbarrier::arrive::one" - ".shared::cluster.b64 [%0];" - :: "r"(mbar_addr) : "memory"); -} - - -__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { - uint32_t result; - asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" - : "=r"(result) : "r"(val)); - return result; -} - -extern "C" { - -__global__ __launch_bounds__(512, 1) void -kernel_cake_dsv4_f89bc6561b38462e9538(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_kv, __nv_bfloat16* __restrict__ O, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int num_head_tiles, int has_sinks) -{ - const int tid = threadIdx.x; - const int warp = make_warp_uniform(tid / 32); - const int lane = tid % 32; - - extern __shared__ __align__(1024) char smem_raw[]; - int smem; - smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); - - const int mbar_base = smem; - #define q_full_addr (mbar_base + 0) - #define kv_full_addr (mbar_base + 8) - #define kv_empty_addr (mbar_base + 40) - #define s_full_addr (mbar_base + 72) - #define p_full_addr (mbar_base + 80) - #define sum_ready_addr (mbar_base + 88) - #define o_done_addr (mbar_base + 96) - #define tmem_dealloc_addr (mbar_base + 104) - - const int bid = blockIdx.x; - const int num_bids = gridDim.x; - - const int cta_rank = 0; - - // Kernel setup ops - __nv_bfloat16* smem_q = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); - const int smem_q_addr = smem + 1024; - __nv_bfloat16* smem_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); - const int smem_kv_addr = smem + 66560; - __nv_bfloat16* smem_v = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); - const int smem_v_addr = smem + 66560; - float* smem_sum = reinterpret_cast(smem_raw + 197632); - const int smem_sum_addr = smem + 197632; - int* smem_index_flags = reinterpret_cast(smem_raw + 198400); - const int smem_index_flags_addr = smem + 198400; - int* smem_indices = reinterpret_cast(smem_raw + 197888); - const int smem_indices_addr = smem + 197888; - - // Mbarrier init (8 pipeline groups, 0 ordered-sequence groups, 14 barriers) - // Mbarriers at smem_raw[0..112) - - if (warp == 0) { - uint32_t leader = elect_sync(); - if (leader) { - // q_full: 1 barriers, init_count=1 - mbarrier_init(smem + 0, 1); - // --- pipeline 'kv_pipe' --- - // kv_full: 4 barriers, init_count=4 - mbarrier_init(smem + 8, 4); - mbarrier_init(smem + 16, 4); - mbarrier_init(smem + 24, 4); - mbarrier_init(smem + 32, 4); - // kv_empty: 4 barriers, init_count=1 - mbarrier_init(smem + 40, 1); - mbarrier_init(smem + 48, 1); - mbarrier_init(smem + 56, 1); - mbarrier_init(smem + 64, 1); - // s_full: 1 barriers, init_count=1 - mbarrier_init(smem + 72, 1); - // p_full: 1 barriers, init_count=128 - mbarrier_init(smem + 80, 128); - // sum_ready: 1 barriers, init_count=128 - mbarrier_init(smem + 88, 128); - // o_done: 1 barriers, init_count=1 - mbarrier_init(smem + 96, 1); - // tmem_dealloc: 1 barriers, init_count=256 - mbarrier_init(smem + 104, 256); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - } - - __syncwarp(); - - // TMEM alloc (512 columns, 384 used) - volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 112); - if (warp == 0) { - int _tmem_hold = smem + 112; - asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); - __syncwarp(); - asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); - } - - __syncthreads(); - asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); - - const int taddr = tmem_addr_storage[0]; - - // Kernel post-init ops - const int tmem_tmem = taddr; - - // ---- Ordered hardware-WG register redistribution ---- - // Dec phase frees registers before any WG attempts inc. - if (warp >= 8 && warp <= 11) { - asm volatile("setmaxnreg.dec.sync.aligned.u32 64;"); - } - - // ---- Role: softmax ---- - if (warp <= 3) { - asm volatile("setmaxnreg.inc.sync.aligned.u32 192;"); - { // softmax_main - float softmax_scale_log2 = bmm1_scale[0] * 1.4426950408889634f; - int query_head_work = blockIdx.x >> 2; - int query_idx = query_head_work / num_head_tiles; - int head_tile = query_head_work % num_head_tiles; - int _max_0 = ((sparse_topk_lens[query_idx] + sparse_topk_lens_offset) > (0) ? (sparse_topk_lens[query_idx] + sparse_topk_lens_offset) : (0)); - int _min_0 = ((_max_0) < (sparse_topk) ? (_max_0) : (sparse_topk)); - int active_topk = _min_0; - const int warp_in_compute = warp; - const int tmem_row_origin = warp_in_compute * 32; - const int logical_row_origin = warp_in_compute * 16; - const int my_row = logical_row_origin + lane % 16; - int head_idx = head_tile * 64 + my_row; - const int col_half = lane / 16; - unsigned int _phase_s_full_0 = 0; - mbarrier_wait(s_full_addr, _phase_s_full_0); - _phase_s_full_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - int score_addr = taddr + (unsigned int)(tmem_row_origin << 16); - float _tmem_load_0[64]; - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_0[0])), "=r"(*reinterpret_cast(&_tmem_load_0[1])), "=r"(*reinterpret_cast(&_tmem_load_0[2])), "=r"(*reinterpret_cast(&_tmem_load_0[3])), "=r"(*reinterpret_cast(&_tmem_load_0[4])), "=r"(*reinterpret_cast(&_tmem_load_0[5])), "=r"(*reinterpret_cast(&_tmem_load_0[6])), "=r"(*reinterpret_cast(&_tmem_load_0[7])), "=r"(*reinterpret_cast(&_tmem_load_0[8])), "=r"(*reinterpret_cast(&_tmem_load_0[9])), "=r"(*reinterpret_cast(&_tmem_load_0[10])), "=r"(*reinterpret_cast(&_tmem_load_0[11])), "=r"(*reinterpret_cast(&_tmem_load_0[12])), "=r"(*reinterpret_cast(&_tmem_load_0[13])), "=r"(*reinterpret_cast(&_tmem_load_0[14])), "=r"(*reinterpret_cast(&_tmem_load_0[15])), "=r"(*reinterpret_cast(&_tmem_load_0[16])), "=r"(*reinterpret_cast(&_tmem_load_0[17])), "=r"(*reinterpret_cast(&_tmem_load_0[18])), "=r"(*reinterpret_cast(&_tmem_load_0[19])), "=r"(*reinterpret_cast(&_tmem_load_0[20])), "=r"(*reinterpret_cast(&_tmem_load_0[21])), "=r"(*reinterpret_cast(&_tmem_load_0[22])), "=r"(*reinterpret_cast(&_tmem_load_0[23])), "=r"(*reinterpret_cast(&_tmem_load_0[24])), "=r"(*reinterpret_cast(&_tmem_load_0[25])), "=r"(*reinterpret_cast(&_tmem_load_0[26])), "=r"(*reinterpret_cast(&_tmem_load_0[27])), "=r"(*reinterpret_cast(&_tmem_load_0[28])), "=r"(*reinterpret_cast(&_tmem_load_0[29])), "=r"(*reinterpret_cast(&_tmem_load_0[30])), "=r"(*reinterpret_cast(&_tmem_load_0[31])) - : "r"(score_addr)); - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_0[32])), "=r"(*reinterpret_cast(&_tmem_load_0[33])), "=r"(*reinterpret_cast(&_tmem_load_0[34])), "=r"(*reinterpret_cast(&_tmem_load_0[35])), "=r"(*reinterpret_cast(&_tmem_load_0[36])), "=r"(*reinterpret_cast(&_tmem_load_0[37])), "=r"(*reinterpret_cast(&_tmem_load_0[38])), "=r"(*reinterpret_cast(&_tmem_load_0[39])), "=r"(*reinterpret_cast(&_tmem_load_0[40])), "=r"(*reinterpret_cast(&_tmem_load_0[41])), "=r"(*reinterpret_cast(&_tmem_load_0[42])), "=r"(*reinterpret_cast(&_tmem_load_0[43])), "=r"(*reinterpret_cast(&_tmem_load_0[44])), "=r"(*reinterpret_cast(&_tmem_load_0[45])), "=r"(*reinterpret_cast(&_tmem_load_0[46])), "=r"(*reinterpret_cast(&_tmem_load_0[47])), "=r"(*reinterpret_cast(&_tmem_load_0[48])), "=r"(*reinterpret_cast(&_tmem_load_0[49])), "=r"(*reinterpret_cast(&_tmem_load_0[50])), "=r"(*reinterpret_cast(&_tmem_load_0[51])), "=r"(*reinterpret_cast(&_tmem_load_0[52])), "=r"(*reinterpret_cast(&_tmem_load_0[53])), "=r"(*reinterpret_cast(&_tmem_load_0[54])), "=r"(*reinterpret_cast(&_tmem_load_0[55])), "=r"(*reinterpret_cast(&_tmem_load_0[56])), "=r"(*reinterpret_cast(&_tmem_load_0[57])), "=r"(*reinterpret_cast(&_tmem_load_0[58])), "=r"(*reinterpret_cast(&_tmem_load_0[59])), "=r"(*reinterpret_cast(&_tmem_load_0[60])), "=r"(*reinterpret_cast(&_tmem_load_0[61])), "=r"(*reinterpret_cast(&_tmem_load_0[62])), "=r"(*reinterpret_cast(&_tmem_load_0[63])) - : "r"(score_addr + 32)); - int valid_cols = active_topk - col_half * 64; - int causal_cols = smem_index_flags[col_half * 2 + 1]; - valid_cols = ((valid_cols < causal_cols) ? valid_cols : causal_cols); - uint32_t _slice_lo_mask_0; - { - int _lim_0 = valid_cols; - if (_lim_0 <= 0) { _slice_lo_mask_0 = 0u; } - else if (_lim_0 >= 32) { _slice_lo_mask_0 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_slice_lo_mask_0) : "r"(_lim_0)); - } - } - if (!(_slice_lo_mask_0 & (1u << 0))) _tmem_load_0[0] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 1))) _tmem_load_0[1] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 2))) _tmem_load_0[2] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 3))) _tmem_load_0[3] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 4))) _tmem_load_0[4] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 5))) _tmem_load_0[5] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 6))) _tmem_load_0[6] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 7))) _tmem_load_0[7] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 8))) _tmem_load_0[8] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 9))) _tmem_load_0[9] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 10))) _tmem_load_0[10] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 11))) _tmem_load_0[11] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 12))) _tmem_load_0[12] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 13))) _tmem_load_0[13] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 14))) _tmem_load_0[14] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 15))) _tmem_load_0[15] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 16))) _tmem_load_0[16] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 17))) _tmem_load_0[17] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 18))) _tmem_load_0[18] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 19))) _tmem_load_0[19] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 20))) _tmem_load_0[20] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 21))) _tmem_load_0[21] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 22))) _tmem_load_0[22] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 23))) _tmem_load_0[23] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 24))) _tmem_load_0[24] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 25))) _tmem_load_0[25] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 26))) _tmem_load_0[26] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 27))) _tmem_load_0[27] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 28))) _tmem_load_0[28] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 29))) _tmem_load_0[29] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 30))) _tmem_load_0[30] = -CAKE_INF; - if (!(_slice_lo_mask_0 & (1u << 31))) _tmem_load_0[31] = -CAKE_INF; - uint32_t _slice_lo_mask_1; - { - int _lim_1 = valid_cols - 32; - if (_lim_1 <= 0) { _slice_lo_mask_1 = 0u; } - else if (_lim_1 >= 32) { _slice_lo_mask_1 = 0xFFFFFFFFu; } - else { - asm volatile("{" - ".reg .u32 t;\n\t" - "shl.b32 t, 1, %1;\n\t" - "add.u32 %0, t, -1;\n\t" - "}" : "=r"(_slice_lo_mask_1) : "r"(_lim_1)); - } - } - if (!(_slice_lo_mask_1 & (1u << 0))) _tmem_load_0[32] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 1))) _tmem_load_0[33] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 2))) _tmem_load_0[34] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 3))) _tmem_load_0[35] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 4))) _tmem_load_0[36] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 5))) _tmem_load_0[37] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 6))) _tmem_load_0[38] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 7))) _tmem_load_0[39] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 8))) _tmem_load_0[40] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 9))) _tmem_load_0[41] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 10))) _tmem_load_0[42] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 11))) _tmem_load_0[43] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 12))) _tmem_load_0[44] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 13))) _tmem_load_0[45] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 14))) _tmem_load_0[46] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 15))) _tmem_load_0[47] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 16))) _tmem_load_0[48] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 17))) _tmem_load_0[49] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 18))) _tmem_load_0[50] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 19))) _tmem_load_0[51] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 20))) _tmem_load_0[52] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 21))) _tmem_load_0[53] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 22))) _tmem_load_0[54] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 23))) _tmem_load_0[55] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 24))) _tmem_load_0[56] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 25))) _tmem_load_0[57] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 26))) _tmem_load_0[58] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 27))) _tmem_load_0[59] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 28))) _tmem_load_0[60] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 29))) _tmem_load_0[61] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 30))) _tmem_load_0[62] = -CAKE_INF; - if (!(_slice_lo_mask_1 & (1u << 31))) _tmem_load_0[63] = -CAKE_INF; - float2 _reg_reduce_max2_2 = {-CAKE_INF, -CAKE_INF}; - row_max_x32_accum(&_tmem_load_0[0], _reg_reduce_max2_2); - row_max_x32_accum(&_tmem_load_0[32], _reg_reduce_max2_2); - float _tmem_load_0_max = row_max_reduce(_reg_reduce_max2_2); - float row_max = _tmem_load_0_max; - float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, row_max, 16); - float _max_1 = max_noftz(row_max, _shfl_xor_0); - row_max = _max_1; - if (has_sinks != 0 && head_idx < num_heads) { - float sink_unscaled = sinks[head_idx] * 1.4426950408889634f / softmax_scale_log2; - float _max_2 = max_noftz(row_max, sink_unscaled); - row_max = _max_2; - } - bool empty_row = row_max == -CAKE_INF; - float safe_max = ((empty_row) ? 0.0f : row_max); - float max_scaled = safe_max * softmax_scale_log2; - float2 _f2_0 = make_float2(softmax_scale_log2, softmax_scale_log2); - float2 _f2_1 = make_float2(-max_scaled, -max_scaled); - float2 _f2_2 = make_float2(0.0f, 0.0f); - float2 sum0 = _f2_2; - float2 _f2_3 = make_float2(0.0f, 0.0f); - float2 sum1 = _f2_3; - float2 _f2_4 = make_float2(0.0f, 0.0f); - float2 sum2 = _f2_4; - float2 _f2_5 = make_float2(0.0f, 0.0f); - float2 sum3 = _f2_5; - #pragma unroll - for (int i = 0; i < 64; i += 8) { - float2 _f2_6 = make_float2(_tmem_load_0[i], _tmem_load_0[i + 1]); - float2 _f2_7 = make_float2(_tmem_load_0[i + 2], _tmem_load_0[i + 3]); - float2 _f2_8 = make_float2(_tmem_load_0[i + 4], _tmem_load_0[i + 5]); - float2 _f2_9 = make_float2(_tmem_load_0[i + 6], _tmem_load_0[i + 7]); - float2 affine0 = fma_f32x2_rn_ftz(_f2_6, _f2_0, _f2_1); - float2 affine1 = fma_f32x2_rn_ftz(_f2_7, _f2_0, _f2_1); - float2 affine2 = fma_f32x2_rn_ftz(_f2_8, _f2_0, _f2_1); - float2 affine3 = fma_f32x2_rn_ftz(_f2_9, _f2_0, _f2_1); - float _exp2_0 = approx_exp2(affine0.x); - float exp0 = _exp2_0; - float _exp2_1 = approx_exp2(affine0.y); - float exp1 = _exp2_1; - float _exp2_2 = approx_exp2(affine1.x); - float exp2 = _exp2_2; - float _exp2_3 = approx_exp2(affine1.y); - float exp3 = _exp2_3; - float _exp2_4 = approx_exp2(affine2.x); - float exp4 = _exp2_4; - float _exp2_5 = approx_exp2(affine2.y); - float exp5 = _exp2_5; - float _exp2_6 = approx_exp2(affine3.x); - float exp6 = _exp2_6; - float _exp2_7 = approx_exp2(affine3.y); - float exp7 = _exp2_7; - _tmem_load_0[i] = exp0; - _tmem_load_0[i + 1] = exp1; - _tmem_load_0[i + 2] = exp2; - _tmem_load_0[i + 3] = exp3; - _tmem_load_0[i + 4] = exp4; - _tmem_load_0[i + 5] = exp5; - _tmem_load_0[i + 6] = exp6; - _tmem_load_0[i + 7] = exp7; - float2 _f2_10 = make_float2(exp0, exp1); - sum0 = add_f32x2(sum0, _f2_10); - float2 _f2_11 = make_float2(exp2, exp3); - sum1 = add_f32x2(sum1, _f2_11); - float2 _f2_12 = make_float2(exp4, exp5); - sum2 = add_f32x2(sum2, _f2_12); - float2 _f2_13 = make_float2(exp6, exp7); - sum3 = add_f32x2(sum3, _f2_13); - } - float2 sum01 = add_f32x2(sum0, sum1); - float2 sum23 = add_f32x2(sum2, sum3); - float2 row_sum_pair = add_f32x2(sum01, sum23); - float row_sum = row_sum_pair.x + row_sum_pair.y; - float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, row_sum, 16); - row_sum = row_sum + _shfl_xor_1; - if (has_sinks != 0 && head_idx < num_heads) { - float _exp2_8 = approx_exp2(sinks[head_idx] * 1.4426950408889634f - max_scaled); - row_sum = row_sum + _exp2_8; - } - unsigned int packed_p[32]; - #pragma unroll - for (int _lp = 0; _lp < 32; _lp++) { - __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_0[_lp*2 + 0], _tmem_load_0[_lp*2+1 + 0])); - packed_p[_lp] = *(uint32_t*)&_bf2; - } - int p_addr = taddr + 128 + (unsigned int)(tmem_row_origin << 16); - asm volatile( - "tcgen05.st.sync.aligned.16x32bx2.x16.b32" - " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" - :: "r"(p_addr), "r"(*reinterpret_cast(&packed_p[0])), "r"(*reinterpret_cast(&packed_p[1])), "r"(*reinterpret_cast(&packed_p[2])), "r"(*reinterpret_cast(&packed_p[3])), "r"(*reinterpret_cast(&packed_p[4])), "r"(*reinterpret_cast(&packed_p[5])), "r"(*reinterpret_cast(&packed_p[6])), "r"(*reinterpret_cast(&packed_p[7])), "r"(*reinterpret_cast(&packed_p[8])), "r"(*reinterpret_cast(&packed_p[9])), "r"(*reinterpret_cast(&packed_p[10])), "r"(*reinterpret_cast(&packed_p[11])), "r"(*reinterpret_cast(&packed_p[12])), "r"(*reinterpret_cast(&packed_p[13])), "r"(*reinterpret_cast(&packed_p[14])), "r"(*reinterpret_cast(&packed_p[15]))); - asm volatile( - "tcgen05.st.sync.aligned.16x32bx2.x16.b32" - " [%0], 32, {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" - :: "r"(p_addr + 16), "r"(*reinterpret_cast(&(packed_p + 16)[0])), "r"(*reinterpret_cast(&(packed_p + 16)[1])), "r"(*reinterpret_cast(&(packed_p + 16)[2])), "r"(*reinterpret_cast(&(packed_p + 16)[3])), "r"(*reinterpret_cast(&(packed_p + 16)[4])), "r"(*reinterpret_cast(&(packed_p + 16)[5])), "r"(*reinterpret_cast(&(packed_p + 16)[6])), "r"(*reinterpret_cast(&(packed_p + 16)[7])), "r"(*reinterpret_cast(&(packed_p + 16)[8])), "r"(*reinterpret_cast(&(packed_p + 16)[9])), "r"(*reinterpret_cast(&(packed_p + 16)[10])), "r"(*reinterpret_cast(&(packed_p + 16)[11])), "r"(*reinterpret_cast(&(packed_p + 16)[12])), "r"(*reinterpret_cast(&(packed_p + 16)[13])), "r"(*reinterpret_cast(&(packed_p + 16)[14])), "r"(*reinterpret_cast(&(packed_p + 16)[15]))); - asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); - if (col_half == 0) { - float _max_3 = max_noftz(row_sum, 1.1754943508222875e-38f); - float published_sum = _max_3; - smem_sum[my_row] = published_sum; - } - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - mbarrier_arrive(p_full_addr); - mbarrier_arrive(sum_ready_addr); - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(tmem_dealloc_addr); - } - } - // ---- Role: epilogue ---- - if (warp >= 4 && warp <= 7) { - asm volatile("setmaxnreg.dec.sync.aligned.u32 128;"); - { // epilogue_main - float output_scale = bmm2_scale[0]; - int work_idx = blockIdx.x; - int query_head_work_1 = work_idx >> 2; - int query_idx_1 = query_head_work_1 / num_head_tiles; - int head_tile_1 = query_head_work_1 % num_head_tiles; - int v_chunk = work_idx & 3; - const int warp_in_role = warp - 4; - const int tmem_row_origin_1 = warp_in_role * 32; - const int logical_row_origin_1 = warp_in_role * 16; - const int my_row_1 = logical_row_origin_1 + lane % 16; - int head_idx_1 = head_tile_1 * 64 + my_row_1; - const int col_half_1 = lane / 16; - const int row_addr = tmem_row_origin_1 << 16; - unsigned int _phase_o_done_0 = 0; - mbarrier_wait(o_done_addr, _phase_o_done_0); - _phase_o_done_0 ^= 1; - unsigned int _phase_sum_ready_0 = 0; - mbarrier_wait(sum_ready_addr, _phase_sum_ready_0); - _phase_sum_ready_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - float row_sum_1 = smem_sum[my_row_1]; - float _rcp_0 = approx_rcp(row_sum_1); - float inv_sum = _rcp_0; - int output_base = (query_idx_1 * num_heads + head_idx_1) * 512 + v_chunk * 128; - float _tmem_load_1[64]; - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_1[0])), "=r"(*reinterpret_cast(&_tmem_load_1[1])), "=r"(*reinterpret_cast(&_tmem_load_1[2])), "=r"(*reinterpret_cast(&_tmem_load_1[3])), "=r"(*reinterpret_cast(&_tmem_load_1[4])), "=r"(*reinterpret_cast(&_tmem_load_1[5])), "=r"(*reinterpret_cast(&_tmem_load_1[6])), "=r"(*reinterpret_cast(&_tmem_load_1[7])), "=r"(*reinterpret_cast(&_tmem_load_1[8])), "=r"(*reinterpret_cast(&_tmem_load_1[9])), "=r"(*reinterpret_cast(&_tmem_load_1[10])), "=r"(*reinterpret_cast(&_tmem_load_1[11])), "=r"(*reinterpret_cast(&_tmem_load_1[12])), "=r"(*reinterpret_cast(&_tmem_load_1[13])), "=r"(*reinterpret_cast(&_tmem_load_1[14])), "=r"(*reinterpret_cast(&_tmem_load_1[15])), "=r"(*reinterpret_cast(&_tmem_load_1[16])), "=r"(*reinterpret_cast(&_tmem_load_1[17])), "=r"(*reinterpret_cast(&_tmem_load_1[18])), "=r"(*reinterpret_cast(&_tmem_load_1[19])), "=r"(*reinterpret_cast(&_tmem_load_1[20])), "=r"(*reinterpret_cast(&_tmem_load_1[21])), "=r"(*reinterpret_cast(&_tmem_load_1[22])), "=r"(*reinterpret_cast(&_tmem_load_1[23])), "=r"(*reinterpret_cast(&_tmem_load_1[24])), "=r"(*reinterpret_cast(&_tmem_load_1[25])), "=r"(*reinterpret_cast(&_tmem_load_1[26])), "=r"(*reinterpret_cast(&_tmem_load_1[27])), "=r"(*reinterpret_cast(&_tmem_load_1[28])), "=r"(*reinterpret_cast(&_tmem_load_1[29])), "=r"(*reinterpret_cast(&_tmem_load_1[30])), "=r"(*reinterpret_cast(&_tmem_load_1[31])) - : "r"(taddr + 256 + (unsigned int)row_addr)); - asm volatile( - "tcgen05.ld.sync.aligned.16x32bx2.x32.b32" - " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32], 64;" - : "=r"(*reinterpret_cast(&_tmem_load_1[32])), "=r"(*reinterpret_cast(&_tmem_load_1[33])), "=r"(*reinterpret_cast(&_tmem_load_1[34])), "=r"(*reinterpret_cast(&_tmem_load_1[35])), "=r"(*reinterpret_cast(&_tmem_load_1[36])), "=r"(*reinterpret_cast(&_tmem_load_1[37])), "=r"(*reinterpret_cast(&_tmem_load_1[38])), "=r"(*reinterpret_cast(&_tmem_load_1[39])), "=r"(*reinterpret_cast(&_tmem_load_1[40])), "=r"(*reinterpret_cast(&_tmem_load_1[41])), "=r"(*reinterpret_cast(&_tmem_load_1[42])), "=r"(*reinterpret_cast(&_tmem_load_1[43])), "=r"(*reinterpret_cast(&_tmem_load_1[44])), "=r"(*reinterpret_cast(&_tmem_load_1[45])), "=r"(*reinterpret_cast(&_tmem_load_1[46])), "=r"(*reinterpret_cast(&_tmem_load_1[47])), "=r"(*reinterpret_cast(&_tmem_load_1[48])), "=r"(*reinterpret_cast(&_tmem_load_1[49])), "=r"(*reinterpret_cast(&_tmem_load_1[50])), "=r"(*reinterpret_cast(&_tmem_load_1[51])), "=r"(*reinterpret_cast(&_tmem_load_1[52])), "=r"(*reinterpret_cast(&_tmem_load_1[53])), "=r"(*reinterpret_cast(&_tmem_load_1[54])), "=r"(*reinterpret_cast(&_tmem_load_1[55])), "=r"(*reinterpret_cast(&_tmem_load_1[56])), "=r"(*reinterpret_cast(&_tmem_load_1[57])), "=r"(*reinterpret_cast(&_tmem_load_1[58])), "=r"(*reinterpret_cast(&_tmem_load_1[59])), "=r"(*reinterpret_cast(&_tmem_load_1[60])), "=r"(*reinterpret_cast(&_tmem_load_1[61])), "=r"(*reinterpret_cast(&_tmem_load_1[62])), "=r"(*reinterpret_cast(&_tmem_load_1[63])) - : "r"(taddr + 256 + (unsigned int)row_addr + 32)); - int col_base = col_half_1 * 64; - if (head_idx_1 < num_heads) { - #pragma unroll - for (int offset = 0; offset < 64; offset += 8) { - { - const float2 _prescale2_0 = {inv_sum * output_scale, inv_sum * output_scale}; - #if __CUDA_ARCH__ >= 1000 - #pragma unroll - for (int _ps = 0; _ps < 4; _ps++) - mul_f32x2_inplace(&reinterpret_cast(&_tmem_load_1[offset])[_ps], _prescale2_0); - #else - #pragma unroll - for (int _ps = 0; _ps < 8; _ps++) - _tmem_load_1[offset + _ps] *= inv_sum * output_scale; - #endif - __nv_bfloat162 _pk[4]; - _pk[0] = __floats2bfloat162_rn(_tmem_load_1[offset + 0], _tmem_load_1[offset + 1]); - _pk[1] = __floats2bfloat162_rn(_tmem_load_1[offset + 2], _tmem_load_1[offset + 3]); - _pk[2] = __floats2bfloat162_rn(_tmem_load_1[offset + 4], _tmem_load_1[offset + 5]); - _pk[3] = __floats2bfloat162_rn(_tmem_load_1[offset + 6], _tmem_load_1[offset + 7]); - *reinterpret_cast(&((__nv_bfloat16*)(O + (output_base + col_base + offset)))[0]) = *reinterpret_cast(&_pk[0]); - } - } - } - asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); - asm volatile("tcgen05.fence::before_thread_sync;"); - mbarrier_arrive(tmem_dealloc_addr); - } - } - // ---- Role: mma_warp ---- - if (warp == 8) { - { // mma_warp_main - unsigned int _phase_q_full_0 = 0; - mbarrier_wait(q_full_addr, _phase_q_full_0); - _phase_q_full_0 ^= 1; - unsigned int _phase_kv_full = 0; - #pragma unroll - for (int k_stage = 0; k_stage < 4; k_stage++) { - mbarrier_wait(kv_full_addr + (k_stage) * 8, _phase_kv_full); - asm volatile("tcgen05.fence::after_thread_sync;"); - int _mma_a_lo_0 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (k_stage * 2) * 512); - int _mma_b_lo_0 = make_warp_uniform((((smem_kv_addr) >> 4) & 0x3FFF) + (k_stage) * 2048); - { - uint64_t _mma_ss_a_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_0); - uint64_t _mma_ss_b_desc_0 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_0); - if (elect_sync()) { - tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, ((k_stage == 0) ? 0 : 1)); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_0, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_0, 2U); - if (elect_sync()) { - tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_0, _mma_ss_b_desc_0, 69207184, 1); - } - } - int _mma_a_lo_1 = make_warp_uniform((((smem_q_addr) >> 4) & 0x3FFF) + (k_stage * 2 + 1) * 512); - int _mma_b_lo_1 = make_warp_uniform((((smem_kv_addr + 16384) >> 4) & 0x3FFF) + (k_stage) * 2048); - { - uint64_t _mma_ss_a_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_a_lo_1); - uint64_t _mma_ss_b_desc_1 = (static_cast(0x40004040U) << 32) | static_cast(_mma_b_lo_1); - if (elect_sync()) { - tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); - } - incr_smem_desc_lo(_mma_ss_a_desc_1, 2U); - incr_smem_desc_lo(_mma_ss_b_desc_1, 2U); - if (elect_sync()) { - tcgen05_mma_f16(tmem_tmem, _mma_ss_a_desc_1, _mma_ss_b_desc_1, 69207184, 1); - } - } - if (k_stage == 3) { - elect_commit(s_full_addr); - } - elect_commit(kv_empty_addr + (k_stage) * 8); - } - int _mma_prepared_b_lo_2 = make_warp_uniform(((((smem_v_addr) >> 4) & 0x3FFF) | 0x4000000) + (0) * 2048); - unsigned int _phase_p_full_0 = 0; - mbarrier_wait(p_full_addr, _phase_p_full_0); - _phase_p_full_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - mbarrier_wait(kv_full_addr, 1); - asm volatile("tcgen05.fence::after_thread_sync;"); - asm volatile( - "{\n\t" - ".reg .pred leader, p0, p1;\n\t" - ".reg .b32 dhi, blo, ta, id;\n\t" - ".reg .b64 db;\n\t" - "elect.sync _|leader, 0xFFFFFFFF;\n\t" - "setp.ne.b32 p0, %3, 0;\n\t" - "setp.ne.b32 p1, 1, 0;\n\t" - "" - "mov.b32 dhi, 0x40004040;\n\t" - "mov.b32 id, 69272720;\n\t" - "mov.b32 ta, %2;\n\t" - "mov.b32 blo, %1;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p0;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "add.u32 ta, ta, 8;\n\t" - "add.u32 blo, blo, 128;\n\t" - "mov.b64 db, {blo, dhi};\n\t" - "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [ta], db, id, p1;\n\t" - "}\n" - :: "r"((tmem_tmem + (256))), "r"(_mma_prepared_b_lo_2), "r"(tmem_tmem + 128), "r"(0)); - elect_commit(o_done_addr); - elect_commit(kv_empty_addr); - unsigned int _phase_tmem_dealloc_0 = 0; - mbarrier_wait(tmem_dealloc_addr, _phase_tmem_dealloc_0); - _phase_tmem_dealloc_0 ^= 1; - asm volatile("tcgen05.fence::after_thread_sync;"); - int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); - asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(512)); - } - } - // ---- Role: empty ---- - if (warp >= 9 && warp <= 11) { - // idle — no tasks assigned - } - // ---- Role: load_warp ---- - if (warp >= 12 && warp <= 15) { - asm volatile("setmaxnreg.dec.sync.aligned.u32 64;"); - { // load_warp_main - const int load_warp_rank = warp - 12; - int work_idx_1 = blockIdx.x; - int query_head_work_2 = work_idx_1 >> 2; - int query_idx_2 = query_head_work_2 / num_head_tiles; - int head_tile_2 = query_head_work_2 % num_head_tiles; - int v_chunk_1 = work_idx_1 & 3; - int* row_ptr = swa_indices + (query_idx_2 * swa_index_stride); - int index_offset = load_warp_rank * 32 + lane; - int sparse_row = -1; - if (index_offset < sparse_topk) { - sparse_row = row_ptr[index_offset]; - } - if (load_warp_rank == 0) { - if (elect_sync()) { - mbarrier_arrive_expect_tx(q_full_addr, 65536); - #pragma unroll - for (int q_stage = 0; q_stage < 8; q_stage++) { - tma_4d_gmem2smem(smem_q_addr + (unsigned int)(q_stage * 8192), (&tmap_q), 0, head_tile_2 * 64, q_stage, query_idx_2, q_full_addr); - } - } - } - smem_indices[index_offset] = sparse_row; - unsigned int _vote_0 = __ballot_sync(0xFFFFFFFF, sparse_row < 0); - unsigned int invalid_indices = _vote_0; - int valid_prefix = 32; - if (invalid_indices != 0) { - int _ffs_0 = __ffs(invalid_indices); - valid_prefix = _ffs_0 - 1; - } - if (elect_sync()) { - smem_index_flags[load_warp_rank] = valid_prefix; - } - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); - asm volatile("barrier.sync 8, 128;" ::: "memory"); - if ((load_warp_rank & 1) != 0) { - if (elect_sync()) { - int lower_prefix = smem_index_flags[load_warp_rank - 1]; - int upper_prefix = smem_index_flags[load_warp_rank]; - smem_index_flags[load_warp_rank] = ((lower_prefix < 32) ? lower_prefix : 32 + upper_prefix); - } - } - unsigned int _phase_kv_empty = 1; - #pragma unroll - for (int k_stage_1 = 0; k_stage_1 < 4; k_stage_1++) { - mbarrier_wait(kv_empty_addr + (k_stage_1) * 8, _phase_kv_empty); - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (k_stage_1) * 8, 8192); - } - #pragma unroll - for (int local_group = 0; local_group < 8; local_group++) { - int dst_k = smem_kv_addr + (unsigned int)(k_stage_1 * 32768); - int group = load_warp_rank * 8 + local_group; - int group_offset = group * 4; - int raw_rows[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows[0])), "=r"(*reinterpret_cast(&raw_rows[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows[(0) + 3])) - : "r"(smem_indices_addr + (unsigned int)(group_offset * 4))); - int raw0 = raw_rows[0]; - int raw1 = raw_rows[1]; - int raw2 = raw_rows[2]; - int raw3 = raw_rows[3]; - int row0 = ((raw0 >= 0) ? raw0 : 0); - int row1 = ((raw1 >= 0) ? raw1 : 0); - int row2 = ((raw2 >= 0) ? raw2 : 0); - int row3 = ((raw3 >= 0) ? raw3 : 0); - if (elect_sync()) { - tma_gather4_gmem2smem(dst_k + group * 512, (&tmap_swa_kv), k_stage_1 * 128, row0, row1, row2, row3, kv_full_addr + (k_stage_1) * 8); - tma_gather4_gmem2smem(dst_k + 16384 + group * 512, (&tmap_swa_kv), k_stage_1 * 128 + 64, row0, row1, row2, row3, kv_full_addr + (k_stage_1) * 8); - } - } - } - #pragma unroll - for (int v_stage = 0; v_stage < 1; v_stage++) { - mbarrier_wait(kv_empty_addr + (v_stage) * 8, 0); - if (elect_sync()) { - mbarrier_arrive_expect_tx(kv_full_addr + (v_stage) * 8, 8192); - } - #pragma unroll - for (int local_group_1 = 0; local_group_1 < 8; local_group_1++) { - int dst_v = smem_v_addr + (unsigned int)(v_stage * 32768); - int v_col = v_chunk_1 * 128 + v_stage * 64; - int group_1 = load_warp_rank * 8 + local_group_1; - int group_offset_1 = group_1 * 4; - int raw_rows_1[4]; - asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" - : "=r"(*reinterpret_cast(&raw_rows_1[0])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 1])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 2])), "=r"(*reinterpret_cast(&raw_rows_1[(0) + 3])) - : "r"(smem_indices_addr + (unsigned int)(group_offset_1 * 4))); - int raw0_1 = raw_rows_1[0]; - int raw1_1 = raw_rows_1[1]; - int raw2_1 = raw_rows_1[2]; - int raw3_1 = raw_rows_1[3]; - int row0_1 = ((raw0_1 >= 0) ? raw0_1 : 0); - int row1_1 = ((raw1_1 >= 0) ? raw1_1 : 0); - int row2_1 = ((raw2_1 >= 0) ? raw2_1 : 0); - int row3_1 = ((raw3_1 >= 0) ? raw3_1 : 0); - if (elect_sync()) { - tma_gather4_gmem2smem(dst_v + group_1 * 512, (&tmap_swa_kv), v_col, row0_1, row1_1, row2_1, row3_1, kv_full_addr + (v_stage) * 8); - tma_gather4_gmem2smem(dst_v + 16384 + group_1 * 512, (&tmap_swa_kv), v_col + 64, row0_1, row1_1, row2_1, row3_1, kv_full_addr + (v_stage) * 8); - } - } - } - } - } - - // Cleanup -} - -} // extern "C" diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_family_956aa0f74f4a82042f2c_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_family_956aa0f74f4a82042f2c_binding.cu deleted file mode 100644 index b5b1f4e43a8..00000000000 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_family_956aa0f74f4a82042f2c_binding.cu +++ /dev/null @@ -1,652 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -// Generated by cake.runtime.ffi_dispatch — do not edit. -// tvm-ffi dispatch module for kernel family 'cake_dsv4_family_956aa0f74f4a82042f2c'. -// dispatch-policy-revision: unpinned -#include -#include -#include -#include - -#include "tvm_ffi_utils.h" - -#include -#include -#include -#include -#include -#include -#include -#include - -extern "C" __global__ void kernel_cake_dsv4_230143c89d42c6c89387(__nv_bfloat16* __restrict__ partial_O, float* __restrict__ partial_lse, __nv_bfloat16* __restrict__ O, int num_q_heads, int num_split); -extern "C" __global__ void kernel_cake_dsv4_125332c2066d89bd7610(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); - -namespace cake_host_shim_3743d63f7b758ba5 { - -using tvm::ffi::Optional; -using tvm::ffi::TensorView; - -class ScopedCudaDevice { - public: - explicit ScopedCudaDevice(int device_id) { - cudaError_t error = cudaGetDevice(&previous_device_); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed before host-shim launch: cudaError=" - << static_cast(error); - if (previous_device_ != device_id) { - error = cudaSetDevice(device_id); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaSetDevice failed before host-shim launch for cuda:" - << device_id << ": cudaError=" << static_cast(error); - restore_ = true; - } - } - - ScopedCudaDevice(const ScopedCudaDevice&) = delete; - ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; - - ~ScopedCudaDevice() noexcept { - if (restore_) { - (void)cudaSetDevice(previous_device_); - } - } - - private: - int previous_device_ = -1; - bool restore_ = false; -}; - -inline int64_t CakeDeviceMultiprocessorCount(int device_id) { - constexpr int kMaxCachedCudaDevices = 64; - TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) - << "physical-SM-count cache does not cover cuda:" << device_id; - static std::atomic count_by_device[kMaxCachedCudaDevices]{}; - int cached = count_by_device[device_id].load(std::memory_order_acquire); - if (cached > 0) return cached; - - int count = 0; - cudaError_t error = cudaDeviceGetAttribute( - &count, cudaDevAttrMultiProcessorCount, device_id); - TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) - << "querying multiProcessorCount failed for cuda:" << device_id - << ": cudaError=" << static_cast(error) << ", count=" << count; - // Concurrent first launches may repeat the immutable device query, but all - // publication is atomic and every later hot-path lookup is one acquire load. - count_by_device[device_id].store(count, std::memory_order_release); - return count; -} - -inline void CheckCudaTensor(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) - << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; -} - -inline void CheckSameCudaDevice( - const TensorView& t, - const TensorView& reference, - const char* name, - const char* reference_name) { - TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) - << name << " must be on the same CUDA device as " << reference_name - << ": got cuda:" << t.device().device_id - << " versus cuda:" << reference.device().device_id; -} - -inline void CheckCurrentCudaDevice( - const TensorView& reference, - const char* reference_name) { - int current_device = -1; - cudaError_t error = cudaGetDevice(¤t_device); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed while validating " << reference_name - << ": cudaError=" << static_cast(error); - TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) - << "current CUDA device must match " << reference_name - << ": current=cuda:" << current_device - << ", tensor=cuda:" << reference.device().device_id; -} - -inline void CheckContiguous(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; -} - -inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { - DLDataType d = t.dtype(); - TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) - << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits - << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits - << ", lanes=" << (int)d.lanes << ")"; -} - -// A logical axis.outer(trailing) folds every source dim above the trailing -// dimensions. Shape products are independent of physical strides, so verify -// the leading dimensions form one dense row-major chain instead of inventing -// a "folded stride". The descriptor reads its exact adjacent physical step -// separately through stride[-(trailing + 1)]. -inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { - TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) - << name << " cannot fold leading dimensions above " << trailing - << " trailing dims from ndim=" << t.ndim(); - int outer_last = t.ndim() - trailing - 1; - if (outer_last <= 0) { - return; - } - int64_t step = t.stride(outer_last); - TVM_FFI_CHECK(step > 0, ValueError) - << name << " physical strides must be positive"; - int64_t expected = step; - for (int axis = outer_last - 1; axis >= 0; --axis) { - expected *= t.size(axis + 1); - if (t.size(axis) > 1) { - TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) - << name << " leading dims are not physically foldable above " << trailing - << " trailing dims: stride(" << axis << ")=" << t.stride(axis) - << ", expected " << expected; - } - } -} - -namespace variant_cake_dsv4_125332c2066d89bd7610_sm_103a_aaa5b590c060 { - -// 4D TMA descriptor for buffer 'tmap_q' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 3, ValueError) - << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - int64_t d2 = t.size(t.ndim() - 2); - int64_t d3 = t.size(t.ndim() - 3); - TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) - << "TMA source 'tmap_q' trailing dims must be positive"; - TVM_FFI_CHECK(d1 % 64 == 0, ValueError) - << "TMA source 'tmap_q' extent " << d1 - << " must divide exactly by " << 64; - uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(d3)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) - << "TMA box (64, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; - int64_t carrier_stride_0 = d1; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; - int64_t carrier_stride_1 = 64; - TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; - TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; - TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; - int64_t carrier_stride_2 = (d2 * d1); - TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; - TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; - TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; - uint64_t global_strides[3] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - (uint64_t)((carrier_stride_1 * 16) / 8), - (uint64_t)((carrier_stride_2 * 16) / 8), - }; - uint32_t box_dim[4] = {64u, 64u, 1u, 1u}; - uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_swa_k' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_swa_k(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_swa_k' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_swa_k' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_swa_k' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_swa_k"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_swa_k' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_swa_k' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_k'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_swa_k') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_compressed_k' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_compressed_k(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_compressed_k' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_compressed_k' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_compressed_k' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_compressed_k"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_compressed_k' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_compressed_k' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_k'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_k') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_swa_v' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_swa_v(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_swa_v' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_swa_v' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_swa_v' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_swa_v"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_swa_v' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_swa_v' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_v'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_swa_v') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_compressed_v' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_compressed_v(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_compressed_v' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_compressed_v' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_compressed_v' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_compressed_v"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_compressed_v' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_compressed_v' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_v'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_v') failed: CUresult=" << (int)r; - return tm; -} - -void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_compressed_k, TensorView arg_tmap_swa_v, TensorView arg_tmap_compressed_v, TensorView arg_O, TensorView arg_partial_lse, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_num_query_tokens, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { - DLDevice dev = arg_tmap_q.device(); - ScopedCudaDevice device_guard(dev.device_id); - CheckCudaTensor(arg_tmap_q, "tmap_q"); - CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); - CheckContiguous(arg_tmap_q, "tmap_q"); - CheckCudaTensor(arg_tmap_swa_k, "tmap_swa_k"); - CheckDtype(arg_tmap_swa_k, "tmap_swa_k", 4, 16, 1); - CheckCudaTensor(arg_tmap_compressed_k, "tmap_compressed_k"); - CheckDtype(arg_tmap_compressed_k, "tmap_compressed_k", 4, 16, 1); - CheckCudaTensor(arg_tmap_swa_v, "tmap_swa_v"); - CheckDtype(arg_tmap_swa_v, "tmap_swa_v", 4, 16, 1); - CheckCudaTensor(arg_tmap_compressed_v, "tmap_compressed_v"); - CheckDtype(arg_tmap_compressed_v, "tmap_compressed_v", 4, 16, 1); - CheckCudaTensor(arg_O, "O"); - CheckDtype(arg_O, "O", 4, 16, 1); - CheckContiguous(arg_O, "O"); - CheckCudaTensor(arg_partial_lse, "partial_lse"); - CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); - CheckContiguous(arg_partial_lse, "partial_lse"); - CheckCudaTensor(arg_swa_indices, "swa_indices"); - CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); - CheckContiguous(arg_swa_indices, "swa_indices"); - CheckCudaTensor(arg_compressed_indices, "compressed_indices"); - CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); - CheckContiguous(arg_compressed_indices, "compressed_indices"); - CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); - CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckCudaTensor(arg_sinks, "sinks"); - CheckDtype(arg_sinks, "sinks", 2, 32, 1); - CheckContiguous(arg_sinks, "sinks"); - CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); - CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); - CheckContiguous(arg_bmm1_scale, "bmm1_scale"); - CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); - CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); - CheckContiguous(arg_bmm2_scale, "bmm2_scale"); - TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) - << "scalar 'num_heads' value " << arg_num_heads - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_num_query_tokens >= -2147483648LL && arg_num_query_tokens <= 2147483647LL, ValueError) - << "scalar 'num_query_tokens' value " << arg_num_query_tokens - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) - << "scalar 'swa_index_stride' value " << arg_swa_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) - << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) - << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) - << "scalar 'sparse_topk' value " << arg_sparse_topk - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) - << "scalar 'has_sinks' value " << arg_has_sinks - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) - << "scalar 'total_work_items' value " << arg_total_work_items - << " is outside i32 range [-2147483648, 2147483647]"; - CheckSameCudaDevice(arg_tmap_swa_k, arg_tmap_q, "tmap_swa_k", "tmap_q"); - CheckSameCudaDevice(arg_tmap_compressed_k, arg_tmap_q, "tmap_compressed_k", "tmap_q"); - CheckSameCudaDevice(arg_tmap_swa_v, arg_tmap_q, "tmap_swa_v", "tmap_q"); - CheckSameCudaDevice(arg_tmap_compressed_v, arg_tmap_q, "tmap_compressed_v", "tmap_q"); - CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); - CheckSameCudaDevice(arg_partial_lse, arg_tmap_q, "partial_lse", "tmap_q"); - CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); - CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); - CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); - CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); - CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); - CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); - TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) - << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y - << ", " << grid_z << ")"; - TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) - << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z - << ") must be divisible by cluster dims (2, 1, 1)"; - - - CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); - CUtensorMap p_tmap_swa_k = EncodeTma_tmap_swa_k(arg_tmap_swa_k); - CUtensorMap p_tmap_compressed_k = EncodeTma_tmap_compressed_k(arg_tmap_compressed_k); - CUtensorMap p_tmap_swa_v = EncodeTma_tmap_swa_v(arg_tmap_swa_v); - CUtensorMap p_tmap_compressed_v = EncodeTma_tmap_compressed_v(arg_tmap_compressed_v); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); - float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); - int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); - int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); - int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); - float* p_sinks = static_cast(arg_sinks.data_ptr()); - float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); - float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); - int32_t v_num_heads = (int32_t)arg_num_heads; - int32_t v_num_query_tokens = (int32_t)arg_num_query_tokens; - int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; - int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; - int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; - int32_t v_sparse_topk = (int32_t)arg_sparse_topk; - int32_t v_has_sinks = (int32_t)arg_has_sinks; - int32_t v_total_work_items = (int32_t)arg_total_work_items; - void* kargs[] = {&p_tmap_q, &p_tmap_swa_k, &p_tmap_compressed_k, &p_tmap_swa_v, &p_tmap_compressed_v, &p_O, &p_partial_lse, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_num_query_tokens, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_has_sinks, &v_total_work_items}; - - static std::mutex smem_mu; - static auto* smem_status_by_device = - new std::unordered_map(); - cudaError_t smem_status = cudaSuccess; - { - std::lock_guard lock(smem_mu); - auto it = smem_status_by_device->find(dev.device_id); - if (it == smem_status_by_device->end()) { - int current_device = -1; - smem_status = cudaGetDevice(¤t_device); - if (smem_status == cudaSuccess && current_device != dev.device_id) { - smem_status = cudaErrorInvalidDevice; - } - if (smem_status == cudaSuccess) { - smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_125332c2066d89bd7610), - cudaFuncAttributeMaxDynamicSharedMemorySize, - 231936); - } - smem_status_by_device->emplace(dev.device_id, smem_status); - } else { - smem_status = it->second; - } - } - TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_125332c2066d89bd7610 failed: " - << cudaGetErrorString(smem_status); - dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); - dim3 block(512u, 1u, 1u); - - // Extended direct-source launch for a linked CUDA kernel symbol. - cudaLaunchAttribute attrs[1]{}; - int n = 0; - attrs[n].id = cudaLaunchAttributeClusterDimension; - attrs[n].val.clusterDim.x = 2u; - attrs[n].val.clusterDim.y = 1u; - attrs[n].val.clusterDim.z = 1u; - ++n; - cudaLaunchConfig_t config{}; - config.gridDim = grid; - config.blockDim = block; - config.dynamicSmemBytes = 231936u; - config.stream = stream; - config.attrs = attrs; - config.numAttrs = n; - cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_125332c2066d89bd7610), kargs); - TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_125332c2066d89bd7610 failed: " - << cudaGetErrorString(launch_status); -} -} // namespace variant_cake_dsv4_125332c2066d89bd7610_sm_103a_aaa5b590c060 - -namespace variant_cake_dsv4_230143c89d42c6c89387_sm_103a_c6d3bf2db9b6 { - -void Run(TensorView arg_partial_O, TensorView arg_partial_lse, TensorView arg_O, int64_t arg_num_q_heads, int64_t arg_num_split, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { - DLDevice dev = arg_partial_O.device(); - ScopedCudaDevice device_guard(dev.device_id); - CheckCudaTensor(arg_partial_O, "partial_O"); - CheckDtype(arg_partial_O, "partial_O", 4, 16, 1); - CheckContiguous(arg_partial_O, "partial_O"); - CheckCudaTensor(arg_partial_lse, "partial_lse"); - CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); - CheckContiguous(arg_partial_lse, "partial_lse"); - CheckCudaTensor(arg_O, "O"); - CheckDtype(arg_O, "O", 4, 16, 1); - CheckContiguous(arg_O, "O"); - TVM_FFI_CHECK(arg_num_q_heads >= -2147483648LL && arg_num_q_heads <= 2147483647LL, ValueError) - << "scalar 'num_q_heads' value " << arg_num_q_heads - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_num_split >= -2147483648LL && arg_num_split <= 2147483647LL, ValueError) - << "scalar 'num_split' value " << arg_num_split - << " is outside i32 range [-2147483648, 2147483647]"; - CheckSameCudaDevice(arg_partial_lse, arg_partial_O, "partial_lse", "partial_O"); - CheckSameCudaDevice(arg_O, arg_partial_O, "O", "partial_O"); - TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) - << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y - << ", " << grid_z << ")"; - - - __nv_bfloat16* p_partial_O = static_cast<__nv_bfloat16*>(arg_partial_O.data_ptr()); - float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); - int32_t v_num_q_heads = (int32_t)arg_num_q_heads; - int32_t v_num_split = (int32_t)arg_num_split; - void* kargs[] = {&p_partial_O, &p_partial_lse, &p_O, &v_num_q_heads, &v_num_split}; - - - dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); - dim3 block(128u, 1u, 1u); - - cudaLaunchConfig_t config{}; - config.gridDim = grid; - config.blockDim = block; - config.dynamicSmemBytes = 0u; - config.stream = stream; - cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_230143c89d42c6c89387), kargs); - TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_230143c89d42c6c89387 failed: " - << cudaGetErrorString(launch_status); -} -} // namespace variant_cake_dsv4_230143c89d42c6c89387_sm_103a_c6d3bf2db9b6 - -int SelectVariant(int32_t num_heads, int32_t num_query_tokens, int32_t swa_index_stride, int32_t compressed_index_stride, int32_t sparse_topk_lens_offset, int32_t sparse_topk, int32_t has_sinks, int32_t total_work_items, int32_t num_split) { - if ((num_split == 5)) return 0; - return -1; -} - -void Dispatch(tvm::ffi::TensorView arg_Q, tvm::ffi::TensorView arg_SWA_cache, tvm::ffi::TensorView arg_compressed_KV_cache, tvm::ffi::TensorView arg_partial_O, tvm::ffi::TensorView arg_partial_lse, tvm::ffi::TensorView arg_swa_indices, tvm::ffi::TensorView arg_compressed_indices, tvm::ffi::TensorView arg_sparse_topk_lens, tvm::ffi::TensorView arg_sinks, tvm::ffi::TensorView arg_bmm1_scale, tvm::ffi::TensorView arg_bmm2_scale, tvm::ffi::TensorView arg_O, int32_t num_heads, int32_t num_query_tokens, int32_t swa_index_stride, int32_t compressed_index_stride, int32_t sparse_topk_lens_offset, int32_t sparse_topk, int32_t has_sinks, int32_t total_work_items, int32_t num_split, int64_t cuda_stream_ptr) { - CheckCudaTensor(arg_Q, "Q"); - CheckCudaTensor(arg_SWA_cache, "SWA_cache"); - CheckCudaTensor(arg_compressed_KV_cache, "compressed_KV_cache"); - CheckCudaTensor(arg_partial_O, "partial_O"); - CheckCudaTensor(arg_partial_lse, "partial_lse"); - CheckCudaTensor(arg_swa_indices, "swa_indices"); - CheckCudaTensor(arg_compressed_indices, "compressed_indices"); - CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckCudaTensor(arg_sinks, "sinks"); - CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); - CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); - CheckCudaTensor(arg_O, "O"); - cudaStream_t stream = reinterpret_cast(static_cast(cuda_stream_ptr)); - int v = SelectVariant(num_heads, num_query_tokens, swa_index_stride, compressed_index_stride, sparse_topk_lens_offset, sparse_topk, has_sinks, total_work_items, num_split); - switch (v) { - case 0: { - { - int64_t gx = (total_work_items * 2); - int64_t gy = 1; - int64_t gz = 1; - cake_host_shim_3743d63f7b758ba5::variant_cake_dsv4_125332c2066d89bd7610_sm_103a_aaa5b590c060::Run(arg_Q, arg_SWA_cache, arg_compressed_KV_cache, arg_SWA_cache, arg_compressed_KV_cache, arg_partial_O, arg_partial_lse, arg_swa_indices, arg_compressed_indices, arg_sparse_topk_lens, arg_sinks, arg_bmm1_scale, arg_bmm2_scale, num_heads, num_query_tokens, swa_index_stride, compressed_index_stride, sparse_topk_lens_offset, sparse_topk, has_sinks, total_work_items, gx, gy, gz, stream); - } - { - int64_t gx = num_query_tokens; - int64_t gy = (num_heads / 4); - int64_t gz = 1; - cake_host_shim_3743d63f7b758ba5::variant_cake_dsv4_230143c89d42c6c89387_sm_103a_c6d3bf2db9b6::Run(arg_partial_O, arg_partial_lse, arg_O, num_heads, num_split, gx, gy, gz, stream); - } - return; - } - default: - TVM_FFI_CHECK(false, ValueError) << "no dispatch route for cake_dsv4_family_956aa0f74f4a82042f2c shape (num_heads=" << num_heads << ", num_query_tokens=" << num_query_tokens << ", swa_index_stride=" << swa_index_stride << ", compressed_index_stride=" << compressed_index_stride << ", sparse_topk_lens_offset=" << sparse_topk_lens_offset << ", sparse_topk=" << sparse_topk << ", has_sinks=" << has_sinks << ", total_work_items=" << total_work_items << ", num_split=" << num_split << ")"; - } -} - -int64_t RouteOf(int32_t num_heads, int32_t num_query_tokens, int32_t swa_index_stride, int32_t compressed_index_stride, int32_t sparse_topk_lens_offset, int32_t sparse_topk, int32_t has_sinks, int32_t total_work_items, int32_t num_split) { return SelectVariant(num_heads, num_query_tokens, swa_index_stride, compressed_index_stride, sparse_topk_lens_offset, sparse_topk, has_sinks, total_work_items, num_split); } - -} // namespace cake_host_shim_3743d63f7b758ba5 - -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_3743d63f7b758ba5::Dispatch); -TVM_FFI_DLL_EXPORT_TYPED_FUNC(route_of, cake_host_shim_3743d63f7b758ba5::RouteOf); diff --git a/csrc/cake_dsv4/sm_103a/cake_dsv4_family_ec44ddd94e3be88cfc83_binding.cu b/csrc/cake_dsv4/sm_103a/cake_dsv4_family_ec44ddd94e3be88cfc83_binding.cu deleted file mode 100644 index f7090b2f731..00000000000 --- a/csrc/cake_dsv4/sm_103a/cake_dsv4_family_ec44ddd94e3be88cfc83_binding.cu +++ /dev/null @@ -1,652 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer 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. - */ - -// Generated by cake.runtime.ffi_dispatch — do not edit. -// tvm-ffi dispatch module for kernel family 'cake_dsv4_family_ec44ddd94e3be88cfc83'. -// dispatch-policy-revision: unpinned -#include -#include -#include -#include - -#include "tvm_ffi_utils.h" - -#include -#include -#include -#include -#include -#include -#include -#include - -extern "C" __global__ void kernel_cake_dsv4_397f817e2931a6cc7c92(__nv_bfloat16* __restrict__ partial_O, float* __restrict__ partial_lse, __nv_bfloat16* __restrict__ O, int num_q_heads, int num_split); -extern "C" __global__ void kernel_cake_dsv4_4c41fb8347f137d4d46f(const __grid_constant__ CUtensorMap tmap_q, const __grid_constant__ CUtensorMap tmap_swa_k, const __grid_constant__ CUtensorMap tmap_compressed_k, const __grid_constant__ CUtensorMap tmap_swa_v, const __grid_constant__ CUtensorMap tmap_compressed_v, __nv_bfloat16* __restrict__ O, float* __restrict__ partial_lse, int* __restrict__ swa_indices, int* __restrict__ compressed_indices, int* __restrict__ sparse_topk_lens, float* __restrict__ sinks, float* __restrict__ bmm1_scale, float* __restrict__ bmm2_scale, int num_heads, int num_query_tokens, int swa_index_stride, int compressed_index_stride, int sparse_topk_lens_offset, int sparse_topk, int has_sinks, int total_work_items); - -namespace cake_host_shim_9b7f43581bd6022f { - -using tvm::ffi::Optional; -using tvm::ffi::TensorView; - -class ScopedCudaDevice { - public: - explicit ScopedCudaDevice(int device_id) { - cudaError_t error = cudaGetDevice(&previous_device_); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed before host-shim launch: cudaError=" - << static_cast(error); - if (previous_device_ != device_id) { - error = cudaSetDevice(device_id); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaSetDevice failed before host-shim launch for cuda:" - << device_id << ": cudaError=" << static_cast(error); - restore_ = true; - } - } - - ScopedCudaDevice(const ScopedCudaDevice&) = delete; - ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete; - - ~ScopedCudaDevice() noexcept { - if (restore_) { - (void)cudaSetDevice(previous_device_); - } - } - - private: - int previous_device_ = -1; - bool restore_ = false; -}; - -inline int64_t CakeDeviceMultiprocessorCount(int device_id) { - constexpr int kMaxCachedCudaDevices = 64; - TVM_FFI_CHECK(device_id >= 0 && device_id < kMaxCachedCudaDevices, RuntimeError) - << "physical-SM-count cache does not cover cuda:" << device_id; - static std::atomic count_by_device[kMaxCachedCudaDevices]{}; - int cached = count_by_device[device_id].load(std::memory_order_acquire); - if (cached > 0) return cached; - - int count = 0; - cudaError_t error = cudaDeviceGetAttribute( - &count, cudaDevAttrMultiProcessorCount, device_id); - TVM_FFI_CHECK(error == cudaSuccess && count > 0, RuntimeError) - << "querying multiProcessorCount failed for cuda:" << device_id - << ": cudaError=" << static_cast(error) << ", count=" << count; - // Concurrent first launches may repeat the immutable device query, but all - // publication is atomic and every later hot-path lookup is one acquire load. - count_by_device[device_id].store(count, std::memory_order_release); - return count; -} - -inline void CheckCudaTensor(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) - << name << " must be a CUDA tensor, got device_type=" << (int)t.device().device_type; -} - -inline void CheckSameCudaDevice( - const TensorView& t, - const TensorView& reference, - const char* name, - const char* reference_name) { - TVM_FFI_CHECK(t.device().device_id == reference.device().device_id, ValueError) - << name << " must be on the same CUDA device as " << reference_name - << ": got cuda:" << t.device().device_id - << " versus cuda:" << reference.device().device_id; -} - -inline void CheckCurrentCudaDevice( - const TensorView& reference, - const char* reference_name) { - int current_device = -1; - cudaError_t error = cudaGetDevice(¤t_device); - TVM_FFI_CHECK(error == cudaSuccess, RuntimeError) - << "cudaGetDevice failed while validating " << reference_name - << ": cudaError=" << static_cast(error); - TVM_FFI_CHECK(current_device == reference.device().device_id, ValueError) - << "current CUDA device must match " << reference_name - << ": current=cuda:" << current_device - << ", tensor=cuda:" << reference.device().device_id; -} - -inline void CheckContiguous(const TensorView& t, const char* name) { - TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; -} - -inline void CheckDtype(const TensorView& t, const char* name, int code, int bits, int lanes) { - DLDataType d = t.dtype(); - TVM_FFI_CHECK((int)d.code == code && (int)d.bits == bits && (int)d.lanes == lanes, TypeError) - << name << " dtype mismatch: expected DLDataType(code=" << code << ", bits=" << bits - << ", lanes=" << lanes << "), got (code=" << (int)d.code << ", bits=" << (int)d.bits - << ", lanes=" << (int)d.lanes << ")"; -} - -// A logical axis.outer(trailing) folds every source dim above the trailing -// dimensions. Shape products are independent of physical strides, so verify -// the leading dimensions form one dense row-major chain instead of inventing -// a "folded stride". The descriptor reads its exact adjacent physical step -// separately through stride[-(trailing + 1)]. -inline void CheckDenseLeadingFold(const TensorView& t, int trailing, const char* name) { - TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) - << name << " cannot fold leading dimensions above " << trailing - << " trailing dims from ndim=" << t.ndim(); - int outer_last = t.ndim() - trailing - 1; - if (outer_last <= 0) { - return; - } - int64_t step = t.stride(outer_last); - TVM_FFI_CHECK(step > 0, ValueError) - << name << " physical strides must be positive"; - int64_t expected = step; - for (int axis = outer_last - 1; axis >= 0; --axis) { - expected *= t.size(axis + 1); - if (t.size(axis) > 1) { - TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) - << name << " leading dims are not physically foldable above " << trailing - << " trailing dims: stride(" << axis << ")=" << t.stride(axis) - << ", expected " << expected; - } - } -} - -namespace variant_cake_dsv4_4c41fb8347f137d4d46f_sm_103a_994fe0d2939f { - -// 4D TMA descriptor for buffer 'tmap_q' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_q(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 3, ValueError) - << "TMA source 'tmap_q' must have at least 3 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_q' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - int64_t d2 = t.size(t.ndim() - 2); - int64_t d3 = t.size(t.ndim() - 3); - TVM_FFI_CHECK(d1 > 0 && d2 > 0 && d3 > 0, ValueError) - << "TMA source 'tmap_q' trailing dims must be positive"; - TVM_FFI_CHECK(d1 % 64 == 0, ValueError) - << "TMA source 'tmap_q' extent " << d1 - << " must divide exactly by " << 64; - uint64_t global_dim[4] = {(uint64_t)(64), (uint64_t)(d2), (uint64_t)((d1 / 64)), (uint64_t)(d3)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0 && global_dim[3] > 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[2] && 1u <= global_dim[3], ValueError) - << "TMA box (64, 64, 1, 1) exceeds resolved global dims for 'tmap_q'"; - int64_t carrier_stride_0 = d1; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 1 to a non-whole-byte offset"; - int64_t carrier_stride_1 = 64; - TVM_FFI_CHECK(carrier_stride_1 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 negative"; - TVM_FFI_CHECK(carrier_stride_1 != 0 || global_dim[2] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 zero while global dimension 2 is not 1"; - TVM_FFI_CHECK((carrier_stride_1 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 2 to a non-whole-byte offset"; - int64_t carrier_stride_2 = (d2 * d1); - TVM_FFI_CHECK(carrier_stride_2 >= 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 negative"; - TVM_FFI_CHECK(carrier_stride_2 != 0 || global_dim[3] == 1, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 zero while global dimension 3 is not 1"; - TVM_FFI_CHECK((carrier_stride_2 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_q' resolved global stride 3 to a non-whole-byte offset"; - uint64_t global_strides[3] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - (uint64_t)((carrier_stride_1 * 16) / 8), - (uint64_t)((carrier_stride_2 * 16) / 8), - }; - uint32_t box_dim[4] = {64u, 64u, 1u, 1u}; - uint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (4D, 'tmap_q') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_swa_k' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_swa_k(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_swa_k' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_swa_k' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_swa_k' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_swa_k"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_swa_k' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_swa_k' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_k'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_swa_k' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_swa_k') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_compressed_k' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_compressed_k(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_compressed_k' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_compressed_k' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_compressed_k' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_compressed_k"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_compressed_k' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_compressed_k' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_k'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_compressed_k' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_k') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_swa_v' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_swa_v(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_swa_v' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_swa_v' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_swa_v' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_swa_v"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_swa_v' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_swa_v' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_swa_v'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_swa_v' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_swa_v') failed: CUresult=" << (int)r; - return tm; -} - -// 2D TMA descriptor for buffer 'tmap_compressed_v' — compiled from the -// descriptor's std.Expr global_dim/global_strides/checks record. -inline CUtensorMap EncodeTma_tmap_compressed_v(const TensorView& t) { - TVM_FFI_CHECK(t.ndim() >= 2, ValueError) - << "TMA source 'tmap_compressed_v' must have at least 2 dimensions, got ndim=" << t.ndim(); - TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) - << "TMA source 'tmap_compressed_v' must have unit innermost stride, got " << t.stride(-1); - int64_t d1 = t.size(t.ndim() - 1); - TVM_FFI_CHECK(d1 > 0, ValueError) - << "TMA source 'tmap_compressed_v' trailing dims must be positive"; - int64_t outer1 = t.numel() / (d1); - CheckDenseLeadingFold(t, 1, "tmap_compressed_v"); - int64_t s2 = t.stride(t.ndim() - 2) * 1; - TVM_FFI_CHECK(s2 > 0, ValueError) - << "TMA source 'tmap_compressed_v' physical strides must be positive"; - uint64_t global_dim[2] = {(uint64_t)(d1), (uint64_t)(outer1)}; - TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0, ValueError) - << "TMA descriptor for 'tmap_compressed_v' resolved a non-positive global dim"; - TVM_FFI_CHECK(64u <= global_dim[0] && 1u <= global_dim[1], ValueError) - << "TMA box (64, 1) exceeds resolved global dims for 'tmap_compressed_v'"; - int64_t carrier_stride_0 = s2; - TVM_FFI_CHECK(carrier_stride_0 >= 0, ValueError) - << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 negative"; - TVM_FFI_CHECK(carrier_stride_0 != 0 || global_dim[1] == 1, ValueError) - << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 zero while global dimension 1 is not 1"; - TVM_FFI_CHECK((carrier_stride_0 * 16) % 8 == 0, ValueError) - << "TMA descriptor for 'tmap_compressed_v' resolved global stride 1 to a non-whole-byte offset"; - uint64_t global_strides[1] = { - (uint64_t)((carrier_stride_0 * 16) / 8), - }; - uint32_t box_dim[2] = {64u, 1u}; - uint32_t elem_strides[2] = {1u, 1u}; - CUtensorMap tm{}; - const void* tensor_base = static_cast(t.data_ptr()) + 0u; - CUresult r = cuTensorMapEncodeTiled( - &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, const_cast(tensor_base), global_dim, global_strides, box_dim, elem_strides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, - CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TVM_FFI_CHECK(r == CUDA_SUCCESS, RuntimeError) - << "cuTensorMapEncodeTiled (2D, 'tmap_compressed_v') failed: CUresult=" << (int)r; - return tm; -} - -void Run(TensorView arg_tmap_q, TensorView arg_tmap_swa_k, TensorView arg_tmap_compressed_k, TensorView arg_tmap_swa_v, TensorView arg_tmap_compressed_v, TensorView arg_O, TensorView arg_partial_lse, TensorView arg_swa_indices, TensorView arg_compressed_indices, TensorView arg_sparse_topk_lens, TensorView arg_sinks, TensorView arg_bmm1_scale, TensorView arg_bmm2_scale, int64_t arg_num_heads, int64_t arg_num_query_tokens, int64_t arg_swa_index_stride, int64_t arg_compressed_index_stride, int64_t arg_sparse_topk_lens_offset, int64_t arg_sparse_topk, int64_t arg_has_sinks, int64_t arg_total_work_items, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { - DLDevice dev = arg_tmap_q.device(); - ScopedCudaDevice device_guard(dev.device_id); - CheckCudaTensor(arg_tmap_q, "tmap_q"); - CheckDtype(arg_tmap_q, "tmap_q", 4, 16, 1); - CheckContiguous(arg_tmap_q, "tmap_q"); - CheckCudaTensor(arg_tmap_swa_k, "tmap_swa_k"); - CheckDtype(arg_tmap_swa_k, "tmap_swa_k", 4, 16, 1); - CheckCudaTensor(arg_tmap_compressed_k, "tmap_compressed_k"); - CheckDtype(arg_tmap_compressed_k, "tmap_compressed_k", 4, 16, 1); - CheckCudaTensor(arg_tmap_swa_v, "tmap_swa_v"); - CheckDtype(arg_tmap_swa_v, "tmap_swa_v", 4, 16, 1); - CheckCudaTensor(arg_tmap_compressed_v, "tmap_compressed_v"); - CheckDtype(arg_tmap_compressed_v, "tmap_compressed_v", 4, 16, 1); - CheckCudaTensor(arg_O, "O"); - CheckDtype(arg_O, "O", 4, 16, 1); - CheckContiguous(arg_O, "O"); - CheckCudaTensor(arg_partial_lse, "partial_lse"); - CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); - CheckContiguous(arg_partial_lse, "partial_lse"); - CheckCudaTensor(arg_swa_indices, "swa_indices"); - CheckDtype(arg_swa_indices, "swa_indices", 0, 32, 1); - CheckContiguous(arg_swa_indices, "swa_indices"); - CheckCudaTensor(arg_compressed_indices, "compressed_indices"); - CheckDtype(arg_compressed_indices, "compressed_indices", 0, 32, 1); - CheckContiguous(arg_compressed_indices, "compressed_indices"); - CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckDtype(arg_sparse_topk_lens, "sparse_topk_lens", 0, 32, 1); - CheckContiguous(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckCudaTensor(arg_sinks, "sinks"); - CheckDtype(arg_sinks, "sinks", 2, 32, 1); - CheckContiguous(arg_sinks, "sinks"); - CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); - CheckDtype(arg_bmm1_scale, "bmm1_scale", 2, 32, 1); - CheckContiguous(arg_bmm1_scale, "bmm1_scale"); - CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); - CheckDtype(arg_bmm2_scale, "bmm2_scale", 2, 32, 1); - CheckContiguous(arg_bmm2_scale, "bmm2_scale"); - TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) - << "scalar 'num_heads' value " << arg_num_heads - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_num_query_tokens >= -2147483648LL && arg_num_query_tokens <= 2147483647LL, ValueError) - << "scalar 'num_query_tokens' value " << arg_num_query_tokens - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_swa_index_stride >= -2147483648LL && arg_swa_index_stride <= 2147483647LL, ValueError) - << "scalar 'swa_index_stride' value " << arg_swa_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_compressed_index_stride >= -2147483648LL && arg_compressed_index_stride <= 2147483647LL, ValueError) - << "scalar 'compressed_index_stride' value " << arg_compressed_index_stride - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk_lens_offset >= -2147483648LL && arg_sparse_topk_lens_offset <= 2147483647LL, ValueError) - << "scalar 'sparse_topk_lens_offset' value " << arg_sparse_topk_lens_offset - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_sparse_topk >= -2147483648LL && arg_sparse_topk <= 2147483647LL, ValueError) - << "scalar 'sparse_topk' value " << arg_sparse_topk - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_has_sinks >= -2147483648LL && arg_has_sinks <= 2147483647LL, ValueError) - << "scalar 'has_sinks' value " << arg_has_sinks - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_total_work_items >= -2147483648LL && arg_total_work_items <= 2147483647LL, ValueError) - << "scalar 'total_work_items' value " << arg_total_work_items - << " is outside i32 range [-2147483648, 2147483647]"; - CheckSameCudaDevice(arg_tmap_swa_k, arg_tmap_q, "tmap_swa_k", "tmap_q"); - CheckSameCudaDevice(arg_tmap_compressed_k, arg_tmap_q, "tmap_compressed_k", "tmap_q"); - CheckSameCudaDevice(arg_tmap_swa_v, arg_tmap_q, "tmap_swa_v", "tmap_q"); - CheckSameCudaDevice(arg_tmap_compressed_v, arg_tmap_q, "tmap_compressed_v", "tmap_q"); - CheckSameCudaDevice(arg_O, arg_tmap_q, "O", "tmap_q"); - CheckSameCudaDevice(arg_partial_lse, arg_tmap_q, "partial_lse", "tmap_q"); - CheckSameCudaDevice(arg_swa_indices, arg_tmap_q, "swa_indices", "tmap_q"); - CheckSameCudaDevice(arg_compressed_indices, arg_tmap_q, "compressed_indices", "tmap_q"); - CheckSameCudaDevice(arg_sparse_topk_lens, arg_tmap_q, "sparse_topk_lens", "tmap_q"); - CheckSameCudaDevice(arg_sinks, arg_tmap_q, "sinks", "tmap_q"); - CheckSameCudaDevice(arg_bmm1_scale, arg_tmap_q, "bmm1_scale", "tmap_q"); - CheckSameCudaDevice(arg_bmm2_scale, arg_tmap_q, "bmm2_scale", "tmap_q"); - TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) - << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y - << ", " << grid_z << ")"; - TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) - << "launch grid (" << grid_x << ", " << grid_y << ", " << grid_z - << ") must be divisible by cluster dims (2, 1, 1)"; - - - CUtensorMap p_tmap_q = EncodeTma_tmap_q(arg_tmap_q); - CUtensorMap p_tmap_swa_k = EncodeTma_tmap_swa_k(arg_tmap_swa_k); - CUtensorMap p_tmap_compressed_k = EncodeTma_tmap_compressed_k(arg_tmap_compressed_k); - CUtensorMap p_tmap_swa_v = EncodeTma_tmap_swa_v(arg_tmap_swa_v); - CUtensorMap p_tmap_compressed_v = EncodeTma_tmap_compressed_v(arg_tmap_compressed_v); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); - float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); - int* p_swa_indices = static_cast(arg_swa_indices.data_ptr()); - int* p_compressed_indices = static_cast(arg_compressed_indices.data_ptr()); - int* p_sparse_topk_lens = static_cast(arg_sparse_topk_lens.data_ptr()); - float* p_sinks = static_cast(arg_sinks.data_ptr()); - float* p_bmm1_scale = static_cast(arg_bmm1_scale.data_ptr()); - float* p_bmm2_scale = static_cast(arg_bmm2_scale.data_ptr()); - int32_t v_num_heads = (int32_t)arg_num_heads; - int32_t v_num_query_tokens = (int32_t)arg_num_query_tokens; - int32_t v_swa_index_stride = (int32_t)arg_swa_index_stride; - int32_t v_compressed_index_stride = (int32_t)arg_compressed_index_stride; - int32_t v_sparse_topk_lens_offset = (int32_t)arg_sparse_topk_lens_offset; - int32_t v_sparse_topk = (int32_t)arg_sparse_topk; - int32_t v_has_sinks = (int32_t)arg_has_sinks; - int32_t v_total_work_items = (int32_t)arg_total_work_items; - void* kargs[] = {&p_tmap_q, &p_tmap_swa_k, &p_tmap_compressed_k, &p_tmap_swa_v, &p_tmap_compressed_v, &p_O, &p_partial_lse, &p_swa_indices, &p_compressed_indices, &p_sparse_topk_lens, &p_sinks, &p_bmm1_scale, &p_bmm2_scale, &v_num_heads, &v_num_query_tokens, &v_swa_index_stride, &v_compressed_index_stride, &v_sparse_topk_lens_offset, &v_sparse_topk, &v_has_sinks, &v_total_work_items}; - - static std::mutex smem_mu; - static auto* smem_status_by_device = - new std::unordered_map(); - cudaError_t smem_status = cudaSuccess; - { - std::lock_guard lock(smem_mu); - auto it = smem_status_by_device->find(dev.device_id); - if (it == smem_status_by_device->end()) { - int current_device = -1; - smem_status = cudaGetDevice(¤t_device); - if (smem_status == cudaSuccess && current_device != dev.device_id) { - smem_status = cudaErrorInvalidDevice; - } - if (smem_status == cudaSuccess) { - smem_status = cudaFuncSetAttribute( - reinterpret_cast(kernel_cake_dsv4_4c41fb8347f137d4d46f), - cudaFuncAttributeMaxDynamicSharedMemorySize, - 220032); - } - smem_status_by_device->emplace(dev.device_id, smem_status); - } else { - smem_status = it->second; - } - } - TVM_FFI_CHECK(smem_status == cudaSuccess, RuntimeError) - << "cudaFuncSetAttribute for kernel_cake_dsv4_4c41fb8347f137d4d46f failed: " - << cudaGetErrorString(smem_status); - dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); - dim3 block(512u, 1u, 1u); - - // Extended direct-source launch for a linked CUDA kernel symbol. - cudaLaunchAttribute attrs[1]{}; - int n = 0; - attrs[n].id = cudaLaunchAttributeClusterDimension; - attrs[n].val.clusterDim.x = 2u; - attrs[n].val.clusterDim.y = 1u; - attrs[n].val.clusterDim.z = 1u; - ++n; - cudaLaunchConfig_t config{}; - config.gridDim = grid; - config.blockDim = block; - config.dynamicSmemBytes = 220032u; - config.stream = stream; - config.attrs = attrs; - config.numAttrs = n; - cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_4c41fb8347f137d4d46f), kargs); - TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_4c41fb8347f137d4d46f failed: " - << cudaGetErrorString(launch_status); -} -} // namespace variant_cake_dsv4_4c41fb8347f137d4d46f_sm_103a_994fe0d2939f - -namespace variant_cake_dsv4_397f817e2931a6cc7c92_sm_103a_aa3105f3b9de { - -void Run(TensorView arg_partial_O, TensorView arg_partial_lse, TensorView arg_O, int64_t arg_num_q_heads, int64_t arg_num_split, int64_t grid_x, int64_t grid_y, int64_t grid_z, cudaStream_t stream) { - DLDevice dev = arg_partial_O.device(); - ScopedCudaDevice device_guard(dev.device_id); - CheckCudaTensor(arg_partial_O, "partial_O"); - CheckDtype(arg_partial_O, "partial_O", 4, 16, 1); - CheckContiguous(arg_partial_O, "partial_O"); - CheckCudaTensor(arg_partial_lse, "partial_lse"); - CheckDtype(arg_partial_lse, "partial_lse", 2, 32, 1); - CheckContiguous(arg_partial_lse, "partial_lse"); - CheckCudaTensor(arg_O, "O"); - CheckDtype(arg_O, "O", 4, 16, 1); - CheckContiguous(arg_O, "O"); - TVM_FFI_CHECK(arg_num_q_heads >= -2147483648LL && arg_num_q_heads <= 2147483647LL, ValueError) - << "scalar 'num_q_heads' value " << arg_num_q_heads - << " is outside i32 range [-2147483648, 2147483647]"; - TVM_FFI_CHECK(arg_num_split >= -2147483648LL && arg_num_split <= 2147483647LL, ValueError) - << "scalar 'num_split' value " << arg_num_split - << " is outside i32 range [-2147483648, 2147483647]"; - CheckSameCudaDevice(arg_partial_lse, arg_partial_O, "partial_lse", "partial_O"); - CheckSameCudaDevice(arg_O, arg_partial_O, "O", "partial_O"); - TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) - << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y - << ", " << grid_z << ")"; - - - __nv_bfloat16* p_partial_O = static_cast<__nv_bfloat16*>(arg_partial_O.data_ptr()); - float* p_partial_lse = static_cast(arg_partial_lse.data_ptr()); - __nv_bfloat16* p_O = static_cast<__nv_bfloat16*>(arg_O.data_ptr()); - int32_t v_num_q_heads = (int32_t)arg_num_q_heads; - int32_t v_num_split = (int32_t)arg_num_split; - void* kargs[] = {&p_partial_O, &p_partial_lse, &p_O, &v_num_q_heads, &v_num_split}; - - - dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); - dim3 block(128u, 1u, 1u); - - cudaLaunchConfig_t config{}; - config.gridDim = grid; - config.blockDim = block; - config.dynamicSmemBytes = 0u; - config.stream = stream; - cudaError_t launch_status = cudaLaunchKernelExC( - &config, reinterpret_cast(kernel_cake_dsv4_397f817e2931a6cc7c92), kargs); - TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) - << "cudaLaunchKernelExC for kernel_cake_dsv4_397f817e2931a6cc7c92 failed: " - << cudaGetErrorString(launch_status); -} -} // namespace variant_cake_dsv4_397f817e2931a6cc7c92_sm_103a_aa3105f3b9de - -int SelectVariant(int32_t num_heads, int32_t num_query_tokens, int32_t swa_index_stride, int32_t compressed_index_stride, int32_t sparse_topk_lens_offset, int32_t sparse_topk, int32_t has_sinks, int32_t total_work_items, int32_t num_split) { - if ((num_split == 4)) return 0; - return -1; -} - -void Dispatch(tvm::ffi::TensorView arg_Q, tvm::ffi::TensorView arg_SWA_cache, tvm::ffi::TensorView arg_compressed_KV_cache, tvm::ffi::TensorView arg_partial_O, tvm::ffi::TensorView arg_partial_lse, tvm::ffi::TensorView arg_swa_indices, tvm::ffi::TensorView arg_compressed_indices, tvm::ffi::TensorView arg_sparse_topk_lens, tvm::ffi::TensorView arg_sinks, tvm::ffi::TensorView arg_bmm1_scale, tvm::ffi::TensorView arg_bmm2_scale, tvm::ffi::TensorView arg_O, int32_t num_heads, int32_t num_query_tokens, int32_t swa_index_stride, int32_t compressed_index_stride, int32_t sparse_topk_lens_offset, int32_t sparse_topk, int32_t has_sinks, int32_t total_work_items, int32_t num_split, int64_t cuda_stream_ptr) { - CheckCudaTensor(arg_Q, "Q"); - CheckCudaTensor(arg_SWA_cache, "SWA_cache"); - CheckCudaTensor(arg_compressed_KV_cache, "compressed_KV_cache"); - CheckCudaTensor(arg_partial_O, "partial_O"); - CheckCudaTensor(arg_partial_lse, "partial_lse"); - CheckCudaTensor(arg_swa_indices, "swa_indices"); - CheckCudaTensor(arg_compressed_indices, "compressed_indices"); - CheckCudaTensor(arg_sparse_topk_lens, "sparse_topk_lens"); - CheckCudaTensor(arg_sinks, "sinks"); - CheckCudaTensor(arg_bmm1_scale, "bmm1_scale"); - CheckCudaTensor(arg_bmm2_scale, "bmm2_scale"); - CheckCudaTensor(arg_O, "O"); - cudaStream_t stream = reinterpret_cast(static_cast(cuda_stream_ptr)); - int v = SelectVariant(num_heads, num_query_tokens, swa_index_stride, compressed_index_stride, sparse_topk_lens_offset, sparse_topk, has_sinks, total_work_items, num_split); - switch (v) { - case 0: { - { - int64_t gx = (total_work_items * 2); - int64_t gy = 1; - int64_t gz = 1; - cake_host_shim_9b7f43581bd6022f::variant_cake_dsv4_4c41fb8347f137d4d46f_sm_103a_994fe0d2939f::Run(arg_Q, arg_SWA_cache, arg_compressed_KV_cache, arg_SWA_cache, arg_compressed_KV_cache, arg_partial_O, arg_partial_lse, arg_swa_indices, arg_compressed_indices, arg_sparse_topk_lens, arg_sinks, arg_bmm1_scale, arg_bmm2_scale, num_heads, num_query_tokens, swa_index_stride, compressed_index_stride, sparse_topk_lens_offset, sparse_topk, has_sinks, total_work_items, gx, gy, gz, stream); - } - { - int64_t gx = num_query_tokens; - int64_t gy = num_heads; - int64_t gz = 1; - cake_host_shim_9b7f43581bd6022f::variant_cake_dsv4_397f817e2931a6cc7c92_sm_103a_aa3105f3b9de::Run(arg_partial_O, arg_partial_lse, arg_O, num_heads, num_split, gx, gy, gz, stream); - } - return; - } - default: - TVM_FFI_CHECK(false, ValueError) << "no dispatch route for cake_dsv4_family_ec44ddd94e3be88cfc83 shape (num_heads=" << num_heads << ", num_query_tokens=" << num_query_tokens << ", swa_index_stride=" << swa_index_stride << ", compressed_index_stride=" << compressed_index_stride << ", sparse_topk_lens_offset=" << sparse_topk_lens_offset << ", sparse_topk=" << sparse_topk << ", has_sinks=" << has_sinks << ", total_work_items=" << total_work_items << ", num_split=" << num_split << ")"; - } -} - -int64_t RouteOf(int32_t num_heads, int32_t num_query_tokens, int32_t swa_index_stride, int32_t compressed_index_stride, int32_t sparse_topk_lens_offset, int32_t sparse_topk, int32_t has_sinks, int32_t total_work_items, int32_t num_split) { return SelectVariant(num_heads, num_query_tokens, swa_index_stride, compressed_index_stride, sparse_topk_lens_offset, sparse_topk, has_sinks, total_work_items, num_split); } - -} // namespace cake_host_shim_9b7f43581bd6022f - -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_9b7f43581bd6022f::Dispatch); -TVM_FFI_DLL_EXPORT_TYPED_FUNC(route_of, cake_host_shim_9b7f43581bd6022f::RouteOf);