diff --git a/cpp/tensorrt_llm/CMakeLists.txt b/cpp/tensorrt_llm/CMakeLists.txt index 06aefb3886ba..afd2d3a1f415 100644 --- a/cpp/tensorrt_llm/CMakeLists.txt +++ b/cpp/tensorrt_llm/CMakeLists.txt @@ -200,6 +200,8 @@ set(TRTLLM_LINK_LIBS layers_src runtime_src testing_src + compressorKernels_src + mhcKernels_src userbuffers_src ${DECODER_SHARED_TARGET_0} ${DECODER_SHARED_TARGET_1}) diff --git a/cpp/tensorrt_llm/kernels/CMakeLists.txt b/cpp/tensorrt_llm/kernels/CMakeLists.txt index e0e498fa89eb..0c6f4f234dec 100644 --- a/cpp/tensorrt_llm/kernels/CMakeLists.txt +++ b/cpp/tensorrt_llm/kernels/CMakeLists.txt @@ -30,6 +30,8 @@ add_subdirectory(dsv3MinLatencyKernels) add_subdirectory(causalConv1d) add_subdirectory(fusedGatedRMSNormQuant) add_subdirectory(mamba2MTPSSMCache) +add_subdirectory(mhcKernels) +add_subdirectory(compressorKernels) file(GLOB_RECURSE SRC_CPP *.cpp) file(GLOB_RECURSE SRC_CU *.cu) @@ -55,6 +57,10 @@ list(FILTER SRC_CU EXCLUDE REGEX "userbuffers/.*") list(FILTER SRC_CU EXCLUDE REGEX "fusedLayernormKernels/.*") list(FILTER SRC_CU EXCLUDE REGEX "fusedGatedRMSNormQuant/.*") list(FILTER SRC_CU EXCLUDE REGEX "mamba2MTPSSMCache/.*") +list(FILTER SRC_CPP EXCLUDE REGEX "mhcKernels/.*") +list(FILTER SRC_CU EXCLUDE REGEX "mhcKernels/.*") +list(FILTER SRC_CPP EXCLUDE REGEX "compressorKernels/.*") +list(FILTER SRC_CU EXCLUDE REGEX "compressorKernels/.*") if(NOT ENABLE_MULTI_DEVICE) list(FILTER SRC_CU EXCLUDE REGEX "customAllReduceKernels*.*cu$") diff --git a/cpp/tensorrt_llm/kernels/compressorKernels/CMakeLists.txt b/cpp/tensorrt_llm/kernels/compressorKernels/CMakeLists.txt new file mode 100644 index 000000000000..46544a47b11f --- /dev/null +++ b/cpp/tensorrt_llm/kernels/compressorKernels/CMakeLists.txt @@ -0,0 +1,25 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. SPDX-License-Identifier: Apache-2.0 +# +# 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. +# + +set(SRC_CU compressorKernels.cu) + +add_library(compressorKernels_src OBJECT ${SRC_CU}) +set_property(TARGET compressorKernels_src PROPERTY POSITION_INDEPENDENT_CODE ON) +set_property(TARGET compressorKernels_src PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS + ON) +target_compile_options(compressorKernels_src + PRIVATE $<$:--use_fast_math>) diff --git a/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu b/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu new file mode 100644 index 000000000000..417ee183b479 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu @@ -0,0 +1,1763 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * 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. + */ + +// ============================================================================ +// Compressor Kernels — DeepSeek-V4 KV Cache Compression +// ============================================================================ +// +// This file implements CUDA kernels for KV cache compression in the DeepSeek-V4 +// sparse attention system. The compressor reduces sequences of input tokens +// into fewer compressed tokens via learned weighted averaging (online softmax), +// then post-processes and scatters results into a paged KV cache. +// +// Three kernels are provided: +// +// 1. pagedKvCompressKernel — Decode path (single/few new tokens per batch). +// Loads prior compressor state from paged memory, performs online softmax +// with the new token(s), writes updated state back, and emits a compressed +// output token when compress_ratio tokens have been accumulated. +// +// 2. prefillReductionKernel — Prefill path (many tokens per batch). +// Processes full chunks of compress_ratio tokens in one shot via online +// softmax reduction over the input sequence. Also saves compressor state +// for any remainder tokens that don't form a complete chunk. +// +// 3. postProcessScatterKernel — Fused post-processing + paged cache write. +// Takes compressed output tokens and applies: RMSNorm → RoPE → Hadamard +// transform → optional V4-Pro QDQ → scatter to paged KV cache. Supports +// default, FP8, and V4-Pro MXFP8/MXFP4 QDQ cache modes. +// Keeps all intermediate values in float32 registers to avoid extra DRAM +// round-trips. +// +// Vectorization strategy: +// All kernels use 128-bit vectorized loads/stores (float4 / 8×bf16). +// VEC = number of elements per thread, chosen so that NTHRD = HEAD_DIM/VEC >= 32. +// For HEAD_DIM=128, bf16: VEC=4, NTHRD=32. For HEAD_DIM=512, bf16: VEC=8, NTHRD=64. +// +// Overlap mode (compress_ratio=4): +// When enabled, state_dim = 2*head_dim and the compressor uses overlapping +// windows: each compressed output is derived from both the previous and current +// chunk of compress_ratio tokens (previous chunk → first head_dim features, +// current chunk → second head_dim features). This doubles the state stored +// per position but improves compression quality. +// +// Template parameters: +// HEAD_DIM — Head dimension (128 or 512) +// KV_SCORE_ELEM_BYTES — kv_score element size (2=bf16, 4=fp32) +// STATE_ELEM_BYTES — compressor state element size (2=bf16, 4=fp32) +// SCALE_TYPE — Output cache scale/dtype for postProcessScatterKernel +// ============================================================================ + +#include "tensorrt_llm/kernels/compressorKernels/compressorKernels.h" + +#include "tensorrt_llm/common/assert.h" +#include +#include +#include +#include +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::compressor +{ + +// ============================================================================ +// Helper functions +// ============================================================================ + +// Full-warp butterfly reductions via __shfl_xor_sync (all 32 lanes participate). +__device__ inline float warpReduceSum(float val) +{ + for (int mask = 16; mask > 0; mask >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, mask); + return val; +} + +__device__ inline float warpReduceMax(float val) +{ + for (int mask = 16; mask > 0; mask >>= 1) + val = fmaxf(val, __shfl_xor_sync(0xFFFFFFFF, val, mask)); + return val; +} + +// Runtime-dispatched element load (bf16 or fp32 → float). Used in the decode +// kernel where elem_bytes is a runtime parameter from paged state buffers. +__device__ inline float loadAsFloat(void const* base, int64_t offset, int elem_bytes) +{ + if (elem_bytes == 2) + return __bfloat162float(reinterpret_cast<__nv_bfloat16 const*>(base)[offset]); + else + return reinterpret_cast(base)[offset]; +} + +__device__ inline void storeFromFloat(void* base, int64_t offset, float val, int elem_bytes) +{ + if (elem_bytes == 2) + reinterpret_cast<__nv_bfloat16*>(base)[offset] = __float2bfloat16_rn(val); + else + reinterpret_cast(base)[offset] = val; +} + +// Bit-hack ceil(log2(x)) for x>0: equivalent to V4 reference fast_log2_ceil. +__device__ inline int fastLog2Ceil(float x) +{ + uint32_t const bits = __float_as_uint(x); + int const exp_part = static_cast((bits >> 23) & 0xFFu) - 127; + uint32_t const man_bits = bits & 0x007FFFFFu; + return exp_part + (man_bits != 0u ? 1 : 0); +} + +// Bit-hack 2^n for integer n: equivalent to V4 reference fast_pow2. +__device__ inline float fastPow2(int n) +{ + uint32_t const bits = static_cast(n + 127) << 23; + return __uint_as_float(bits); +} + +// V4-Pro fast_round_scale: 2^ceil(log2(amax * max_value_inv)). Operand order +// (multiplication vs `amax / max_value`) matches V4 byte-for-byte; using fp32 +// bit hacks avoids log2f/exp2f rounding so the resulting power-of-2 is exact. +__device__ inline float roundedPow2Scale(float amax, float max_value_inv, float min_amax) +{ + float const clamped_amax = fmaxf(amax, min_amax); + return fastPow2(fastLog2Ceil(clamped_amax * max_value_inv)); +} + +// Hardware FP4 (e2m1) round-trip via __nv_fp4_e2m1 (cuda_fp4.h). The +// constructor uses the SM100 PTX `cvt.rn.satfinite.e2m1x2.f32` cast +// (round-to-nearest-even with finite saturation), matching V4-Pro's +// Cast(FP4) byte-for-byte. Verified against the V4-Pro reference LUT +// (e2m1 levels {0, 0.5, 1, 1.5, 2, 3, 4, 6} with midpoint thresholds +// {0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0}) over 1.1M sweep inputs covering +// every exact level, every midpoint ±2 ULPs, out-of-range, subnormals, and +// dense random fp32 -- zero mismatches. The Python test reference still +// uses the explicit LUT so the kernel is checked against an independent +// software model rather than the HW cast itself. +__device__ inline uint8_t toUe8m0(float val) +{ + __nv_fp8_e8m0 out; + out.__x = __nv_cvt_float_to_e8m0(val, __NV_SATFINITE, cudaRoundPosInf); + return out.__x; +} + +__device__ inline uint8_t packE2M1x2(float lo, float hi) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + uint32_t val; + asm volatile( + "{\n" + ".reg .b8 byte0;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n" + "mov.b32 %0, {byte0, byte0, byte0, byte0};\n" + "}" + : "=r"(val) + : "f"(lo), "f"(hi)); + return static_cast(val); +#else + return 0; +#endif +} + +// Vectorized load/store types: maps byte-width to CUDA vector type. +template +struct VecType; + +template <> +struct VecType<4> +{ + using type = unsigned int; +}; // 32-bit: 2 bf16 or 1 fp32 + +template <> +struct VecType<8> +{ + using type = uint2; +}; // 64-bit: 4 bf16 or 2 fp32 + +template <> +struct VecType<16> +{ + using type = uint4; +}; // 128-bit: 8 bf16 or 4 fp32 + +// Cache scale / pre-store quantization type for postProcessScatterKernel. +// The store dtype is implied: kNone keeps the input dtype (elem_bytes +// controls bf16 vs fp32), kFP8* writes one byte per element, and +// kMXFP4Blockwise writes packed FP4 (two values per byte) plus per-32 +// UE8M0 scale bytes. +enum class CacheScaleType +{ + kNone = 0, + kFP8PerTensor = 1, // FP8 E4M3 with static scale=1.0 + kFP8Blockwise = 2, // FP8 E4M3 with per-128-element fp32 scales + kMXFP4Blockwise = 3 // packed FP4 E2M1 with per-32 UE8M0 scales +}; + +// ============================================================================ +// Decode Kernel: pagedKvCompressKernel +// +// Template: +// NEXT_N: number of new tokens per sequence in this decode step (1-4) +// +// Grid: (batch_size) — one block per batch element +// Block: (NTHRD) where NTHRD = HEAD_DIM / VEC (>= 32 threads) +// +// Algorithm per batch element: +// For each new token in the decode step: +// 1. Load existing compressor state (partial kv/score) from paged cache +// 2. Perform online softmax: accumulate new token's contribution using +// the numerically stable running max + weighted sum formulation +// 3. Write updated state back to paged cache +// 4. If compress_ratio tokens accumulated → emit compressed output, +// reset state for next compression window +// +// Each thread handles VEC contiguous elements of head_dim. In overlap mode +// (state_dim = 2*head_dim), Phase 1 iterates over 2 column halves. +// +// Memory layout: +// kv_score: [total_tokens, 2 * state_dim] — interleaved KV and score projections +// paged_kv: paged cache for compressor KV state +// paged_score: paged cache for compressor score state (with APE bias) +// output: [total_comp_tokens, head_dim] — compressed output tokens +// ============================================================================ + +// Helper: vectorized online softmax step reading from paged KV/score state. +// Loads one position's KV and score from paged memory and updates the running +// online softmax accumulators (rmax, rsum, rwsum) per element. +// APE is already baked into paged_score (added during Phase 1), so no APE +// addition is performed here. +template +__device__ __forceinline__ void decodeSoftmaxVec(void const* __restrict__ paged_kv_raw, + void const* __restrict__ paged_score_raw, + int64_t page_sd, // page_size * state_dim (in elements) + int state_dim, + int phys_kv, // physical page index for kv + int phys_sc, // physical page index for score + int blk_off, // offset within page + int kv_col_off, // column offset (0 or HEAD_DIM) + int tid, float* __restrict__ rmax, float* __restrict__ rsum, float* __restrict__ rwsum) +{ + using StateElemT = typename std::conditional::type; + using StateVecT = typename VecType::type; + + auto const* kv = reinterpret_cast(paged_kv_raw); + auto const* sc = reinterpret_cast(paged_score_raw); + + int64_t base_kv = static_cast(phys_kv) * page_sd + blk_off * state_dim + kv_col_off; + int64_t base_sc = static_cast(phys_sc) * page_sd + blk_off * state_dim + kv_col_off; + + StateVecT k_raw = reinterpret_cast(&kv[base_kv])[tid]; + StateVecT s_raw = reinterpret_cast(&sc[base_sc])[tid]; + StateElemT const* ke = reinterpret_cast(&k_raw); + StateElemT const* se = reinterpret_cast(&s_raw); + +#pragma unroll + for (int i = 0; i < VEC; i += 4) + { + float kf[4] = {static_cast(ke[i]), static_cast(ke[i + 1]), static_cast(ke[i + 2]), + static_cast(ke[i + 3])}; + // score already includes APE (added during Phase 1 store) + float sf[4] = {static_cast(se[i]), static_cast(se[i + 1]), static_cast(se[i + 2]), + static_cast(se[i + 3])}; + // Online softmax: maintain running (max, sum_exp, weighted_sum) per element. + // nm = new max, sc_f = rescale factor for old accumulators, tm = exp(score - new_max). + // Final output: rwsum / rsum = weighted average of KV values. +#pragma unroll + for (int j = 0; j < 4; j++) + { + float nm = fmaxf(rmax[i + j], sf[j]); + float sc_f = expf(rmax[i + j] - nm); + float tm = expf(sf[j] - nm); + rsum[i + j] = rsum[i + j] * sc_f + tm; + rwsum[i + j] = rwsum[i + j] * sc_f + kf[j] * tm; + rmax[i + j] = nm; + } + } +} + +template +__global__ void pagedKvCompressKernel(void const* __restrict__ kv_score_raw, float const* __restrict__ ape, + void* __restrict__ paged_kv_raw, void* __restrict__ paged_score_raw, int32_t const* __restrict__ block_table_kv, + int32_t const* __restrict__ block_table_score, void* __restrict__ output_raw, int32_t const* __restrict__ kv_lens, + int32_t const* __restrict__ cu_seq_lens, int32_t const* __restrict__ cu_kv_comp, int page_size, int max_blocks, + int out_elem_bytes) +{ + using KvScoreElemT = typename std::conditional::type; + using StateElemT = typename std::conditional::type; + // DeepSeek-V4 model configures compress_ratio to be 4 or 128. + static_assert(COMPRESS_RATIO == 4 || COMPRESS_RATIO == 128, "Unsupported COMPRESS_RATIO"); + constexpr bool IS_OVERLAP = (COMPRESS_RATIO == 4); + constexpr int ELEM_BYTES_FOR_VEC + = (KV_SCORE_ELEM_BYTES > STATE_ELEM_BYTES) ? KV_SCORE_ELEM_BYTES : STATE_ELEM_BYTES; + constexpr int MAX_VEC = 16 / ELEM_BYTES_FOR_VEC; + constexpr int VEC = (HEAD_DIM / MAX_VEC >= 32) ? MAX_VEC : (HEAD_DIM / 32); + using KvScoreVecT = typename VecType::type; + using StateVecT = typename VecType::type; + static_assert(VEC >= 4, "VEC must be >= 4 for float4 ape loads"); + + // HEAD_BLOCKS: split head_dim across blockIdx.y for better SM utilisation. + // For HD=512 and max elem size 2: NTHRD_BASE=64 → HEAD_BLOCKS=2, NTHRD_INNER=32. + // For HD=128 and max elem size 2/4: NTHRD_BASE=32 → HEAD_BLOCKS=1, NTHRD_INNER=32. + // For HD=512 and max elem size 4: NTHRD_BASE=128 → HEAD_BLOCKS=4, NTHRD_INNER=32. + constexpr int NTHRD_BASE = HEAD_DIM / VEC; + constexpr int HEAD_BLOCKS = (NTHRD_BASE > 32) ? (NTHRD_BASE / 32) : 1; + constexpr int NTHRD_INNER = NTHRD_BASE / HEAD_BLOCKS; // always <= 32 + // ELEM_PER_BLOCK: head_dim elements handled by one blockIdx.y block. + // = NTHRD_INNER * VEC = HEAD_DIM / HEAD_BLOCKS. + // Used for the multi-warp shared-memory merge layout so that each block + // only allocates storage for its own head slice, not the full HEAD_DIM. + constexpr int ELEM_PER_BLOCK = NTHRD_INNER * VEC; + + // state_dim is fully determined by template parameters. + constexpr int STATE_DIM = IS_OVERLAP ? 2 * HEAD_DIM : HEAD_DIM; + constexpr int64_t TWO_SD = 2 * STATE_DIM; + constexpr int COFF = IS_OVERLAP ? 2 : 1; + + int const tid = threadIdx.x % NTHRD_INNER; + int const warp_id = threadIdx.x / NTHRD_INNER; // 0..NUM_RED_WARPS-1 + int const batch_idx = blockIdx.x; + int const head_blk = blockIdx.y; + int const eff_tid = head_blk * NTHRD_INNER + tid; + + int const kv_len = kv_lens[batch_idx]; + int const sp = kv_len - NEXT_N; + int const in_off = cu_seq_lens[batch_idx]; + int const out_off = cu_kv_comp[batch_idx]; + int64_t const page_sd = static_cast(page_size) * STATE_DIM; + + auto const* kv_score = reinterpret_cast(kv_score_raw); + auto* paged_kv = reinterpret_cast(paged_kv_raw); + auto* paged_score = reinterpret_cast(paged_score_raw); + + // ================================================================ + // Phase 1: Write NEXT_N new tokens' KV and score state to paged cache. + // + // Only warp 0 participates (all warps share the same eff_tid mapping). + // When NUM_RED_WARPS == 1, the guard compiles away. + // ================================================================ + if (warp_id == 0) + { +#pragma unroll + for (int t = 0; t < NEXT_N; t++) + { + int token_idx = sp + t; + if (token_idx < kv_len) + { + int ape_idx = token_idx % COMPRESS_RATIO; + int log_blk = token_idx / page_size; + int blk_off = token_idx % page_size; + int phys_kv = block_table_kv[batch_idx * max_blocks + log_blk]; + int phys_sc = block_table_score[batch_idx * max_blocks + log_blk]; + + for (int col_idx = 0; col_idx < COFF; col_idx++) + { + int const col = col_idx * HEAD_DIM; + int64_t const src = static_cast(in_off + t) * TWO_SD + col; + int64_t const dkv = static_cast(phys_kv) * page_sd + blk_off * STATE_DIM + col; + int64_t const dsc = static_cast(phys_sc) * page_sd + blk_off * STATE_DIM + col; + + KvScoreVecT kv_raw = reinterpret_cast(&kv_score[src])[eff_tid]; + KvScoreVecT sc_raw = reinterpret_cast(&kv_score[src + STATE_DIM])[eff_tid]; + + KvScoreElemT const* kv_e = reinterpret_cast(&kv_raw); + StateVecT kv_out; + StateElemT* kv_o = reinterpret_cast(&kv_out); +#pragma unroll + for (int i = 0; i < VEC; i++) + { + kv_o[i] = static_cast(static_cast(kv_e[i])); + } + reinterpret_cast(&paged_kv[dkv])[eff_tid] = kv_out; + + KvScoreElemT const* sc_e = reinterpret_cast(&sc_raw); + StateVecT sc_out; + StateElemT* sc_o = reinterpret_cast(&sc_out); +#pragma unroll + for (int i = 0; i < VEC; i += 4) + { + float4 av + = *reinterpret_cast(&ape[ape_idx * STATE_DIM + col + eff_tid * VEC + i]); + sc_o[i] = static_cast(static_cast(sc_e[i]) + av.x); + sc_o[i + 1] = static_cast(static_cast(sc_e[i + 1]) + av.y); + sc_o[i + 2] = static_cast(static_cast(sc_e[i + 2]) + av.z); + sc_o[i + 3] = static_cast(static_cast(sc_e[i + 3]) + av.w); + } + reinterpret_cast(&paged_score[dsc])[eff_tid] = sc_out; + } + } + } + } + + if constexpr (NUM_RED_WARPS > 1) + { + __syncthreads(); + } + + // ================================================================ + // Phase 2: Count how many complete compression windows finished. + // ================================================================ + int last_token_idx = sp + NEXT_N - 1; + int num_compressions = (last_token_idx + 1) / COMPRESS_RATIO - sp / COMPRESS_RATIO; + + // ================================================================ + // Phase 3: Online softmax reduction over each complete chunk. + // + // When NUM_RED_WARPS > 1, the compress_ratio positions are split + // across warps. Each warp reduces its partition independently, then + // partial (rmax, rsum, rwsum) accumulators are merged via shared + // memory using the log-sum-exp identity. + // ================================================================ + for (int c = 0; c < NEXT_N; c++) + { + if (c >= num_compressions) + break; + + int compress_idx = sp / COMPRESS_RATIO + c; + int curr_chunk_start = compress_idx * COMPRESS_RATIO; + + float rmax[VEC], rsum[VEC], rwsum[VEC]; +#pragma unroll + for (int i = 0; i < VEC; i++) + { + rmax[i] = -INFINITY; + rsum[i] = 0.0f; + rwsum[i] = 0.0f; + } + + constexpr int positions_per_warp = COMPRESS_RATIO / NUM_RED_WARPS; + int const my_r_start = warp_id * positions_per_warp; + int const my_r_end = (warp_id == NUM_RED_WARPS - 1) ? COMPRESS_RATIO : (my_r_start + positions_per_warp); + + if constexpr (IS_OVERLAP) + { + int prev_start = curr_chunk_start - COMPRESS_RATIO; + if (prev_start >= 0) + { + if (page_size >= COMPRESS_RATIO) + { + int log_blk_prev = prev_start / page_size; + int phys_kv_prev = block_table_kv[batch_idx * max_blocks + log_blk_prev]; + int phys_sc_prev = block_table_score[batch_idx * max_blocks + log_blk_prev]; + int chunk_off_prev = prev_start % page_size; + for (int r = my_r_start; r < my_r_end; r++) + { + decodeSoftmaxVec(paged_kv_raw, paged_score_raw, page_sd, + STATE_DIM, phys_kv_prev, phys_sc_prev, chunk_off_prev + r, 0, eff_tid, rmax, rsum, rwsum); + } + } + else + { + for (int r = my_r_start; r < my_r_end; r++) + { + int pos = prev_start + r; + int log_blk = pos / page_size; + int blk_off = pos % page_size; + int phys_kv = block_table_kv[batch_idx * max_blocks + log_blk]; + int phys_sc = block_table_score[batch_idx * max_blocks + log_blk]; + decodeSoftmaxVec(paged_kv_raw, paged_score_raw, page_sd, + STATE_DIM, phys_kv, phys_sc, blk_off, 0, eff_tid, rmax, rsum, rwsum); + } + } + } + + if (page_size >= COMPRESS_RATIO) + { + int log_blk_cur = curr_chunk_start / page_size; + int phys_kv_cur = block_table_kv[batch_idx * max_blocks + log_blk_cur]; + int phys_sc_cur = block_table_score[batch_idx * max_blocks + log_blk_cur]; + int chunk_off_cur = curr_chunk_start % page_size; + for (int r = my_r_start; r < my_r_end; r++) + { + decodeSoftmaxVec(paged_kv_raw, paged_score_raw, page_sd, STATE_DIM, + phys_kv_cur, phys_sc_cur, chunk_off_cur + r, HEAD_DIM, eff_tid, rmax, rsum, rwsum); + } + } + else + { + for (int r = my_r_start; r < my_r_end; r++) + { + int pos = curr_chunk_start + r; + int log_blk = pos / page_size; + int blk_off = pos % page_size; + int phys_kv = block_table_kv[batch_idx * max_blocks + log_blk]; + int phys_sc = block_table_score[batch_idx * max_blocks + log_blk]; + decodeSoftmaxVec(paged_kv_raw, paged_score_raw, page_sd, STATE_DIM, + phys_kv, phys_sc, blk_off, HEAD_DIM, eff_tid, rmax, rsum, rwsum); + } + } + } + else + { + if (page_size >= COMPRESS_RATIO) + { + int log_blk = curr_chunk_start / page_size; + int phys_kv = block_table_kv[batch_idx * max_blocks + log_blk]; + int phys_sc = block_table_score[batch_idx * max_blocks + log_blk]; + int chunk_off = curr_chunk_start % page_size; + for (int r = my_r_start; r < my_r_end; r++) + { + decodeSoftmaxVec(paged_kv_raw, paged_score_raw, page_sd, STATE_DIM, + phys_kv, phys_sc, chunk_off + r, 0, eff_tid, rmax, rsum, rwsum); + } + } + else + { + for (int r = my_r_start; r < my_r_end; r++) + { + int pos = curr_chunk_start + r; + int log_blk = pos / page_size; + int blk_off = pos % page_size; + int phys_kv = block_table_kv[batch_idx * max_blocks + log_blk]; + int phys_sc = block_table_score[batch_idx * max_blocks + log_blk]; + decodeSoftmaxVec(paged_kv_raw, paged_score_raw, page_sd, STATE_DIM, + phys_kv, phys_sc, blk_off, 0, eff_tid, rmax, rsum, rwsum); + } + } + } + + // Multi-warp merge epilogue (compiled away when NUM_RED_WARPS == 1). + if constexpr (NUM_RED_WARPS > 1) + { + // Shared-memory layout: [NUM_RED_WARPS * ELEM_PER_BLOCK] per array. + // ELEM_PER_BLOCK = NTHRD_INNER * VEC = HEAD_DIM / HEAD_BLOCKS, i.e. + // the number of head_dim elements covered by this block (blockIdx.y). + // Using ELEM_PER_BLOCK (not HEAD_DIM) avoids 4x over-allocation when + // HEAD_BLOCKS > 1 (e.g. HD=512 fp32 has HEAD_BLOCKS=4). + extern __shared__ float smem[]; + float* s_rmax = smem; + float* s_rsum = s_rmax + NUM_RED_WARPS * ELEM_PER_BLOCK; + float* s_rwsum = s_rsum + NUM_RED_WARPS * ELEM_PER_BLOCK; + + // local_elem: index within this block's head slice [0, ELEM_PER_BLOCK). + // = tid * VEC + i (tid = threadIdx.x % NTHRD_INNER, same as eff_tid - head_blk*NTHRD_INNER) +#pragma unroll + for (int i = 0; i < VEC; i++) + { + int const local_elem = tid * VEC + i; + s_rmax[warp_id * ELEM_PER_BLOCK + local_elem] = rmax[i]; + s_rsum[warp_id * ELEM_PER_BLOCK + local_elem] = rsum[i]; + s_rwsum[warp_id * ELEM_PER_BLOCK + local_elem] = rwsum[i]; + } + __syncthreads(); + + if (warp_id == 0) + { + for (int w = 1; w < NUM_RED_WARPS; w++) + { +#pragma unroll + for (int i = 0; i < VEC; i++) + { + int const local_elem = tid * VEC + i; + float const m2 = s_rmax[w * ELEM_PER_BLOCK + local_elem]; + float const s2 = s_rsum[w * ELEM_PER_BLOCK + local_elem]; + float const ws2 = s_rwsum[w * ELEM_PER_BLOCK + local_elem]; + + float const nm = fmaxf(rmax[i], m2); + float const sc1 = expf(rmax[i] - nm); + float const sc2 = expf(m2 - nm); + rsum[i] = rsum[i] * sc1 + s2 * sc2; + rwsum[i] = rwsum[i] * sc1 + ws2 * sc2; + rmax[i] = nm; + } + } + } + __syncthreads(); + } + + bool const should_write = (NUM_RED_WARPS == 1) || (warp_id == 0); + if (should_write) + { + int64_t const out_base = static_cast(out_off + c) * HEAD_DIM + eff_tid * VEC; + if (out_elem_bytes == 2) + { + __nv_bfloat16 packed[VEC]; +#pragma unroll + for (int i = 0; i < VEC; i++) + packed[i] = __float2bfloat16_rn(rwsum[i] / rsum[i]); + using OutVecT = typename VecType::type; + *reinterpret_cast(&reinterpret_cast<__nv_bfloat16*>(output_raw)[out_base]) + = *reinterpret_cast(packed); + } + else + { + float result[VEC]; +#pragma unroll + for (int i = 0; i < VEC; i++) + result[i] = rwsum[i] / rsum[i]; +#pragma unroll + for (int i = 0; i < VEC; i += 4) + *reinterpret_cast(&reinterpret_cast(output_raw)[out_base + i]) + = *reinterpret_cast(&result[i]); + } + } + } +} + +// ============================================================================ +// Decode kernel configuration matrix — single source of truth. +// +// X-macro listing every supported (HD, KV_EB, STATE_EB, CR, NN, NRW) tuple. +// Both the explicit template instantiations below AND the runtime dispatcher +// in pagedKvCompressLaunch() walk this list, so adding/removing a config is +// a one-line edit. +// +// HD — HEAD_DIM in {128, 512} +// KV_EB — kv_score element bytes in {2 (bf16), 4 (fp32)} +// STATE_EB — paged state element bytes in {2 (bf16), 4 (fp32)} +// CR — COMPRESS_RATIO in {4, 128} +// NN — NEXT_N (new tokens / decode step) in {1..4} +// NRW — NUM_RED_WARPS — 4 only when CR=128 (multi-warp Phase 3 reduce +// hides DRAM latency for the heavier R=128 chunk); 1 otherwise. +// +// Multi-warp SMEM budget (per block): 3 * NRW * ELEM_PER_BLOCK * sizeof(float). +// HD=128: ELEM_PER_BLOCK=128 → 6 KB +// HD=512 bf16: ELEM_PER_BLOCK=256 → 12 KB +// HD=512 fp32: ELEM_PER_BLOCK=128 → 6 KB +// ============================================================================ + +// Per-axis fan-outs (used to keep the master list compact). +#define FOREACH_DECODE_NN(F, HD, KV, ST, CR, NRW) \ + F(HD, KV, ST, CR, 1, NRW) F(HD, KV, ST, CR, 2, NRW) F(HD, KV, ST, CR, 3, NRW) F(HD, KV, ST, CR, 4, NRW) +#define FOREACH_DECODE_DTYPE(F, HD, CR, NRW) \ + FOREACH_DECODE_NN(F, HD, 2, 2, CR, NRW) \ + FOREACH_DECODE_NN(F, HD, 2, 4, CR, NRW) \ + FOREACH_DECODE_NN(F, HD, 4, 2, CR, NRW) FOREACH_DECODE_NN(F, HD, 4, 4, CR, NRW) + +// Master list. Order does not matter; the dispatcher walks linearly. +// clang-format off +#define FOREACH_DECODE_CONFIG(F) \ + /* CR=4: single-warp only (small reduction; multi-warp would over-subscribe). */ \ + FOREACH_DECODE_DTYPE(F, 128, 4, 1) FOREACH_DECODE_DTYPE(F, 512, 4, 1) \ + /* CR=128: single-warp fallback (covers next_n>4 path which currently isn't reached). */ \ + FOREACH_DECODE_DTYPE(F, 128, 128, 1) FOREACH_DECODE_DTYPE(F, 512, 128, 1) \ + /* CR=128: multi-warp fast path. Used whenever next_n <= 4 (i.e. MTP-3 and below). */ \ + FOREACH_DECODE_DTYPE(F, 128, 128, 4) FOREACH_DECODE_DTYPE(F, 512, 128, 4) +// clang-format on + +// Generate explicit template instantiations. +#define INST_DECODE(HD, KV_EB, STATE_EB, CR, NN, NRW) \ + template __global__ void pagedKvCompressKernel(void const*, float const*, void*, \ + void*, int32_t const*, int32_t const*, void*, int32_t const*, int32_t const*, int32_t const*, int, int, int); +FOREACH_DECODE_CONFIG(INST_DECODE) +#undef INST_DECODE + +// ============================================================================ +// Decode Launch Wrapper +// +// Dispatches to the correct template instantiation based on head_dim, elem_bytes, +// and next_n (number of new tokens per decode step, capped at 4). +// Grid is 2D: (batch_size, head_blocks) where head_blocks = NTHRD_BASE / 32. +// For HD=512 bf16: head_blocks=2; for HD=128 bf16: head_blocks=1. +// ============================================================================ + +// Forward declaration (defined in prefill section below). +static inline int prefillNthreads(int head_dim, int elem_bytes_for_vec); + +void pagedKvCompressLaunch(void const* kv_score, float const* ape, void* paged_kv, void* paged_score, + int32_t const* block_table_kv, int32_t const* block_table_score, void* output, int32_t const* kv_lens, + int32_t const* cu_seq_lens, int32_t const* cu_kv_comp, int batch_size, int page_size, int max_blocks, int head_dim, + int compress_ratio, int next_n, int kv_score_elem_bytes, int state_elem_bytes, int out_elem_bytes, + cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO( + compress_ratio == 4 || compress_ratio == 128, "pagedKvCompressLaunch only supports compress_ratio 4 or 128"); + TLLM_CHECK_WITH_INFO( + (kv_score_elem_bytes == 2 || kv_score_elem_bytes == 4) && (state_elem_bytes == 2 || state_elem_bytes == 4), + "pagedKvCompressLaunch only supports bf16/fp32 kv_score and paged state"); + + // Compute HEAD_BLOCKS: mirrors the compile-time constant in the kernel. + // VEC = max_vec if HEAD_DIM/max_vec >= 32, else HEAD_DIM/32. + // NTHRD_BASE = HEAD_DIM / VEC; HEAD_BLOCKS = NTHRD_BASE / 32 (or 1 if <= 32). + // This spreads the head_dim across multiple blocks for better SM utilisation. + int const elem_bytes_for_vec = max(kv_score_elem_bytes, state_elem_bytes); + int const max_vec_elem = 16 / elem_bytes_for_vec; + int const vec = (head_dim / max_vec_elem >= 32) ? max_vec_elem : (head_dim / 32); + int const nthrd_base = head_dim / vec; // mirrors kernel's NTHRD_BASE = HEAD_DIM / VEC + int const head_blocks = (nthrd_base > 32) ? (nthrd_base / 32) : 1; + int const nthreads_inner = nthrd_base / head_blocks; // = min(32, nthrd_base) = always 32 + + // For large compress_ratio, use 4-warp parallel reduction to cut the serial + // softmax loop from COMPRESS_RATIO iterations to COMPRESS_RATIO/4 per warp. + // Supported configs: CR=128, (HD=128 or HD=512), NEXT_N in 1..4. NEXT_N>2 + // is required for MTP-3 decode (each step accepts up to 4 tokens per request); + // without multi-warp the slow path is a single warp doing 128 serial paged + // loads, which is DRAM-latency-bound (no other warps to hide it). + // + // smem per block = 3 * MULTI_WARP * ELEM_PER_BLOCK * sizeof(float) + // where ELEM_PER_BLOCK = nthreads_inner * vec = HEAD_DIM / HEAD_BLOCKS. + // HD=128: ELEM_PER_BLOCK=128 → 6 KB. + // HD=512 with max elem size 2 (vec=8, HEAD_BLOCKS=2): ELEM_PER_BLOCK=256 → 12 KB. + // HD=512 with max elem size 4 (vec=4, HEAD_BLOCKS=4): ELEM_PER_BLOCK=128 → 6 KB. + constexpr int MULTI_WARP = 4; + bool const use_multi_warp = (compress_ratio == 128 && next_n <= 4); + int const num_red_warps = use_multi_warp ? MULTI_WARP : 1; + int const nthreads = nthreads_inner * num_red_warps; + int const elem_per_block = nthreads_inner * vec; // = HEAD_DIM / HEAD_BLOCKS + int const smem_bytes = use_multi_warp ? (3 * MULTI_WARP * elem_per_block * static_cast(sizeof(float))) : 0; + + dim3 grid(batch_size, head_blocks); + + // Clamp the runtime next_n into the supported range; configs above 4 fall + // back to the NN=4 instantiation (matches the prior `default:` arm). + int const next_n_dispatch = std::min(next_n, 4); + + // Walk FOREACH_DECODE_CONFIG until we find a matching (HD, KV, ST, CR, NN, NRW) + // tuple, then launch that instantiation. Any unsupported tuple bails via TLLM_THROW. +#define TRY_LAUNCH(HD, KV_EB, STATE_EB, CR, NN, NRW) \ + if (head_dim == HD && kv_score_elem_bytes == KV_EB && state_elem_bytes == STATE_EB && compress_ratio == CR \ + && next_n_dispatch == NN && num_red_warps == NRW) \ + { \ + pagedKvCompressKernel<<>>(kv_score, ape, \ + paged_kv, paged_score, block_table_kv, block_table_score, output, kv_lens, cu_seq_lens, cu_kv_comp, \ + page_size, max_blocks, out_elem_bytes); \ + return; \ + } + FOREACH_DECODE_CONFIG(TRY_LAUNCH) +#undef TRY_LAUNCH + + TLLM_THROW( + "pagedKvCompressLaunch: no matching instantiation for HD=%d, kv_eb=%d, state_eb=%d, CR=%d, NN=%d, NRW=%d", + head_dim, kv_score_elem_bytes, state_elem_bytes, compress_ratio, next_n_dispatch, num_red_warps); +} + +#undef FOREACH_DECODE_CONFIG +#undef FOREACH_DECODE_DTYPE +#undef FOREACH_DECODE_NN + +// ============================================================================ +// Prefill Kernel: prefillReductionKernel +// +// Template: +// +// Grid: (batch_size, max_outputs_per_batch) — one block per compressed output +// Block: (NTHRD) where NTHRD = HEAD_DIM / VEC (>= 32 threads) +// +// Unlike the decode kernel (which operates token-by-token from paged state), +// the prefill kernel processes the full input sequence at once. Each block +// reads compress_ratio consecutive input rows from kv_score and reduces them +// via online softmax to produce one compressed output. +// +// The last block (local_output_idx == num_outputs - 1) also handles saving +// compressor state for any remainder tokens that don't form a full chunk. +// This state is written to paged kv/score caches for use in future decode steps. +// +// Memory layout: +// kv_score: [total_tokens, 2*state_dim] — interleaved KV and score from linear projection +// paged_kv: paged cache for compressor state (remainder) +// paged_score: paged cache for compressor score state (remainder, with APE) +// output: [total_comp_tokens, head_dim] — compressed output tokens +// ============================================================================ + +// Per-element online softmax step on VEC elements via 128-bit vectorized loads. +// Reads directly from the kv_score input buffer (not paged state) since prefill +// has the full sequence available. +template +__device__ __forceinline__ void prefillSoftmaxVec(void const* __restrict__ kv_score_raw, float const* __restrict__ ape, + int64_t row_elem, // (input_offset + row_idx) * two_sd + int kv_col_off, // column offset into kv_score row (0 or HEAD_DIM) + int ape_base, // r * state_dim + ape_col_off + int state_dim, int tid, float* __restrict__ rmax, float* __restrict__ rsum, float* __restrict__ rwsum) +{ + using KvScoreElemT = typename std::conditional::type; + using KvScoreVecT = typename VecType::type; + + auto const* kv = reinterpret_cast(kv_score_raw); + + KvScoreVecT k_raw = reinterpret_cast(&kv[row_elem + kv_col_off])[tid]; + KvScoreVecT s_raw = reinterpret_cast(&kv[row_elem + state_dim + kv_col_off])[tid]; + KvScoreElemT const* ke = reinterpret_cast(&k_raw); + KvScoreElemT const* se = reinterpret_cast(&s_raw); + +#pragma unroll + for (int i = 0; i < VEC; i += 4) + { + float4 av = *reinterpret_cast(&ape[ape_base + tid * VEC + i]); + float kf[4] = {static_cast(ke[i]), static_cast(ke[i + 1]), static_cast(ke[i + 2]), + static_cast(ke[i + 3])}; + float sf[4] = {static_cast(se[i]) + av.x, static_cast(se[i + 1]) + av.y, + static_cast(se[i + 2]) + av.z, static_cast(se[i + 3]) + av.w}; +#pragma unroll + for (int j = 0; j < 4; j++) + { + float nm = fmaxf(rmax[i + j], sf[j]); + float sc = expf(rmax[i + j] - nm); + float tm = expf(sf[j] - nm); + rsum[i + j] = rsum[i + j] * sc + tm; + rwsum[i + j] = rwsum[i + j] * sc + kf[j] * tm; + rmax[i + j] = nm; + } + } +} + +template +__global__ void prefillReductionKernel(void const* __restrict__ kv_score_raw, float const* __restrict__ ape, + void* __restrict__ paged_kv_raw, void* __restrict__ paged_score_raw, int32_t const* __restrict__ block_table_kv, + int32_t const* __restrict__ block_table_score, void* __restrict__ output_raw, int32_t const* __restrict__ kv_lens, + int32_t const* __restrict__ start_pos_arr, int32_t const* __restrict__ cu_seq_lens, + int32_t const* __restrict__ cu_kv_comp, int page_size, int state_dim, int max_blocks, int out_elem_bytes) +{ + using KvScoreElemT = typename std::conditional::type; + using StateElemT = typename std::conditional::type; + static_assert(COMPRESS_RATIO == 4 || COMPRESS_RATIO == 128, "Unsupported COMPRESS_RATIO"); + constexpr bool IS_OVERLAP = (COMPRESS_RATIO == 4); + + constexpr int ELEM_BYTES_FOR_VEC + = (KV_SCORE_ELEM_BYTES > STATE_ELEM_BYTES) ? KV_SCORE_ELEM_BYTES : STATE_ELEM_BYTES; + constexpr int MAX_VEC = 16 / ELEM_BYTES_FOR_VEC; + constexpr int VEC = (HEAD_DIM / MAX_VEC >= 32) ? MAX_VEC : (HEAD_DIM / 32); + using KvScoreVecT = typename VecType::type; + using StateVecT = typename VecType::type; + static_assert(VEC >= 4, "VEC must be >= 4 for float4 ape loads"); + + int const tid = threadIdx.x; + int const batch_idx = blockIdx.x; + int const local_output_idx = blockIdx.y; + + int const sp = start_pos_arr[batch_idx]; + int const kv_len = kv_lens[batch_idx]; + int const input_offset = cu_seq_lens[batch_idx]; + int const output_offset = cu_kv_comp[batch_idx]; + + // Absolute compression index range. Window k covers [k*R, (k+1)*R). + // We emit one output per complete window that gains at least one new token. + int const first_abs_idx = sp / COMPRESS_RATIO; + int const last_abs_idx = kv_len / COMPRESS_RATIO; // exclusive + int const actual_num_outputs = last_abs_idx - first_abs_idx; + // Keep at least one block so the last CTA can still persist remainder state + // even when this chunk has no full compression window. + int const num_outputs = max(actual_num_outputs, 1); + + if (local_output_idx >= num_outputs) + return; + + constexpr int coff = IS_OVERLAP ? 2 : 1; + bool const should_compress = (local_output_idx < actual_num_outputs); + + // Absolute window for this output block. + int const abs_idx = first_abs_idx + local_output_idx; + int const win_start = abs_idx * COMPRESS_RATIO; + + auto const* kv_score = reinterpret_cast(kv_score_raw); + auto* paged_kv = reinterpret_cast(paged_kv_raw); + auto* paged_score = reinterpret_cast(paged_score_raw); + + int64_t const two_sd = 2 * state_dim; + int64_t const page_sd = static_cast(page_size) * state_dim; + + // ================================================================ + // Phase 1: State Update (all output blocks, for block reuse support) + // + // 1a runs on every block that has a full chunk (should_compress), writing + // each block's own window to paged cache so any prefix slice is valid for + // block reuse. 1b (remainder) still runs only on the last block. + // ================================================================ + + // Helper: write a contiguous block of positions to paged KV/score cache. + // Only new positions (>= sp) are written; positions already persisted from + // prior calls are skipped by starting the loop at write_r_start. + // APE index is r (position within the compression window), matching the + // index used during the original decode/prefill that first established + // the window alignment. + auto write_to_paged = [&](int range_start, int range_end, int write_r_start) + { + for (int r = write_r_start; r < range_end - range_start; r++) + { + int const pos = range_start + r; + int const input_row = pos - sp; // >= 0 since pos >= sp (loop starts at write_r_start) + int const log_blk = pos / page_size; + int const blk_off = pos % page_size; + int const phys_kv = block_table_kv[batch_idx * max_blocks + log_blk]; + int const phys_sc = block_table_score[batch_idx * max_blocks + log_blk]; + + for (int col_idx = 0; col_idx < coff; col_idx++) + { + int const col = col_idx * HEAD_DIM; + int64_t const src = static_cast(input_offset + input_row) * two_sd + col; + int64_t const dkv = static_cast(phys_kv) * page_sd + blk_off * state_dim + col; + int64_t const dsc = static_cast(phys_sc) * page_sd + blk_off * state_dim + col; + + KvScoreVecT kv_raw = reinterpret_cast(&kv_score[src])[tid]; + KvScoreVecT sc_raw = reinterpret_cast(&kv_score[src + state_dim])[tid]; + + KvScoreElemT const* kv_e = reinterpret_cast(&kv_raw); + StateVecT kv_out; + StateElemT* kv_o = reinterpret_cast(&kv_out); +#pragma unroll + for (int i = 0; i < VEC; i++) + { + kv_o[i] = static_cast(static_cast(kv_e[i])); + } + reinterpret_cast(&paged_kv[dkv])[tid] = kv_out; + + KvScoreElemT const* sc_e = reinterpret_cast(&sc_raw); + StateVecT sc_out; + StateElemT* sc_o = reinterpret_cast(&sc_out); +#pragma unroll + for (int i = 0; i < VEC; i += 4) + { + float4 av = *reinterpret_cast(&ape[r * state_dim + col + tid * VEC + i]); + sc_o[i] = static_cast(static_cast(sc_e[i]) + av.x); + sc_o[i + 1] = static_cast(static_cast(sc_e[i + 1]) + av.y); + sc_o[i + 2] = static_cast(static_cast(sc_e[i + 2]) + av.z); + sc_o[i + 3] = static_cast(static_cast(sc_e[i + 3]) + av.w); + } + reinterpret_cast(&paged_score[dsc])[tid] = sc_out; + } + } + }; + + // 1a. Full chunk for this output block. + // Positions [win_start, sp) are already in paged cache from a prior call; + // start the loop at write_r_start to skip them without a per-iteration branch. + if (should_compress) + { + int const write_r_start = (win_start < sp) ? (sp - win_start) : 0; + write_to_paged(win_start, win_start + COMPRESS_RATIO, write_r_start); + } + + // 1b. Remainder tokens (last block only). + // Tokens past the last complete window are persisted so a later call can + // continue the same compression window. + // rem_start_pos < sp when the chunk has no full window (actual_num_outputs == 0); + // rem_write_start skips those already-paged positions without a per-iteration branch. + if (local_output_idx == num_outputs - 1) + { + int const rem_start_pos = last_abs_idx * COMPRESS_RATIO; + int const rem_count = kv_len - rem_start_pos; + int const rem_write_start = (rem_start_pos < sp) ? (sp - rem_start_pos) : 0; + write_to_paged(rem_start_pos, rem_start_pos + rem_count, rem_write_start); + } + + // ================================================================ + // Phase 2: Online softmax reduction (vectorized) + // + // Each block reduces compress_ratio rows into one output via per-element + // online softmax: output[d] = sum_r(kv[r,d] * softmax(score[r,d] + ape[r,d])) + // In overlap mode, combines previous chunk's first-half and current chunk's + // second-half features (same as decode kernel). + // ================================================================ + if (!should_compress) + return; + + float rmax[VEC], rsum[VEC], rwsum[VEC]; +#pragma unroll + for (int i = 0; i < VEC; i++) + { + rmax[i] = -INFINITY; + rsum[i] = 0.0f; + rwsum[i] = 0.0f; + } + + // Helper: online-softmax reduction over a contiguous window of COMPRESS_RATIO positions. + // Positions [range_start, range_start + new_start) come from paged cache (APE already + // fused during the call that wrote them); positions [range_start + new_start, range_start + // + COMPRESS_RATIO) are new tokens read from kv_score with live APE addition. + // Two branch-free loops avoid a per-iteration if/else on pos < sp. + // kv_col_off — column offset into kv_score / paged state (0 or HEAD_DIM for overlap) + // ape_col_off — column offset into APE table for the score column (same as kv_col_off) + auto reduce_window = [&](int range_start, int kv_col_off, int ape_col_off) + { + // Precompute split once for the whole window. + int const new_start = (range_start < sp) ? min(sp - range_start, COMPRESS_RATIO) : 0; + + // Paged portion — APE already fused when tokens were stored. + for (int r = 0; r < new_start; r++) + { + int const pos = range_start + r; + int log_blk = pos / page_size; + int blk_off = pos % page_size; + int phys_kv = block_table_kv[batch_idx * max_blocks + log_blk]; + int phys_sc = block_table_score[batch_idx * max_blocks + log_blk]; + decodeSoftmaxVec(paged_kv_raw, paged_score_raw, page_sd, state_dim, + phys_kv, phys_sc, blk_off, kv_col_off, tid, rmax, rsum, rwsum); + } + + // New-token portion — read from kv_score and add APE live. + for (int r = new_start; r < COMPRESS_RATIO; r++) + { + int const pos = range_start + r; + int const input_row = pos - sp; + int64_t const row = static_cast(input_offset + input_row) * two_sd; + prefillSoftmaxVec( + kv_score_raw, ape, row, kv_col_off, r * state_dim + ape_col_off, state_dim, tid, rmax, rsum, rwsum); + } + }; + + if constexpr (IS_OVERLAP) + { + // Overlap mode: each output combines + // prev-segment (first head_dim, kv_col=0) from window (abs_idx-1) + // curr-segment (second head_dim, kv_col=HD) from window abs_idx + if (abs_idx > 0) + reduce_window((abs_idx - 1) * COMPRESS_RATIO, 0, 0); // prev window, first half + reduce_window(win_start, HEAD_DIM, HEAD_DIM); // curr window, second half + } + else + { + // Non-overlap mode: reduce one ratio-sized window. + reduce_window(win_start, 0, 0); + } + + // ================================================================ + // Store output (vectorized) + // ================================================================ + int64_t const out_base = static_cast(output_offset + local_output_idx) * HEAD_DIM + tid * VEC; + + if (out_elem_bytes == 2) + { + __nv_bfloat16 packed[VEC]; +#pragma unroll + for (int i = 0; i < VEC; i++) + packed[i] = __float2bfloat16_rn(rwsum[i] / rsum[i]); + // VEC * 2 bytes: VEC=4 → 8B (uint2), VEC=8 → 16B (uint4) + using OutVecT = typename VecType::type; + *reinterpret_cast(&reinterpret_cast<__nv_bfloat16*>(output_raw)[out_base]) + = *reinterpret_cast(packed); + } + else + { + float result[VEC]; +#pragma unroll + for (int i = 0; i < VEC; i++) + result[i] = rwsum[i] / rsum[i]; + // VEC * 4 bytes: VEC=4 → 16B (uint4/float4), VEC=8 → 32B (2×float4) +#pragma unroll + for (int i = 0; i < VEC; i += 4) + *reinterpret_cast(&reinterpret_cast(output_raw)[out_base + i]) + = *reinterpret_cast(&result[i]); + } +} + +// Explicit instantiations +#define INST_PREFILL(HD, KV_EB, STATE_EB, CR) \ + template __global__ void prefillReductionKernel(void const*, float const*, void*, void*, \ + int32_t const*, int32_t const*, void*, int32_t const*, int32_t const*, int32_t const*, int32_t const*, int, \ + int, int, int); + +#define INST_PREFILL_DTYPES(HD, CR) \ + INST_PREFILL(HD, 2, 2, CR) \ + INST_PREFILL(HD, 2, 4, CR) INST_PREFILL(HD, 4, 2, CR) INST_PREFILL(HD, 4, 4, CR) + +INST_PREFILL_DTYPES(128, 4) +INST_PREFILL_DTYPES(128, 128) +INST_PREFILL_DTYPES(512, 4) +INST_PREFILL_DTYPES(512, 128) +#undef INST_PREFILL_DTYPES +#undef INST_PREFILL + +// ============================================================================ +// Prefill Launch Wrapper +// +// Grid is (batch_size, max(max_outputs, 1)). Blocks for local_output_idx >= num_outputs +// early-exit inside the kernel. +// ============================================================================ + +// Compute threads per block: mirrors compile-time NTHRD = HEAD_DIM / VEC. +static inline int prefillNthreads(int head_dim, int elem_bytes_for_vec) +{ + int max_vec = 16 / elem_bytes_for_vec; + int vec = (head_dim / max_vec >= 32) ? max_vec : (head_dim / 32); + return head_dim / vec; +} + +void prefillReductionLaunch(void const* kv_score, float const* ape, void* paged_kv, void* paged_score, + int32_t const* block_table_kv, int32_t const* block_table_score, void* output, int32_t const* kv_lens, + int32_t const* start_pos, int32_t const* cu_seq_lens, int32_t const* cu_kv_comp, int batch_size, int page_size, + int max_blocks, int head_dim, int compress_ratio, int max_outputs, int kv_score_elem_bytes, int state_elem_bytes, + int out_elem_bytes, cudaStream_t stream) +{ + bool const overlap = (compress_ratio == 4); + TLLM_CHECK_WITH_INFO( + compress_ratio == 4 || compress_ratio == 128, "prefillReductionLaunch only supports compress_ratio 4 or 128"); + TLLM_CHECK_WITH_INFO( + (kv_score_elem_bytes == 2 || kv_score_elem_bytes == 4) && (state_elem_bytes == 2 || state_elem_bytes == 4), + "prefillReductionLaunch only supports bf16/fp32 kv_score and paged state"); + int const elem_bytes_for_vec = max(kv_score_elem_bytes, state_elem_bytes); + int const nthreads = prefillNthreads(head_dim, elem_bytes_for_vec); + int const coff = overlap ? 2 : 1; + int const state_dim = coff * head_dim; + dim3 grid(batch_size, max(max_outputs, 1)); + +#define LAUNCH_PREFILL(HD, KV_EB, STATE_EB, CR) \ + prefillReductionKernel<<>>(kv_score, ape, paged_kv, \ + paged_score, block_table_kv, block_table_score, output, kv_lens, start_pos, cu_seq_lens, cu_kv_comp, \ + page_size, state_dim, max_blocks, out_elem_bytes) + +#define DISPATCH_PREFILL_DTYPE(HD, CR) \ + do \ + { \ + if (kv_score_elem_bytes == 4 && state_elem_bytes == 4) \ + { \ + LAUNCH_PREFILL(HD, 4, 4, CR); \ + } \ + else if (kv_score_elem_bytes == 2 && state_elem_bytes == 4) \ + { \ + LAUNCH_PREFILL(HD, 2, 4, CR); \ + } \ + else if (kv_score_elem_bytes == 4 && state_elem_bytes == 2) \ + { \ + LAUNCH_PREFILL(HD, 4, 2, CR); \ + } \ + else \ + { \ + LAUNCH_PREFILL(HD, 2, 2, CR); \ + } \ + } while (false) + + if (head_dim == 512) + { + if (compress_ratio == 4) + DISPATCH_PREFILL_DTYPE(512, 4); + else + DISPATCH_PREFILL_DTYPE(512, 128); + } + else + { + if (compress_ratio == 4) + DISPATCH_PREFILL_DTYPE(128, 4); + else + DISPATCH_PREFILL_DTYPE(128, 128); + } + +#undef DISPATCH_PREFILL_DTYPE +#undef LAUNCH_PREFILL +} + +// ============================================================================ +// Postprocess + Scatter Kernel: postProcessScatterKernel +// +// Template: +// +// Grid: (total_tokens) — one block per compressed token +// Block: (NTHRD = HEAD_DIM / VEC) threads, always >= 32 +// Smem: HEAD_DIM * sizeof(float) — used for cross-warp Hadamard butterfly +// +// This kernel fuses all post-compression processing with the paged cache write +// into a single kernel launch, keeping data in float32 registers throughout. +// This eliminates the DRAM round-trip that a split postprocess+scatter would need. +// +// SCALE_TYPE alone determines the output cache layout: +// - kNone: ELEM_BYTES per value (bf16 / fp32) into kv_cache. +// - kFP8PerTensor: one fp8 byte per value, implicit scale=1.0. +// - kFP8Blockwise: one fp8 byte per value + one fp32 scale per 128 values. +// - kMXFP4Blockwise: packed fp4 (two values per byte) + one ue8m0 byte +// per 32 values. +// +// Pipeline (10 steps, all in fp32 registers): +// 1. Vectorized load compressed token from kv_comp +// 2. RMSNorm: compute sum-of-squares → cross-warp reduce → rsqrt → scale +// 3. Apply RMSNorm weights +// 4. RoPE: interleaved even/odd rotation on rope_dim elements (skip nope_dim) +// 5. Hadamard butterfly transform (3 phases: local → warp shuffle → shared mem) +// 6. Scale by 1/sqrt(HEAD_DIM) (Hadamard normalization) +// 7. Optionally write postprocessed result to kv_out (for callers that need it) +// 8. Binary search cu_kv_comp to find batch_idx for this token +// 9. Compute paged cache destination (logical→physical block via block table) +// 10. Store to cache (layout by SCALE_TYPE). +// ============================================================================ + +template +__global__ void postProcessScatterKernel(void const* __restrict__ kv_comp, // [total_tokens, head_dim] input + void* __restrict__ kv_out, // [total_tokens, head_dim] postprocessed output (may be nullptr) + void const* __restrict__ rms_weight, // [head_dim] + float rms_eps, + float const* __restrict__ cos_sin_table, // [max_pos, 2, rope_dim/2] + int32_t const* __restrict__ position_ids, // [total_tokens] + int nope_dim, int rope_dim, + // scatter params + void* __restrict__ kv_cache, // paged cache buffer + int32_t const* __restrict__ num_outputs_arr, int32_t const* __restrict__ cu_kv_comp, + int32_t const* __restrict__ start_pos_arr, int32_t const* __restrict__ block_offsets, + bool const* __restrict__ compressed_mask, int batch_size, int tokens_per_block, int max_blocks, + int cache_stride_blk_bytes, int total_tokens, int num_scale_blocks, void* __restrict__ quant_output, + void* __restrict__ scale_output) +{ + using ElementT = typename std::conditional::type; + constexpr int MAX_VEC = 16 / ELEM_BYTES; + constexpr int VEC = (HEAD_DIM / MAX_VEC >= 32) ? MAX_VEC : (HEAD_DIM / 32); + constexpr int NTHRD = HEAD_DIM / VEC; + constexpr int VEC_BYTES = VEC * ELEM_BYTES; + using VecT = typename VecType::type; + + int const token_idx = blockIdx.x; + if (token_idx >= total_tokens) + return; + + // Per-token mask: precomputed on host, skips padded generation slots + // and batches that produced no compressed tokens. + if (!compressed_mask[token_idx]) + return; + + // ================================================================ + // Step 0: Find owning batch via binary search on cu_kv_comp. + // ================================================================ + int batch_idx, local_output_idx; + if (batch_size <= 1) + { + batch_idx = 0; + local_output_idx = token_idx; + } + else + { + int lo = 0, hi = batch_size; + while (lo < hi) + { + int mid = (lo + hi) >> 1; + if (cu_kv_comp[mid + 1] <= token_idx) + lo = mid + 1; + else + hi = mid; + } + batch_idx = lo; + if (batch_idx >= batch_size) + return; + local_output_idx = token_idx - cu_kv_comp[batch_idx]; + } + + if (local_output_idx >= num_outputs_arr[batch_idx]) + return; + + int const tid = threadIdx.x; + extern __shared__ float smem[]; + + // ================================================================ + // Step 1: Vectorized load from kv_comp + // ================================================================ + auto const* src = reinterpret_cast( + reinterpret_cast(kv_comp) + static_cast(token_idx) * HEAD_DIM); + VecT raw_in = src[tid]; + ElementT const* in_elems = reinterpret_cast(&raw_in); + + float v[VEC]; +#pragma unroll + for (int i = 0; i < VEC; i++) + v[i] = static_cast(in_elems[i]); + + // ================================================================ + // Step 2: RMSNorm + // ================================================================ + float local_sq = 0.f; +#pragma unroll + for (int i = 0; i < VEC; i++) + local_sq += v[i] * v[i]; + + float warp_sum = warpReduceSum(local_sq); + + constexpr int NUM_WARPS = (NTHRD + 31) / 32; + int const warp_id = tid / 32; + int const lane_id = tid % 32; + + if (lane_id == 0) + smem[warp_id] = warp_sum; + __syncthreads(); + + if (warp_id == 0) + { + float s = (lane_id < NUM_WARPS) ? smem[lane_id] : 0.0f; + for (int offset = 16; offset > 0; offset >>= 1) + s += __shfl_xor_sync(0xFFFFFFFF, s, offset); + if (lane_id == 0) + smem[0] = s; + } + __syncthreads(); + float const rms_scale = rsqrtf(smem[0] / static_cast(HEAD_DIM) + rms_eps); + + // ================================================================ + // Step 3: Load weight, apply RMSNorm + // ================================================================ + auto const* wt_src = reinterpret_cast(reinterpret_cast(rms_weight)); + VecT raw_w = wt_src[tid]; + ElementT const* w_elems = reinterpret_cast(&raw_w); + +#pragma unroll + for (int i = 0; i < VEC; i++) + v[i] = v[i] * rms_scale * static_cast(w_elems[i]); + + // ================================================================ + // Step 4: RoPE (Rotary Positional Embedding) + // + // Applied only to elements in [nope_dim, nope_dim+rope_dim). + // Uses interleaved even/odd pairs: (x_even, x_odd) → rotated by (cos, sin). + // cos_sin_table layout: [max_pos, rope_dim] where first half is cos, second is sin. + // ================================================================ + int const half_rope = rope_dim / 2; + int const pos_id = position_ids[token_idx]; + +#pragma unroll + for (int i = 0; i < VEC; i += 2) + { + int const elem_idx = tid * VEC + i; + if (elem_idx >= nope_dim) + { + int const rope_idx = elem_idx - nope_dim; + int const d = rope_idx >> 1; + float const cos_v = cos_sin_table[pos_id * rope_dim + d]; + float const sin_v = cos_sin_table[pos_id * rope_dim + half_rope + d]; + float const x_even = v[i]; + float const x_odd = v[i + 1]; + v[i] = x_even * cos_v - x_odd * sin_v; + v[i + 1] = x_odd * cos_v + x_even * sin_v; + } + } + + // ================================================================ + // Step 5: Hadamard butterfly transform (rotate activation) + // + // Implements the Walsh-Hadamard transform H_n * v via butterfly network. + // H_n has the recursive structure: H_n = [[H_{n/2}, H_{n/2}], [H_{n/2}, -H_{n/2}]] + // which decomposes into log2(HEAD_DIM) butterfly stages. + // + // Three phases handle increasing stride lengths: + // A) Local: strides < VEC — within each thread's register file + // B) Warp shuffle: strides VEC..32*VEC-1 — via __shfl_xor_sync + // C) Shared memory: strides >= 32*VEC — via XOR-swizzled smem + // + // The XOR swizzle pattern `idx ^ ((idx >> 3) & 0x1F)` ensures bank-conflict- + // free access to shared memory across all butterfly stride patterns. + // + // Skipped entirely when ROTATE_ACTIVATION=false. + // ================================================================ + if constexpr (ROTATE_ACTIVATION) + { + + // Phase A: local butterfly (strides 1..VEC-1, within each thread's VEC registers) +#pragma unroll + for (int stride = 1; stride < VEC; stride <<= 1) + { +#pragma unroll + for (int i = 0; i < VEC; i++) + { + if ((i & stride) == 0) + { + float a = v[i], b = v[i ^ stride]; + v[i] = a + b; + v[i ^ stride] = a - b; + } + } + } + + // Phase B: warp shuffle butterfly (strides VEC..32*VEC-1, within a single warp) + if constexpr (NTHRD > 1) + { + constexpr int SHFL_END = (NTHRD <= 32) ? NTHRD : 32; +#pragma unroll + for (int ts = 1; ts < SHFL_END; ts <<= 1) + { + int const stride = ts * VEC; +#pragma unroll + for (int i = 0; i < VEC; i++) + { + float partner = __shfl_xor_sync(0xFFFFFFFF, v[i], ts); + int const elem_idx = tid * VEC + i; + v[i] = (elem_idx & stride) ? (partner - v[i]) : (v[i] + partner); + } + } + } + + // Phase C: cross-warp butterfly via XOR-swizzled shared memory (strides >= 32*VEC) + // Only needed when NTHRD > 32 (i.e., multiple warps, e.g., HEAD_DIM=512, VEC=8 → 64 threads) + if constexpr (NTHRD > 32) + { +#pragma unroll + for (int i = 0; i < VEC; i++) + { + int const idx = tid * VEC + i; + smem[idx ^ ((idx >> 3) & 0x1F)] = v[i]; + } + __syncthreads(); + + for (int stride = 32 * VEC; stride < HEAD_DIM; stride <<= 1) + { +#pragma unroll + for (int i = 0; i < VEC; i++) + { + int const idx = tid * VEC + i; + int const partner_idx = idx ^ stride; + float const a = smem[idx ^ ((idx >> 3) & 0x1F)]; + float const b = smem[partner_idx ^ ((partner_idx >> 3) & 0x1F)]; + v[i] = (idx & stride) ? (b - a) : (a + b); + } + __syncthreads(); +#pragma unroll + for (int i = 0; i < VEC; i++) + { + int const idx = tid * VEC + i; + smem[idx ^ ((idx >> 3) & 0x1F)] = v[i]; + } + __syncthreads(); + } + } + + // ================================================================ + // Step 6: Scale by Hadamard factor + // ================================================================ + float const had_scale = rsqrtf(static_cast(HEAD_DIM)); + +#pragma unroll + for (int i = 0; i < VEC; i++) + v[i] *= had_scale; + + } // ROTATE_ACTIVATION + + // ================================================================ + // Step 7: Write postprocessed output to kv_out (if requested) + // ================================================================ + if (kv_out != nullptr) + { + VecT raw_out; + ElementT* out_elems = reinterpret_cast(&raw_out); +#pragma unroll + for (int i = 0; i < VEC; i++) + out_elems[i] = static_cast(v[i]); + + auto* dst + = reinterpret_cast(reinterpret_cast(kv_out) + static_cast(token_idx) * HEAD_DIM); + dst[tid] = raw_out; + } + + // ================================================================ + // Step 9: Compute paged cache destination address. + // Map (batch_idx, local_output_idx) → logical block → physical block + // via the block table. block_base points to the start of the physical + // page; token_offset is the slot within that page. + // ================================================================ + int const start_pos = start_pos_arr[batch_idx]; + int const cache_pos = start_pos + local_output_idx; + int const logical_block = cache_pos / tokens_per_block; + int const token_offset = cache_pos % tokens_per_block; + int const phys_block = block_offsets[batch_idx * max_blocks + logical_block]; + + uint8_t* block_base + = reinterpret_cast(kv_cache) + static_cast(phys_block) * cache_stride_blk_bytes; + + // ================================================================ + // Step 11: Store to cache (compile-time dispatch on cache dtype/scale type) + // + // Cache addressing is byte-based: block_base points to the start of + // the physical block, cache_stride_blk_bytes is the total block size. + // ================================================================ + if constexpr (SCALE_TYPE == CacheScaleType::kNone) + { + // Default mode: float→bf16/fp32 pack + vectorized store. + // Cache layout per block: [tokens_per_block * HEAD_DIM] elements of ElementT. + VecT raw_out; + ElementT* out_elems = reinterpret_cast(&raw_out); +#pragma unroll + for (int i = 0; i < VEC; i++) + out_elems[i] = static_cast(v[i]); + + ElementT* row_base = reinterpret_cast(block_base) + token_offset * HEAD_DIM; + reinterpret_cast(row_base)[tid] = raw_out; + } + else if constexpr (SCALE_TYPE == CacheScaleType::kFP8PerTensor) + { + // FP8 per-tensor: direct float→fp8_e4m3fn cast (implicit scale=1.0). + // Cache layout per block: [tokens_per_block * HEAD_DIM] bytes of fp8. + uint8_t fp8_bytes[VEC]; +#pragma unroll + for (int i = 0; i < VEC; i++) + { + __nv_fp8_e4m3 fp8_val(v[i]); + fp8_bytes[i] = *reinterpret_cast(&fp8_val); + } + + using Fp8VecT = typename VecType::type; + uint8_t* fp8_dst = block_base + token_offset * HEAD_DIM + tid * VEC; + *reinterpret_cast(fp8_dst) = *reinterpret_cast(fp8_bytes); + } + else if constexpr (SCALE_TYPE == CacheScaleType::kFP8Blockwise) + { + // FP8 blockwise: per-128-element quantization with explicit scales. + // GROUP_SIZE = number of threads that share one scale factor. + // For HD=512, VEC=8: GROUP_SIZE=16 threads → 128 elements per scale block. + // + // Cache layout per block: + // [fp8_data: tokens_per_block * HEAD_DIM bytes] + // [scales: tokens_per_block * (HEAD_DIM/128) * 4 bytes] + constexpr int GROUP_SIZE = 128 / VEC; + + // Step 11a: Compute per-group amax via warp shuffle reduction. + // GROUP_SIZE <= 16 (< warp), so shuffle is sufficient (no smem needed). + float local_amax = 0.f; +#pragma unroll + for (int i = 0; i < VEC; i++) + local_amax = fmaxf(local_amax, fabsf(v[i])); + +#pragma unroll + for (int offset = GROUP_SIZE / 2; offset > 0; offset >>= 1) + local_amax = fmaxf(local_amax, __shfl_xor_sync(0xFFFFFFFF, local_amax, offset)); + + // Step 11b: Compute scale and inverse scale for quantization. + // 448.0 is the max representable value for fp8_e4m3fn. + float const scale = local_amax / 448.0f; + float const inv_scale = (local_amax > 0.f) ? (448.0f / local_amax) : 1.0f; + + // Step 11c: Quantize to FP8 and store data. + uint8_t fp8_bytes[VEC]; +#pragma unroll + for (int i = 0; i < VEC; i++) + { + __nv_fp8_e4m3 fp8_val(v[i] * inv_scale); + fp8_bytes[i] = *reinterpret_cast(&fp8_val); + } + + using Fp8VecT = typename VecType::type; + uint8_t* fp8_dst = block_base + token_offset * HEAD_DIM + tid * VEC; + *reinterpret_cast(fp8_dst) = *reinterpret_cast(fp8_bytes); + + // Step 11d: Store scale factor (one thread per 128-element group writes it). + if (tid % GROUP_SIZE == 0) + { + int const scale_idx = tid / GROUP_SIZE; + float* scale_dst = reinterpret_cast(block_base + tokens_per_block * HEAD_DIM + + (token_offset * num_scale_blocks + scale_idx) * sizeof(float)); + *scale_dst = scale; + } + + // Step 11e: Optionally write FP8 data and scales to output buffers. + // Used by the indexer compressor which returns (fp8_data, scales) to Python + // for downstream sparse attention indexing. + if (quant_output != nullptr) + { + uint8_t* fp8_out_dst + = reinterpret_cast(quant_output) + static_cast(token_idx) * HEAD_DIM + tid * VEC; + *reinterpret_cast(fp8_out_dst) = *reinterpret_cast(fp8_bytes); + } + if (scale_output != nullptr && tid % GROUP_SIZE == 0) + { + int const scale_idx = tid / GROUP_SIZE; + reinterpret_cast(scale_output)[static_cast(token_idx) * num_scale_blocks + scale_idx] + = scale; + } + } + else if constexpr (SCALE_TYPE == CacheScaleType::kMXFP4Blockwise) + { + constexpr int GROUP_SIZE = 32 / VEC; + constexpr int PACKED_VEC_BYTES = VEC / 2; + constexpr float kFp4Max = 6.0f; + constexpr float kFp4MaxInv = 1.0f / kFp4Max; + constexpr float kFp4MinAmax = kFp4Max * 1.1754943508222875e-38f; + + float local_amax = 0.f; +#pragma unroll + for (int i = 0; i < VEC; i++) + local_amax = fmaxf(local_amax, fabsf(v[i])); + +#pragma unroll + for (int offset = GROUP_SIZE / 2; offset > 0; offset >>= 1) + local_amax = fmaxf(local_amax, __shfl_xor_sync(0xFFFFFFFF, local_amax, offset)); + + float const scale = roundedPow2Scale(local_amax, kFp4MaxInv, kFp4MinAmax); + + uint8_t fp4_bytes[PACKED_VEC_BYTES]; +#pragma unroll + for (int i = 0; i < VEC; i += 2) + { + fp4_bytes[i / 2] = packE2M1x2(v[i] / scale, v[i + 1] / scale); + } + + int const packed_head_dim = HEAD_DIM / 2; + uint8_t* fp4_dst = block_base + token_offset * packed_head_dim + tid * PACKED_VEC_BYTES; +#pragma unroll + for (int i = 0; i < PACKED_VEC_BYTES; ++i) + fp4_dst[i] = fp4_bytes[i]; + + if (tid % GROUP_SIZE == 0) + { + int const scale_idx = tid / GROUP_SIZE; + uint8_t* scale_dst + = block_base + tokens_per_block * packed_head_dim + token_offset * num_scale_blocks + scale_idx; + *scale_dst = toUe8m0(scale); + } + + if (quant_output != nullptr) + { + uint8_t* fp4_out_dst = reinterpret_cast(quant_output) + + static_cast(token_idx) * packed_head_dim + tid * PACKED_VEC_BYTES; +#pragma unroll + for (int i = 0; i < PACKED_VEC_BYTES; ++i) + fp4_out_dst[i] = fp4_bytes[i]; + } + if (scale_output != nullptr && tid % GROUP_SIZE == 0) + { + int const scale_idx = tid / GROUP_SIZE; + reinterpret_cast(scale_output)[static_cast(token_idx) * num_scale_blocks + scale_idx] + = toUe8m0(scale); + } + } +} + +// Explicit instantiations — fused postprocess+scatter. +// kNone supports bf16 (EB=2) and fp32 (EB=4) input types; the quantized +// scale types only support bf16 input since the compressor output is bf16. +// Each combination is instantiated with ROTATE_ACTIVATION=true and false. +#define INST_PPS(HD, EB, CST, AR) \ + template __global__ void postProcessScatterKernel(void const*, void*, void const*, float, \ + float const*, int32_t const*, int, int, void*, int32_t const*, int32_t const*, int32_t const*, int32_t const*, \ + bool const*, int, int, int, int, int, int, void*, void*); + +#define INST_PPS_AR(HD, EB, CST) \ + INST_PPS(HD, EB, CST, true) \ + INST_PPS(HD, EB, CST, false) + +INST_PPS_AR(128, 2, CacheScaleType::kNone) +INST_PPS_AR(128, 4, CacheScaleType::kNone) +INST_PPS_AR(512, 2, CacheScaleType::kNone) +INST_PPS_AR(512, 4, CacheScaleType::kNone) +INST_PPS_AR(128, 2, CacheScaleType::kFP8PerTensor) +INST_PPS_AR(512, 2, CacheScaleType::kFP8PerTensor) +INST_PPS_AR(128, 2, CacheScaleType::kFP8Blockwise) +INST_PPS_AR(512, 2, CacheScaleType::kFP8Blockwise) +INST_PPS_AR(128, 2, CacheScaleType::kMXFP4Blockwise) +INST_PPS_AR(512, 2, CacheScaleType::kMXFP4Blockwise) +#undef INST_PPS_AR +#undef INST_PPS + +// ============================================================================ +// Postprocess + Scatter Launch Wrapper +// +// Derives cache layout parameters (cache_stride_blk_bytes, num_scale_blocks) +// from cache dtype / scale type, then dispatches to the appropriate template instantiation. +// ============================================================================ + +// Compute number of threads per block, mirroring the compile-time VEC/NTHRD logic. +// Ensures NTHRD >= 32 by reducing VEC when HEAD_DIM is small. +static inline int compressorNthreads(int head_dim, int elem_bytes) +{ + int max_vec = 16 / elem_bytes; + int vec = (head_dim / max_vec >= 32) ? max_vec : (head_dim / 32); + return head_dim / vec; +} + +void postProcessScatterLaunch(void const* kv_comp, void* kv_out, void const* rms_weight, float rms_eps, + float const* cos_sin_table, int32_t const* position_ids, int nope_dim, int rope_dim, void* kv_cache, + int32_t const* num_outputs, int32_t const* cu_kv_comp, int32_t const* start_pos, int32_t const* block_offsets, + bool const* compressed_mask, int batch_size, int tokens_per_block, int head_dim, int max_blocks_per_seq, + int elem_bytes, int total_tokens, int cache_scale_type, bool rotate_activation, void* quant_output, + void* scale_output, cudaStream_t stream) +{ + if (total_tokens == 0) + { + return; + } + + TLLM_CHECK_WITH_INFO( + cache_scale_type >= 0 && cache_scale_type <= 3, "Invalid cache_scale_type: %d", cache_scale_type); + auto const cst = static_cast(cache_scale_type); + + bool const is_quantized_store = (cst != CacheScaleType::kNone); + int const nthreads = compressorNthreads(head_dim, elem_bytes); + int const smem_bytes = head_dim * sizeof(float); + TLLM_CHECK_WITH_INFO(cst != CacheScaleType::kMXFP4Blockwise || head_dim % 32 == 0, + "MXFP4 cache requires head_dim divisible by 32, got %d", head_dim); + TLLM_CHECK_WITH_INFO(cst != CacheScaleType::kMXFP4Blockwise || head_dim % 2 == 0, + "FP4 packed cache requires even head_dim, got %d", head_dim); + TLLM_CHECK_WITH_INFO(!is_quantized_store || elem_bytes == 2, + "Quantized cache modes require bf16 compressor output, got elem_bytes=%d", elem_bytes); + + // Derive cache block stride (in bytes) and scale block count from the + // scale type. Each physical cache block stores tokens_per_block tokens: + // none: tpb * HD * elem_bytes + // fp8 pertensor: tpb * HD + // fp8 blockwise: tpb * HD + tpb * (HD/128)*4 + // mxfp4: tpb * (HD/2) + tpb * (HD/32) + int num_scale_blocks = 0; + int cache_stride_blk_bytes = 0; + switch (cst) + { + case CacheScaleType::kFP8PerTensor: cache_stride_blk_bytes = tokens_per_block * head_dim; break; + case CacheScaleType::kFP8Blockwise: + num_scale_blocks = head_dim / 128; + cache_stride_blk_bytes = tokens_per_block * head_dim + tokens_per_block * num_scale_blocks * 4; + break; + case CacheScaleType::kMXFP4Blockwise: + num_scale_blocks = head_dim / 32; + cache_stride_blk_bytes = tokens_per_block * (head_dim / 2) + tokens_per_block * num_scale_blocks; + break; + default: cache_stride_blk_bytes = tokens_per_block * head_dim * elem_bytes; break; + } + +#define LAUNCH_PPS(HD, EB, CST, AR) \ + postProcessScatterKernel<<>>(kv_comp, kv_out, \ + rms_weight, rms_eps, cos_sin_table, position_ids, nope_dim, rope_dim, kv_cache, num_outputs, cu_kv_comp, \ + start_pos, block_offsets, compressed_mask, batch_size, tokens_per_block, max_blocks_per_seq, \ + cache_stride_blk_bytes, total_tokens, num_scale_blocks, quant_output, scale_output) + +#define DISPATCH_ROTATE(HD, EB, CST) \ + if (rotate_activation) \ + { \ + LAUNCH_PPS(HD, EB, CST, true); \ + } \ + else \ + { \ + LAUNCH_PPS(HD, EB, CST, false); \ + } +#define DISPATCH_HD_EB(CST) \ + if (elem_bytes == 4) \ + { \ + switch (head_dim) \ + { \ + case 128: DISPATCH_ROTATE(128, 4, CST); break; \ + default: DISPATCH_ROTATE(512, 4, CST); break; \ + } \ + } \ + else \ + { \ + switch (head_dim) \ + { \ + case 128: DISPATCH_ROTATE(128, 2, CST); break; \ + default: DISPATCH_ROTATE(512, 2, CST); break; \ + } \ + } +#define DISPATCH_HD_BF16(CST) \ + switch (head_dim) \ + { \ + case 128: DISPATCH_ROTATE(128, 2, CST); break; \ + default: DISPATCH_ROTATE(512, 2, CST); break; \ + } + + if (cst == CacheScaleType::kFP8PerTensor) + { + DISPATCH_HD_BF16(CacheScaleType::kFP8PerTensor); + } + else if (cst == CacheScaleType::kFP8Blockwise) + { + DISPATCH_HD_BF16(CacheScaleType::kFP8Blockwise); + } + else if (cst == CacheScaleType::kMXFP4Blockwise) + { + DISPATCH_HD_BF16(CacheScaleType::kMXFP4Blockwise); + } + else + { + DISPATCH_HD_EB(CacheScaleType::kNone); + } + +#undef DISPATCH_HD_BF16 +#undef DISPATCH_HD_EB +#undef DISPATCH_ROTATE +#undef LAUNCH_PPS +} + +} // namespace kernels::compressor + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.h b/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.h new file mode 100644 index 000000000000..bff4a1691bc9 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.h @@ -0,0 +1,99 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include +#include +#include +#include + +#include "tensorrt_llm/common/config.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::compressor +{ + +// Decode kernel: write NEXT_N tokens to paged cache + conditional compression +// via online softmax. Overlap is derived from compress_ratio (ratio=4). +// +// Grid: (batch_size, cdiv(state_dim, block_size)) +// One thread per state_dim element across all phases. +// state_dim is a constexpr derived from COMPRESS_RATIO and HEAD_DIM inside the kernel. +void pagedKvCompressLaunch(void const* kv_score, // [m, 2*state_dim] (bf16 or fp32) + float const* ape, // [compress_ratio, state_dim] + void* paged_kv, // [num_blocks, page_size, state_dim] + void* paged_score, // [num_blocks, page_size, state_dim] + int32_t const* block_table_kv, // [bsz, max_blocks] + int32_t const* block_table_score, // [bsz, max_blocks] + void* output, // [total_outputs, head_dim] + int32_t const* kv_lens, // [bsz] + int32_t const* cu_seq_lens, // [bsz+1] + int32_t const* cu_kv_comp, // [bsz+1] + int batch_size, int page_size, int max_blocks, int head_dim, int compress_ratio, int next_n, + int kv_score_elem_bytes, // bytes per element for kv_score (2=bf16, 4=fp32) + int state_elem_bytes, // bytes per element for paged state (2=bf16, 4=fp32) + int out_elem_bytes, // bytes per element for output + cudaStream_t stream); + +// Prefill kernel: bulk compression with per-token gather/scatter + state update. +// Writes remainder tokens to paged cache, then performs online softmax reduction. +// +// Grid: (batch_size, max_outputs_per_batch, num_head_chunks) +// Each block computes one compressed output for one head_dim chunk. +void prefillReductionLaunch(void const* kv_score, // [m, 2*state_dim] (bf16 or fp32) + float const* ape, // [compress_ratio, state_dim] + void* paged_kv, // [num_blocks, page_size, state_dim] + void* paged_score, // [num_blocks, page_size, state_dim] + int32_t const* block_table_kv, // [bsz, max_blocks] + int32_t const* block_table_score, // [bsz, max_blocks] + void* output, // [total_outputs, head_dim] + int32_t const* kv_lens, // [bsz] + int32_t const* start_pos, // [bsz] + int32_t const* cu_seq_lens, // [bsz+1] + int32_t const* cu_kv_comp, // [bsz+1] + int batch_size, int page_size, int max_blocks, int head_dim, int compress_ratio, int max_outputs, + int kv_score_elem_bytes, int state_elem_bytes, int out_elem_bytes, cudaStream_t stream); + +// RMSNorm + RoPE + Hadamard + paged scatter in a single kernel launch. +// Optionally writes postprocessed result to kv_out (nullptr to skip). +// +// Grid: (total_tokens) -- one block per compressed token +// Block: (head_dim / VEC), always >= 32 +void postProcessScatterLaunch(void const* kv_comp, // [total_tokens, head_dim] input + void* kv_out, // [total_tokens, head_dim] postprocessed output (nullptr to skip) + void const* rms_weight, // [head_dim] + float rms_eps, + float const* cos_sin_table, // [max_pos, 2, rope_dim/2] + int32_t const* position_ids, // [total_tokens] + int nope_dim, int rope_dim, + void* kv_cache, // paged cache buffer + int32_t const* num_outputs, // [bsz] + int32_t const* cu_kv_comp, // [bsz+1] + int32_t const* start_pos, // [bsz] + int32_t const* block_offsets, // [bsz, max_blocks] + bool const* compressed_mask, // [total_tokens] — per-token mask, false ⇒ skip + int batch_size, int tokens_per_block, int head_dim, int max_blocks_per_seq, int elem_bytes, int total_tokens, + int cache_scale_type, // 0=none (bf16/fp32 by elem_bytes), 1=fp8_pertensor, + // 2=fp8_blockwise, 3=mxfp4 (packed FP4) + bool rotate_activation, // whether to apply Hadamard transform (false to skip) + void* quant_output, // optional fp8/fp4 packed output (nullptr if unused) + void* scale_output, // optional scale output (float* for fp8, uint8_t* for fp4) + cudaStream_t stream); + +} // namespace kernels::compressor + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/CMakeLists.txt b/cpp/tensorrt_llm/kernels/mhcKernels/CMakeLists.txt new file mode 100644 index 000000000000..eb07c6f165d3 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/mhcKernels/CMakeLists.txt @@ -0,0 +1,38 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. SPDX-License-Identifier: Apache-2.0 +# +# 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. +# + +# Library sources are listed explicitly. Profiling/benchmark/test harnesses in +# this directory (bench_*.cu, probe_*.cu, profile_*.cu, test_*.cu) are +# standalone binaries with their own main() and must NOT be compiled into the +# trtllm shared library. +set(SRC_CU mhcKernels.cu mhcFusedHcKernel.cu) +add_library(mhcKernels_src OBJECT ${SRC_CU}) + +set_property(TARGET mhcKernels_src PROPERTY POSITION_INDEPENDENT_CODE ON) +set_property(TARGET mhcKernels_src PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS ON) +target_compile_options(mhcKernels_src + PRIVATE $<$:--use_fast_math>) + +if(BUILD_DEEP_GEMM) + # Expose DeepGEMM helper headers (sm100_utils.cuh, utils.cuh, tma_utils.cuh, + # reduction.cuh) used by the tcgen05.mma-based fused post-mapping + GEMM + # kernel on SM100. + target_include_directories( + mhcKernels_src + PRIVATE ${CMAKE_BINARY_DIR}/_deps/deepgemm-src/deep_gemm/include) + target_compile_definitions(mhcKernels_src PRIVATE TRTLLM_MHC_ENABLE_FUSED_HC) +endif() diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh b/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh new file mode 100644 index 000000000000..ed2841f63ce2 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh @@ -0,0 +1,1602 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * 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. + */ + +// Fused post-mapping + TF32 HC prenorm GEMM on B200 (SM100). +// +// Mathematical formula: +// new_r[i, hc, h] = post_mix[i, hc] * x[i, h] +// + sum_j comb_mix[i, j, hc] * residual[i, j, h] +// D[i, n] = sum_{hc, h} new_r[i, hc, h] * W[n, hc, h] +// sqr[i] = sum_{hc, h} new_r[i, hc, h]^2 (bf16-rounded) +// +// Shape: +// residual [M, HC_MULT, hidden] bf16 +// x [M, hidden] bf16 +// post_mix [M, HC_MULT] fp32 +// comb_mix [M, HC_MULT, HC_MULT] fp32 +// W [N, HC_MULT*hidden] fp32 (TF32) +// D [M, N] fp32 +// sqr [M] fp32 +// +// Kernel architecture (derived from DeepGEMM sm100_tf32_hc_prenorm_gemm): +// - BLOCK_M x BLOCK_N x BLOCK_K = 64 x 32 x 64, TF32 MMA on tcgen05 +// - 256 threads per CTA (warps 0..3 = MMA group, warps 4..7 = pmap group) +// - new_r is NEVER materialized in GMEM - it is computed on-chip directly +// into TMEM (fp32) where UMMA consumes it. +// - Iteration order: outer h_tile (slow), inner hc_idx (fast). residual+x +// SMEM is reused for HC_MULT=4 consecutive MMA stages. +// +// Barriers: +// full_B[N_B_STAGES] TMA -> MMA (B arrived in SMEM) +// empty_B[N_B_STAGES] MMA -> TMA (B slot empty) +// full_input[N_INPUT] TMA -> pmap (residual+x arrived) +// empty_input[N_INPUT] pmap -> TMA (input slot empty) +// full_cast[2] pmap -> MMA (A ready in TMEM) +// empty_cast[2] MMA -> pmap (TMEM slot empty) +// tmem_full[1] MMA -> epilogue + +#pragma once +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wunknown-attributes" + +#include +#include +#include + +namespace deep_gemm::sm90 +{ +using cuda::std::swap; +} + +namespace deep_gemm::sm100 +{ +using cuda::std::swap; +} + +#include +#include +#include +#include +#include +#include +#include + +namespace fused_mhc +{ + +// Reuse DeepGEMM's swizzle helper +using deep_gemm::sm100::make_umma_desc; +using deep_gemm::sm100::get_num_aligned_tmem_cols; +using deep_gemm::sm100::tcgen05_before_thread_sync; +using deep_gemm::sm100::tcgen05_after_thread_sync; +using deep_gemm::sm100::advance_umma_desc_lo; +using deep_gemm::tma_copy; +using deep_gemm::utils::PatternVisitor; +using deep_gemm::ptx::get_lane_idx; + +template +__device__ __forceinline__ uint32_t get_swizzled_smem_offset(uint32_t const& offset, uint32_t const& lane_idx) +{ + auto const& bank_group_idx = offset + lane_idx * (kSwizzleMode / kSwizzleBase); + constexpr uint32_t kNumBankGroups = 128 / kSwizzleBase; + constexpr bool kHasShortcut = (kSwizzleMode / kSwizzleBase) == kNumBankGroups; + auto row = kHasShortcut ? (offset / kNumBankGroups + lane_idx) : (bank_group_idx / kNumBankGroups); + auto col = kHasShortcut ? (offset) : (bank_group_idx % kNumBankGroups); + col ^= row % (kSwizzleMode / kSwizzleBase); + return row * 128 + col * kSwizzleBase; +} + +__device__ __forceinline__ void stsm_x4_b16_rout(void* smem_dst, uint32_t a, uint32_t b, uint32_t c, uint32_t d) +{ + asm volatile( + "stmatrix.sync.aligned.x4.m8n8.shared.b16 [%0], {%1, %2, %3, %4};\n" ::"l"(__cvta_generic_to_shared(smem_dst)), + "r"(a), "r"(b), "r"(c), "r"(d)); +} + +template +__global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) fused_tf32_pmap_gemm_rout_atomic_impl( + const uint32_t shape_m, const __grid_constant__ cute::TmaDescriptor tensor_map_residual, + const __grid_constant__ cute::TmaDescriptor tensor_map_x, const __grid_constant__ cute::TmaDescriptor tensor_map_b, + const __grid_constant__ cute::TmaDescriptor tensor_map_residual_out, + float* __restrict__ D, // [M, SHAPE_N] (caller memsets to 0) + float const* __restrict__ post_mix, float const* __restrict__ comb_mix, float* __restrict__ sqr_sum) +{ // [M] (caller memsets to 0) +#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000) and (__CUDA_ARCH__ < 1100)) or defined(__CLION_IDE__) + using Barrier = cutlass::arch::ClusterTransactionBarrier; + + constexpr uint32_t SHAPE_K = HC_MULT * HIDDEN; + constexpr uint32_t H_TILES_PER_HC = HIDDEN / BLOCK_K; + static_assert(H_TILES_PER_HC % kNumSplits == 0, "H_TILES_PER_HC must be divisible by kNumSplits"); + constexpr uint32_t H_TILES_PER_SPLIT = H_TILES_PER_HC / kNumSplits; + constexpr uint32_t kNumCastStages = 4; + constexpr uint32_t kSwizzleAMode = cute::min(BLOCK_K * sizeof(nv_bfloat16), 128); + constexpr uint32_t kSwizzleBMode = cute::min(BLOCK_K * sizeof(float), 128); + constexpr uint32_t kSwizzleXMode = kSwizzleAMode; + constexpr uint32_t kSwizzleResMode = kSwizzleAMode; + constexpr uint32_t kSwizzleRoutMode = kSwizzleAMode; + constexpr auto kMajorA = cute::UMMA::Major::K; + constexpr auto kMajorB = cute::UMMA::Major::K; + static_assert(HIDDEN % BLOCK_K == 0, "HIDDEN must be multiple of BLOCK_K"); + static_assert(N_B_STAGES >= HC_MULT, "N_B_STAGES must be >= HC_MULT"); + static_assert(kSwizzleCDMode / sizeof(float) == BLOCK_N, "Invalid block N"); + static_assert(kNumMMAThreads == 128, "Invalid MMA threads"); + static_assert(kNumPmapThreads == 128, "Invalid pmap threads"); + static_assert(BLOCK_M == 64, "Invalid block M"); + static_assert(HC_MULT == 4, "Only HC_MULT=4 supported"); + static_assert(kSwizzleCDMode == 128, "Atomic variant expects kSwizzleCDMode=128"); + static_assert(SHAPE_N <= BLOCK_N, "SHAPE_N must fit within BLOCK_N"); + + auto const warp_idx = cutlass::canonical_warp_idx_sync(); + auto const lane_idx = get_lane_idx(); + + extern __shared__ __align__(1024) uint8_t smem_buffer[]; + + constexpr uint32_t SMEM_CD_SIZE = BLOCK_M * kSwizzleCDMode; + constexpr uint32_t SMEM_B_PER_STAGE = BLOCK_N * BLOCK_K * sizeof(float); + constexpr uint32_t SMEM_RES_PER_ISTG = BLOCK_M * HC_MULT * BLOCK_K * sizeof(nv_bfloat16); + constexpr uint32_t SMEM_X_PER_ISTG = BLOCK_M * BLOCK_K * sizeof(nv_bfloat16); + constexpr uint32_t SMEM_POST_SIZE = BLOCK_M * HC_MULT * sizeof(float); + constexpr uint32_t SMEM_COMB_SIZE = BLOCK_M * HC_MULT * HC_MULT * sizeof(float); + constexpr uint32_t SMEM_RC_PER_HC = BLOCK_M * BLOCK_K * sizeof(nv_bfloat16); // 8 KB + constexpr uint32_t SMEM_RC_SIZE = HC_MULT * SMEM_RC_PER_HC; // 32 KB + + constexpr uint32_t kNumTmemCols = get_num_aligned_tmem_cols(); + + // Prefetch TMA descriptors + if (warp_idx == 0 and cute::elect_one_sync()) + { + cute::prefetch_tma_descriptor(&tensor_map_residual); + cute::prefetch_tma_descriptor(&tensor_map_x); + cute::prefetch_tma_descriptor(&tensor_map_b); + cute::prefetch_tma_descriptor(&tensor_map_residual_out); + } + + // SMEM layout: [cd, B stages, res stages, x stages, post, comb, rc (HC_MULT slices)] + auto smem_cd = reinterpret_cast(smem_buffer); + uint8_t* cursor = smem_buffer + SMEM_CD_SIZE; + auto smem_b = PatternVisitor( + [&, base = cursor](uint32_t const& i) { return reinterpret_cast(base + i * SMEM_B_PER_STAGE); }); + cursor += N_B_STAGES * SMEM_B_PER_STAGE; + auto smem_res = PatternVisitor( + [&, base = cursor](uint32_t const& i) { return reinterpret_cast(base + i * SMEM_RES_PER_ISTG); }); + cursor += N_INPUT_STAGES * SMEM_RES_PER_ISTG; + auto smem_x_stg = PatternVisitor( + [&, base = cursor](uint32_t const& i) { return reinterpret_cast(base + i * SMEM_X_PER_ISTG); }); + cursor += N_INPUT_STAGES * SMEM_X_PER_ISTG; + auto smem_post = reinterpret_cast(cursor); + cursor += SMEM_POST_SIZE; + auto smem_comb = reinterpret_cast(cursor); + cursor += SMEM_COMB_SIZE; + auto smem_rc = reinterpret_cast(cursor); // [HC_MULT][BLOCK_M][BLOCK_K] bf16, single-buffered + cursor += SMEM_RC_SIZE; + + cursor = reinterpret_cast((reinterpret_cast(cursor) + 7) & ~uintptr_t(7)); + auto barrier_start_ptr = reinterpret_cast(cursor); + auto full_B = PatternVisitor([=](uint32_t const& i) { return barrier_start_ptr + i; }); + auto empty_B = PatternVisitor([=](uint32_t const& i) { return barrier_start_ptr + N_B_STAGES + i; }); + auto full_input = PatternVisitor([=](uint32_t const& i) { return barrier_start_ptr + 2 * N_B_STAGES + i; }); + auto empty_input + = PatternVisitor([=](uint32_t const& i) { return barrier_start_ptr + 2 * N_B_STAGES + N_INPUT_STAGES + i; }); + auto full_cast = PatternVisitor( + [=](uint32_t const& i) { return barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + i; }); + auto empty_cast = PatternVisitor([=](uint32_t const& i) + { return barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + kNumCastStages + i; }); + auto tmem_full_barrier = barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + 2 * kNumCastStages; + + cursor += (2 * N_B_STAGES + 2 * N_INPUT_STAGES + 2 * kNumCastStages + 1) * sizeof(Barrier); + auto tmem_ptr_in_smem = reinterpret_cast(cursor); + + if (warp_idx == 1 and cute::elect_one_sync()) + { +#pragma unroll + for (uint32_t i = 0; i < N_B_STAGES; ++i) + { + full_B[i]->init(1); + empty_B[i]->init(1); + } +#pragma unroll + for (uint32_t i = 0; i < N_INPUT_STAGES; ++i) + { + full_input[i]->init(1); + empty_input[i]->init(kNumPmapThreads); + } +#pragma unroll + for (uint32_t i = 0; i < kNumCastStages; ++i) + { + full_cast[i]->init(kNumPmapThreads); + empty_cast[i]->init(1); + } + tmem_full_barrier->init(1); + cutlass::arch::fence_barrier_init(); + } + else if (warp_idx == 2) + { + cute::TMEM::Allocator1Sm().allocate(kNumTmemCols, tmem_ptr_in_smem); + } + __syncthreads(); + + const uint32_t block_idx = __shfl_sync(0xffffffff, blockIdx.x, 0); + const uint32_t m_block_idx = block_idx / kNumSplits; + const uint32_t k_split_idx = block_idx % kNumSplits; + const uint32_t m_offset = m_block_idx * BLOCK_M; + const uint32_t h_tile_start = k_split_idx * H_TILES_PER_SPLIT; + constexpr uint32_t num_total_stages = H_TILES_PER_SPLIT * HC_MULT; + + // Prologue: pmap warp group loads post_mix, comb_mix into SMEM + if (warp_idx >= kNumMMAThreads / 32) + { + const uint32_t pmap_tid = threadIdx.x - kNumMMAThreads; +#pragma unroll + for (uint32_t t = 0; t < 2; ++t) + { + uint32_t idx = pmap_tid + t * kNumPmapThreads; + if (idx < BLOCK_M * HC_MULT) + { + uint32_t m = idx / HC_MULT; + uint32_t hc = idx % HC_MULT; + uint32_t gmem_m = m_offset + m; + float v = (gmem_m < shape_m) ? post_mix[gmem_m * HC_MULT + hc] : 0.f; + smem_post[idx] = v; + } + } +#pragma unroll + for (uint32_t t = 0; t < 8; ++t) + { + uint32_t idx = pmap_tid + t * kNumPmapThreads; + if (idx < BLOCK_M * HC_MULT * HC_MULT) + { + uint32_t m = idx / (HC_MULT * HC_MULT); + uint32_t jk = idx % (HC_MULT * HC_MULT); + uint32_t gmem_m = m_offset + m; + float v = (gmem_m < shape_m) ? comb_mix[gmem_m * HC_MULT * HC_MULT + jk] : 0.f; + smem_comb[idx] = v; + } + } + } + __syncthreads(); + + if (warp_idx < kNumMMAThreads / 32) + { + // ----- TMA warp (warp 0) ----- + if (warp_idx == 0 and cute::elect_one_sync()) + { + uint32_t b_stage = 0; + uint32_t i_stage = 0; + uint32_t s = 0; + for (uint32_t ht = 0; ht < H_TILES_PER_SPLIT; ++ht) + { + const uint32_t h_tile = h_tile_start + ht; + empty_input[i_stage]->wait(((ht / N_INPUT_STAGES) & 1) ^ 1); + uint32_t m_idx = m_block_idx * BLOCK_M; + uint32_t h_idx = h_tile * BLOCK_K; +#pragma unroll + for (uint32_t j = 0; j < HC_MULT; ++j) + { + tma_copy(&tensor_map_residual, full_input[i_stage], + smem_res[i_stage] + j * BLOCK_M * BLOCK_K, j * HIDDEN + h_idx, m_idx); + } + tma_copy( + &tensor_map_x, full_input[i_stage], smem_x_stg[i_stage], h_idx, m_idx); + constexpr uint32_t kInputBytes = SMEM_RES_PER_ISTG + SMEM_X_PER_ISTG; + full_input[i_stage]->arrive_and_expect_tx(kInputBytes); + +#pragma unroll + for (uint32_t hc = 0; hc < HC_MULT; ++hc) + { + empty_B[b_stage]->wait(((s / N_B_STAGES) & 1) ^ 1); + uint32_t k_idx = hc * HIDDEN + h_idx; + tma_copy( + &tensor_map_b, full_B[b_stage], smem_b[b_stage], k_idx, 0); + full_B[b_stage]->arrive_and_expect_tx(SMEM_B_PER_STAGE); + b_stage = (b_stage + 1) % N_B_STAGES; + ++s; + } + i_stage = (i_stage + 1) % N_INPUT_STAGES; + } + } + + // ----- MMA issue warp (warp 1) ----- + if (warp_idx == 1) + { + constexpr uint32_t UMMA_M = BLOCK_M; + constexpr uint32_t UMMA_N = BLOCK_N; + constexpr uint32_t UMMA_K = 32 / sizeof(float); + constexpr uint32_t BLOCK_SWIZZLED_BK = kSwizzleBMode / sizeof(float); + using umma_t = cute::SM100_MMA_TF32_TS; + auto instr_desc = cute::UMMA::make_instr_desc(); + auto const& runtime_instr_desc = cute::UMMA::make_runtime_instr_desc(instr_desc); + static_assert(N_B_STAGES <= 32, "Too many B stages"); + auto b_desc = make_umma_desc(smem_b[0], 0, 0); + uint32_t const& b_desc_lo = lane_idx < N_B_STAGES ? b_desc.lo + lane_idx * SMEM_B_PER_STAGE / 16 : 0u; + + for (uint32_t s = 0; s < num_total_stages; ++s) + { + const uint32_t b_stage = s % N_B_STAGES; + const uint32_t cast_stage_idx = s % kNumCastStages; + full_cast[cast_stage_idx]->wait((s / kNumCastStages) & 1); + full_B[b_stage]->wait((s / N_B_STAGES) & 1); + tcgen05_after_thread_sync(); + auto const& b_desc_base_lo = __shfl_sync(0xffffffff, b_desc_lo, static_cast(b_stage)); +#pragma unroll + for (uint32_t k = 0; k < BLOCK_K / UMMA_K; ++k) + { + uint32_t const& atom_idx = (k * UMMA_K) / BLOCK_SWIZZLED_BK; + uint32_t const& in_atom_idx = (k * UMMA_K) % BLOCK_SWIZZLED_BK; + uint32_t const& offset = atom_idx * BLOCK_N * BLOCK_SWIZZLED_BK; + b_desc.lo = advance_umma_desc_lo( + b_desc_base_lo, offset, in_atom_idx); + umma_t::fma(BLOCK_K * cast_stage_idx + k * UMMA_K, b_desc, BLOCK_K * kNumCastStages, s > 0 or k > 0, + runtime_instr_desc); + } + cutlass::arch::umma_arrive(reinterpret_cast(empty_cast[cast_stage_idx])); + cutlass::arch::umma_arrive(reinterpret_cast(empty_B[b_stage])); + } + cutlass::arch::umma_arrive(reinterpret_cast(tmem_full_barrier)); + } + + // ----- Epilogue (warps 0..3, 128 threads) ----- + constexpr uint32_t kNumBankGroupBytes = 16; + constexpr uint32_t kNumElemsPerBankGroup = kNumBankGroupBytes / sizeof(float); + static_assert(BLOCK_N % kNumElemsPerBankGroup == 0, "Invalid swizzling"); + + tmem_full_barrier->wait(0); + tcgen05_after_thread_sync(); + +#pragma unroll + for (uint32_t i = 0; i < BLOCK_N / kNumElemsPerBankGroup; ++i) + { + uint32_t tmem_addr = BLOCK_K * kNumCastStages + i * kNumElemsPerBankGroup; + auto smem_ptr = reinterpret_cast(smem_cd) + warp_idx * BLOCK_M / 4 * kSwizzleCDMode + + get_swizzled_smem_offset(i, lane_idx); + uint32_t values[kNumElemsPerBankGroup]; + static_assert(kNumElemsPerBankGroup == 4, "Invalid type"); + cute::SM100_TMEM_LOAD_32dp32b4x::copy(tmem_addr, values[0], values[1], values[2], values[3]); + cutlass::arch::fence_view_async_tmem_load(); + if (BLOCK_M == 128 or (BLOCK_M == 64 and lane_idx < 16)) + deep_gemm::ptx::st_shared(smem_ptr, values[0], values[1], values[2], values[3]); + if constexpr (BLOCK_M == 64) + __syncwarp(); + } + cutlass::arch::NamedBarrier::sync(kNumMMAThreads, 0); + + constexpr uint32_t kTotalOut = BLOCK_M * SHAPE_N; + const uint32_t tid = threadIdx.x; +#pragma unroll + for (uint32_t k = tid; k < kTotalOut; k += kNumMMAThreads) + { + uint32_t m = k / SHAPE_N; + uint32_t n = k - m * SHAPE_N; + uint32_t gm = m_block_idx * BLOCK_M + m; + if (gm < shape_m) + { + uint32_t col_group = n >> 2; + uint32_t in_group = n & 3; + uint32_t phys_col_grp = col_group ^ (m & 7); + uint32_t byte = m * kSwizzleCDMode + phys_col_grp * kNumBankGroupBytes + in_group * sizeof(float); + float val = *reinterpret_cast(reinterpret_cast(smem_cd) + byte); + if constexpr (kNumSplits == 1) + { + // Single-CTA-per-(m,n) at KS=1 → safe to store directly. + // Caller no longer needs to pre-zero D (see launcher). + D[gm * SHAPE_N + n] = val; + } + else + { + atomicAdd(&D[gm * SHAPE_N + n], val); + } + } + } + + if (warp_idx == 1) + cute::TMEM::Allocator1Sm().free(0, kNumTmemCols); + } + else + { + // ----- Pmap warp group (warps 4..7, 128 threads) ----- + const uint32_t sub_warp_idx = warp_idx - kNumMMAThreads / 32; + const uint32_t upper_row = sub_warp_idx * 16 + lane_idx / 4; + const uint32_t lower_row = upper_row + 8; + const uint32_t col_lane = lane_idx % 4; + + float pm_u[HC_MULT], pm_l[HC_MULT]; + float cm_u[HC_MULT][HC_MULT], cm_l[HC_MULT][HC_MULT]; +#pragma unroll + for (uint32_t hc = 0; hc < HC_MULT; ++hc) + { + pm_u[hc] = smem_post[upper_row * HC_MULT + hc]; + pm_l[hc] = smem_post[lower_row * HC_MULT + hc]; + } +#pragma unroll + for (uint32_t j = 0; j < HC_MULT; ++j) + { +#pragma unroll + for (uint32_t hc = 0; hc < HC_MULT; ++hc) + { + cm_u[j][hc] = smem_comb[upper_row * HC_MULT * HC_MULT + j * HC_MULT + hc]; + cm_l[j][hc] = smem_comb[lower_row * HC_MULT * HC_MULT + j * HC_MULT + hc]; + } + } + + float sqr_u = 0.f, sqr_l = 0.f; + constexpr uint32_t kNumBankGroupBytes = 16; + constexpr uint32_t kNumElemsPerBankGroup = kNumBankGroupBytes / sizeof(nv_bfloat16); + constexpr uint32_t kNumLoads = BLOCK_K / kNumElemsPerBankGroup; + constexpr uint32_t BLOCK_M_PER_WARP = BLOCK_M / 4; + static_assert(BLOCK_K * sizeof(nv_bfloat16) == kSwizzleAMode, "BLOCK_K must match swizzle A mode"); + static_assert(kNumLoads % 2 == 0, "kNumLoads must be even for LDSM.x4"); + + uint32_t s = 0; + for (uint32_t ht = 0; ht < H_TILES_PER_SPLIT; ++ht) + { + const uint32_t i_stage = ht % N_INPUT_STAGES; + full_input[i_stage]->wait((ht / N_INPUT_STAGES) & 1); + + uint32_t x_vals[2][kNumLoads]; + { + uint8_t const* x_base + = reinterpret_cast(smem_x_stg[i_stage]) + sub_warp_idx * BLOCK_M_PER_WARP * kSwizzleXMode; +#pragma unroll + for (uint32_t i = 0; i < kNumLoads; i += 2) + { + auto smem_ptr = x_base + get_swizzled_smem_offset(i + lane_idx / 16, lane_idx % 16); + deep_gemm::sm90::SM90_U32x4_LDSM_N::copy(x_vals[0][i + 0], x_vals[1][i + 0], x_vals[0][i + 1], + x_vals[1][i + 1], const_cast(smem_ptr)); + } + } + + uint32_t r_vals[HC_MULT][2][kNumLoads]; +#pragma unroll + for (uint32_t j = 0; j < HC_MULT; ++j) + { + uint8_t const* r_base = reinterpret_cast(smem_res[i_stage]) + + j * BLOCK_M * BLOCK_K * sizeof(nv_bfloat16) + sub_warp_idx * BLOCK_M_PER_WARP * kSwizzleResMode; +#pragma unroll + for (uint32_t i = 0; i < kNumLoads; i += 2) + { + auto smem_ptr + = r_base + get_swizzled_smem_offset(i + lane_idx / 16, lane_idx % 16); + deep_gemm::sm90::SM90_U32x4_LDSM_N::copy(r_vals[j][0][i + 0], r_vals[j][1][i + 0], + r_vals[j][0][i + 1], r_vals[j][1][i + 1], const_cast(smem_ptr)); + } + } + + float2 xf[2][kNumLoads]; +#pragma unroll + for (uint32_t u = 0; u < 2; ++u) + { +#pragma unroll + for (uint32_t i = 0; i < kNumLoads; ++i) + { + xf[u][i] = __bfloat1622float2(*reinterpret_cast(&x_vals[u][i])); + } + } + + if constexpr (kEarlyRelease) + { + empty_input[i_stage]->arrive(); + } + + // Wait for previous ht's residual_out TMA_STOREs to drain before we + // overwrite single-buffered smem_rc with new hc values. + if (ht > 0) + { + cute::tma_store_wait<0>(); + } + +#pragma unroll + for (uint32_t hc = 0; hc < HC_MULT; ++hc) + { + const uint32_t cast_stage_idx = s % kNumCastStages; + empty_cast[cast_stage_idx]->wait(((s / kNumCastStages) & 1) ^ 1); + + uint32_t rc_u_buf[kNumLoads], rc_l_buf[kNumLoads]; +#pragma unroll + for (uint32_t i = 0; i < kNumLoads; ++i) + { + float2 nu{pm_u[hc] * xf[0][i].x, pm_u[hc] * xf[0][i].y}; + float2 nl{pm_l[hc] * xf[1][i].x, pm_l[hc] * xf[1][i].y}; +#pragma unroll + for (uint32_t j = 0; j < HC_MULT; ++j) + { + float2 ruj = __bfloat1622float2(*reinterpret_cast(&r_vals[j][0][i])); + float2 rlj = __bfloat1622float2(*reinterpret_cast(&r_vals[j][1][i])); + nu.x = fmaf(cm_u[j][hc], ruj.x, nu.x); + nu.y = fmaf(cm_u[j][hc], ruj.y, nu.y); + nl.x = fmaf(cm_l[j][hc], rlj.x, nl.x); + nl.y = fmaf(cm_l[j][hc], rlj.y, nl.y); + } + nv_bfloat162 b_up = __float22bfloat162_rn(nu); + nv_bfloat162 b_lo = __float22bfloat162_rn(nl); + uint32_t b_up_bits = *reinterpret_cast(&b_up); + uint32_t b_lo_bits = *reinterpret_cast(&b_lo); + rc_u_buf[i] = b_up_bits; + rc_l_buf[i] = b_lo_bits; + float2 ru = __bfloat1622float2(b_up); + float2 rl = __bfloat1622float2(b_lo); + sqr_u = fmaf(ru.x, ru.x, sqr_u); + sqr_u = fmaf(ru.y, ru.y, sqr_u); + sqr_l = fmaf(rl.x, rl.x, sqr_l); + sqr_l = fmaf(rl.y, rl.y, sqr_l); + cute::SM100_TMEM_STORE_16dp256b1x::copy(*reinterpret_cast(&ru.x), + *reinterpret_cast(&ru.y), *reinterpret_cast(&rl.x), + *reinterpret_cast(&rl.y), cast_stage_idx * BLOCK_K + i * 8); + } + cutlass::arch::fence_view_async_tmem_store(); + tcgen05_before_thread_sync(); + full_cast[cast_stage_idx]->arrive(); + ++s; + + // STSM bf16 new_r values into smem_rc[hc] sub-region for this warp. + uint8_t* rc_base = reinterpret_cast(smem_rc) + hc * SMEM_RC_PER_HC + + sub_warp_idx * BLOCK_M_PER_WARP * kSwizzleRoutMode; +#pragma unroll + for (uint32_t i = 0; i < kNumLoads; i += 2) + { + auto smem_ptr + = rc_base + get_swizzled_smem_offset(i + lane_idx / 16, lane_idx % 16); + stsm_x4_b16_rout(smem_ptr, rc_u_buf[i + 0], rc_l_buf[i + 0], rc_u_buf[i + 1], rc_l_buf[i + 1]); + } + } + if constexpr (!kEarlyRelease) + { + empty_input[i_stage]->arrive(); + } + + // Emit HC_MULT TMA_STOREs of residual_cur: one per hc slice, per-warp rows. + cute::tma_store_fence(); + if (cute::elect_one_sync()) + { + const uint32_t h_idx = (h_tile_start + ht) * BLOCK_K; +#pragma unroll + for (uint32_t hc = 0; hc < HC_MULT; ++hc) + { + uint8_t* rc_base = reinterpret_cast(smem_rc) + hc * SMEM_RC_PER_HC + + sub_warp_idx * BLOCK_M_PER_WARP * kSwizzleRoutMode; + cute::SM90_TMA_STORE_2D::copy(&tensor_map_residual_out, rc_base, hc * HIDDEN + h_idx, + m_offset + sub_warp_idx * BLOCK_M_PER_WARP); + cute::tma_store_arrive(); + } + } + } + + // Drain any in-flight residual_out TMA stores before exit. + cute::tma_store_wait<0>(); + + // Warp-reduce sqr across 4 col_lanes then atomicAdd to global. + sqr_u += __shfl_xor_sync(0xffffffff, sqr_u, 1); + sqr_u += __shfl_xor_sync(0xffffffff, sqr_u, 2); + sqr_l += __shfl_xor_sync(0xffffffff, sqr_l, 1); + sqr_l += __shfl_xor_sync(0xffffffff, sqr_l, 2); + if (col_lane == 0) + { + uint32_t gm_u = m_block_idx * BLOCK_M + upper_row; + uint32_t gm_l = m_block_idx * BLOCK_M + lower_row; + if constexpr (kNumSplits == 1) + { + // KS=1 → only this CTA writes (gm_u, gm_l), no race possible. + if (gm_u < shape_m) + sqr_sum[gm_u] = sqr_u; + if (gm_l < shape_m) + sqr_sum[gm_l] = sqr_l; + } + else + { + if (gm_u < shape_m) + atomicAdd(&sqr_sum[gm_u], sqr_u); + if (gm_l < shape_m) + atomicAdd(&sqr_sum[gm_l], sqr_l); + } + } + } +#else + if (blockIdx.x == 0 and threadIdx.x == 0) + DG_DEVICE_ASSERT(false and "This kernel only supports sm_100a"); +#endif +} + +// ============================================================================ +// ALL-IN-ONE variant (Path D, tf32 tcgen05 MMA analogue of Path F). +// +// Single-kernel fusion of: +// 1) post_mapping : residual_cur[j,h] = post_mix_prev[j]*x_prev[h] +// + sum_k comb_mix_prev[k,j]*residual_prev[k,h] +// 2) pre_GEMM : D[i,n] = sum_{hc,h} residual_cur[i,hc,h] * W_T[n, hc*HIDDEN+h] +// sqr[i] = sum_{hc,h} residual_cur[i,hc,h]^2 +// 3) bigFuse : rmsnorm+sigmoid+sinkhorn on D/sqr -> post_mix_out, +// comb_mix_out, pre_mix; layer_input = pre_mix @ residual_cur. +// +// Semantics match Path F (fused_pmap_gemm_fma_allinone) exactly. The only +// algorithmic difference vs Path F is that we use tf32 tcgen05 MMA for the +// residual_cur @ W_T GEMM (instead of CUDA-core FMA). +// +// Pipelining, warp layout, and SMEM layout inherit from Path B +// (fused_tf32_pmap_gemm_rout_atomic_impl): the pmap warp group computes +// residual_cur, TMA-stores it to GMEM, and STSM-casts it into TMEM; the MMA +// warp group consumes TMEM for the GEMM; the epilogue atomicAdd's into +// y_acc / sqr_sum. Phase 3 elects the last-home CTA (per m_block) via +// atomicAdd on done_counter; Phase 4 runs bigfuse inline on that CTA only, +// reloading residual_cur from GMEM and writing layer_input + post_mix_out + +// comb_mix_out. +// +// Caller MUST zero D (y_acc), sqr_sum (r_acc), done_counter before launch. +// ============================================================================ + +template +__global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) + fused_allinone_tf32_pmap_gemm_atomic_impl(const uint32_t shape_m, + const __grid_constant__ cute::TmaDescriptor tensor_map_residual, // residual_prev, bf16 + const __grid_constant__ cute::TmaDescriptor tensor_map_x, // x_prev, bf16 + const __grid_constant__ cute::TmaDescriptor tensor_map_b, // W_T, tf32 + const __grid_constant__ cute::TmaDescriptor tensor_map_residual_out, // residual_cur, bf16 (TMA store) + __nv_bfloat16 const* __restrict__ residual_cur_ptr, // same buffer as TMA target + __nv_bfloat16* __restrict__ layer_input_out, // [M, HIDDEN] bf16 + float* __restrict__ D, // [M, SHAPE_N] fp32 (y_acc, caller zeros) + float* __restrict__ sqr_sum, // [M] fp32 (r_acc, caller zeros) + int* __restrict__ done_counter, // [ceil(M/BLOCK_M)] int (caller zeros) + float const* __restrict__ post_mix_prev, // [M, HC_MULT] + float const* __restrict__ comb_mix_prev, // [M, HC_MULT, HC_MULT] + float const* __restrict__ hc_scale, // [3] + float const* __restrict__ hc_base, // [HC_MULT*(2+HC_MULT)] + float* __restrict__ post_mix_out, // [M, HC_MULT] + float* __restrict__ comb_mix_out, // [M, HC_MULT, HC_MULT] + // When kFuseNorm: layer_input_out receives the RMSNorm-normalized + // values: out[t,h] = bf16(li[t,h] * rsqrt(mean(li²)+norm_eps) * w[h]). + // norm_weight must be bf16 [HIDDEN]; norm_eps is the RMSNorm epsilon. + // When kFuseNorm is false, norm_weight/norm_eps are ignored. + __nv_bfloat16 const* __restrict__ norm_weight, float norm_eps, float rms_eps, float hc_pre_eps, + float hc_sinkhorn_eps, float hc_post_mult_value, uint32_t sinkhorn_repeat) +{ +#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000) and (__CUDA_ARCH__ < 1100)) or defined(__CLION_IDE__) + using Barrier = cutlass::arch::ClusterTransactionBarrier; + + constexpr uint32_t HC_MULT2 = HC_MULT * HC_MULT; + constexpr uint32_t HC_MULT3 = HC_MULT * (2 + HC_MULT); + constexpr uint32_t SHAPE_K = HC_MULT * HIDDEN; + constexpr uint32_t H_TILES_PER_HC = HIDDEN / BLOCK_K; + static_assert(H_TILES_PER_HC % kNumSplits == 0, "H_TILES_PER_HC must be divisible by kNumSplits"); + constexpr uint32_t H_TILES_PER_SPLIT = H_TILES_PER_HC / kNumSplits; + constexpr uint32_t kNumCastStages = 4; + constexpr uint32_t kSwizzleAMode = cute::min(BLOCK_K * sizeof(nv_bfloat16), 128); + constexpr uint32_t kSwizzleBMode = cute::min(BLOCK_K * sizeof(float), 128); + constexpr uint32_t kSwizzleXMode = kSwizzleAMode; + constexpr uint32_t kSwizzleResMode = kSwizzleAMode; + constexpr uint32_t kSwizzleRoutMode = kSwizzleAMode; + constexpr auto kMajorA = cute::UMMA::Major::K; + constexpr auto kMajorB = cute::UMMA::Major::K; + static_assert(HIDDEN % BLOCK_K == 0, "HIDDEN must be multiple of BLOCK_K"); + static_assert(N_B_STAGES >= HC_MULT, "N_B_STAGES must be >= HC_MULT"); + static_assert(kSwizzleCDMode / sizeof(float) == BLOCK_N, "Invalid block N"); + static_assert(kNumMMAThreads == 128, "Invalid MMA threads"); + static_assert(kNumPmapThreads == 128, "Invalid pmap threads"); + static_assert(BLOCK_M == 64, "Invalid block M"); + static_assert(HC_MULT == 4, "Only HC_MULT=4 supported"); + static_assert(kSwizzleCDMode == 128, "Atomic variant expects kSwizzleCDMode=128"); + static_assert(SHAPE_N <= BLOCK_N, "SHAPE_N must fit within BLOCK_N"); + static_assert(SHAPE_N == HC_MULT3, "Path D expects SHAPE_N == HC_MULT*(2+HC_MULT)=24"); + + auto const warp_idx = cutlass::canonical_warp_idx_sync(); + auto const lane_idx = get_lane_idx(); + + extern __shared__ __align__(1024) uint8_t smem_buffer[]; + + constexpr uint32_t SMEM_CD_SIZE = BLOCK_M * kSwizzleCDMode; + constexpr uint32_t SMEM_B_PER_STAGE = BLOCK_N * BLOCK_K * sizeof(float); + constexpr uint32_t SMEM_RES_PER_ISTG = BLOCK_M * HC_MULT * BLOCK_K * sizeof(nv_bfloat16); + constexpr uint32_t SMEM_X_PER_ISTG = BLOCK_M * BLOCK_K * sizeof(nv_bfloat16); + constexpr uint32_t SMEM_POST_SIZE = BLOCK_M * HC_MULT * sizeof(float); + constexpr uint32_t SMEM_COMB_SIZE = BLOCK_M * HC_MULT * HC_MULT * sizeof(float); + constexpr uint32_t SMEM_RC_PER_HC = BLOCK_M * BLOCK_K * sizeof(nv_bfloat16); // 8 KB + constexpr uint32_t SMEM_RC_SIZE = HC_MULT * SMEM_RC_PER_HC; // 32 KB + + constexpr uint32_t kNumTmemCols = get_num_aligned_tmem_cols(); + + // Prefetch TMA descriptors + if (warp_idx == 0 and cute::elect_one_sync()) + { + cute::prefetch_tma_descriptor(&tensor_map_residual); + cute::prefetch_tma_descriptor(&tensor_map_x); + cute::prefetch_tma_descriptor(&tensor_map_b); + cute::prefetch_tma_descriptor(&tensor_map_residual_out); + } + + // SMEM layout: [cd, B stages, res stages, x stages, post, comb, rc (HC_MULT slices)] + auto smem_cd = reinterpret_cast(smem_buffer); + uint8_t* cursor = smem_buffer + SMEM_CD_SIZE; + auto smem_b = PatternVisitor( + [&, base = cursor](uint32_t const& i) { return reinterpret_cast(base + i * SMEM_B_PER_STAGE); }); + cursor += N_B_STAGES * SMEM_B_PER_STAGE; + auto smem_res = PatternVisitor( + [&, base = cursor](uint32_t const& i) { return reinterpret_cast(base + i * SMEM_RES_PER_ISTG); }); + cursor += N_INPUT_STAGES * SMEM_RES_PER_ISTG; + auto smem_x_stg = PatternVisitor( + [&, base = cursor](uint32_t const& i) { return reinterpret_cast(base + i * SMEM_X_PER_ISTG); }); + cursor += N_INPUT_STAGES * SMEM_X_PER_ISTG; + auto smem_post = reinterpret_cast(cursor); + cursor += SMEM_POST_SIZE; + auto smem_comb = reinterpret_cast(cursor); + cursor += SMEM_COMB_SIZE; + auto smem_rc = reinterpret_cast(cursor); // [HC_MULT][BLOCK_M][BLOCK_K] bf16 + cursor += SMEM_RC_SIZE; + + cursor = reinterpret_cast((reinterpret_cast(cursor) + 7) & ~uintptr_t(7)); + auto barrier_start_ptr = reinterpret_cast(cursor); + auto full_B = PatternVisitor([=](uint32_t const& i) { return barrier_start_ptr + i; }); + auto empty_B = PatternVisitor([=](uint32_t const& i) { return barrier_start_ptr + N_B_STAGES + i; }); + auto full_input = PatternVisitor([=](uint32_t const& i) { return barrier_start_ptr + 2 * N_B_STAGES + i; }); + auto empty_input + = PatternVisitor([=](uint32_t const& i) { return barrier_start_ptr + 2 * N_B_STAGES + N_INPUT_STAGES + i; }); + auto full_cast = PatternVisitor( + [=](uint32_t const& i) { return barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + i; }); + auto empty_cast = PatternVisitor([=](uint32_t const& i) + { return barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + kNumCastStages + i; }); + auto tmem_full_barrier = barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + 2 * kNumCastStages; + + cursor += (2 * N_B_STAGES + 2 * N_INPUT_STAGES + 2 * kNumCastStages + 1) * sizeof(Barrier); + auto tmem_ptr_in_smem = reinterpret_cast(cursor); + + if (warp_idx == 1 and cute::elect_one_sync()) + { +#pragma unroll + for (uint32_t i = 0; i < N_B_STAGES; ++i) + { + full_B[i]->init(1); + empty_B[i]->init(1); + } +#pragma unroll + for (uint32_t i = 0; i < N_INPUT_STAGES; ++i) + { + full_input[i]->init(1); + empty_input[i]->init(kNumPmapThreads); + } +#pragma unroll + for (uint32_t i = 0; i < kNumCastStages; ++i) + { + full_cast[i]->init(kNumPmapThreads); + empty_cast[i]->init(1); + } + tmem_full_barrier->init(1); + cutlass::arch::fence_barrier_init(); + } + else if (warp_idx == 2) + { + cute::TMEM::Allocator1Sm().allocate(kNumTmemCols, tmem_ptr_in_smem); + } + __syncthreads(); + + const uint32_t block_idx = __shfl_sync(0xffffffff, blockIdx.x, 0); + const uint32_t m_block_idx = block_idx / kNumSplits; + const uint32_t k_split_idx = block_idx % kNumSplits; + const uint32_t m_offset = m_block_idx * BLOCK_M; + const uint32_t h_tile_start = k_split_idx * H_TILES_PER_SPLIT; + constexpr uint32_t num_total_stages = H_TILES_PER_SPLIT * HC_MULT; + + // Prologue: pmap warp group loads post_mix_prev, comb_mix_prev into SMEM + if (warp_idx >= kNumMMAThreads / 32) + { + const uint32_t pmap_tid = threadIdx.x - kNumMMAThreads; +#pragma unroll + for (uint32_t t = 0; t < 2; ++t) + { + uint32_t idx = pmap_tid + t * kNumPmapThreads; + if (idx < BLOCK_M * HC_MULT) + { + uint32_t m = idx / HC_MULT; + uint32_t hc = idx % HC_MULT; + uint32_t gmem_m = m_offset + m; + float v = (gmem_m < shape_m) ? post_mix_prev[gmem_m * HC_MULT + hc] : 0.f; + smem_post[idx] = v; + } + } +#pragma unroll + for (uint32_t t = 0; t < 8; ++t) + { + uint32_t idx = pmap_tid + t * kNumPmapThreads; + if (idx < BLOCK_M * HC_MULT * HC_MULT) + { + uint32_t m = idx / (HC_MULT * HC_MULT); + uint32_t jk = idx % (HC_MULT * HC_MULT); + uint32_t gmem_m = m_offset + m; + float v = (gmem_m < shape_m) ? comb_mix_prev[gmem_m * HC_MULT * HC_MULT + jk] : 0.f; + smem_comb[idx] = v; + } + } + } + __syncthreads(); + + if (warp_idx < kNumMMAThreads / 32) + { + // ----- TMA warp (warp 0) ----- + if (warp_idx == 0 and cute::elect_one_sync()) + { + uint32_t b_stage = 0; + uint32_t i_stage = 0; + uint32_t s = 0; + for (uint32_t ht = 0; ht < H_TILES_PER_SPLIT; ++ht) + { + const uint32_t h_tile = h_tile_start + ht; + empty_input[i_stage]->wait(((ht / N_INPUT_STAGES) & 1) ^ 1); + uint32_t m_idx = m_block_idx * BLOCK_M; + uint32_t h_idx = h_tile * BLOCK_K; +#pragma unroll + for (uint32_t j = 0; j < HC_MULT; ++j) + { + tma_copy(&tensor_map_residual, full_input[i_stage], + smem_res[i_stage] + j * BLOCK_M * BLOCK_K, j * HIDDEN + h_idx, m_idx); + } + tma_copy( + &tensor_map_x, full_input[i_stage], smem_x_stg[i_stage], h_idx, m_idx); + constexpr uint32_t kInputBytes = SMEM_RES_PER_ISTG + SMEM_X_PER_ISTG; + full_input[i_stage]->arrive_and_expect_tx(kInputBytes); + +#pragma unroll + for (uint32_t hc = 0; hc < HC_MULT; ++hc) + { + empty_B[b_stage]->wait(((s / N_B_STAGES) & 1) ^ 1); + uint32_t k_idx = hc * HIDDEN + h_idx; + tma_copy( + &tensor_map_b, full_B[b_stage], smem_b[b_stage], k_idx, 0); + full_B[b_stage]->arrive_and_expect_tx(SMEM_B_PER_STAGE); + b_stage = (b_stage + 1) % N_B_STAGES; + ++s; + } + i_stage = (i_stage + 1) % N_INPUT_STAGES; + } + } + + // ----- MMA issue warp (warp 1) ----- + if (warp_idx == 1) + { + constexpr uint32_t UMMA_M = BLOCK_M; + constexpr uint32_t UMMA_N = BLOCK_N; + constexpr uint32_t UMMA_K = 32 / sizeof(float); + constexpr uint32_t BLOCK_SWIZZLED_BK = kSwizzleBMode / sizeof(float); + using umma_t = cute::SM100_MMA_TF32_TS; + auto instr_desc = cute::UMMA::make_instr_desc(); + auto const& runtime_instr_desc = cute::UMMA::make_runtime_instr_desc(instr_desc); + static_assert(N_B_STAGES <= 32, "Too many B stages"); + auto b_desc = make_umma_desc(smem_b[0], 0, 0); + uint32_t const& b_desc_lo = lane_idx < N_B_STAGES ? b_desc.lo + lane_idx * SMEM_B_PER_STAGE / 16 : 0u; + + for (uint32_t s = 0; s < num_total_stages; ++s) + { + const uint32_t b_stage = s % N_B_STAGES; + const uint32_t cast_stage_idx = s % kNumCastStages; + full_cast[cast_stage_idx]->wait((s / kNumCastStages) & 1); + full_B[b_stage]->wait((s / N_B_STAGES) & 1); + tcgen05_after_thread_sync(); + auto const& b_desc_base_lo = __shfl_sync(0xffffffff, b_desc_lo, static_cast(b_stage)); +#pragma unroll + for (uint32_t k = 0; k < BLOCK_K / UMMA_K; ++k) + { + uint32_t const& atom_idx = (k * UMMA_K) / BLOCK_SWIZZLED_BK; + uint32_t const& in_atom_idx = (k * UMMA_K) % BLOCK_SWIZZLED_BK; + uint32_t const& offset = atom_idx * BLOCK_N * BLOCK_SWIZZLED_BK; + b_desc.lo = advance_umma_desc_lo( + b_desc_base_lo, offset, in_atom_idx); + umma_t::fma(BLOCK_K * cast_stage_idx + k * UMMA_K, b_desc, BLOCK_K * kNumCastStages, s > 0 or k > 0, + runtime_instr_desc); + } + cutlass::arch::umma_arrive(reinterpret_cast(empty_cast[cast_stage_idx])); + cutlass::arch::umma_arrive(reinterpret_cast(empty_B[b_stage])); + } + cutlass::arch::umma_arrive(reinterpret_cast(tmem_full_barrier)); + } + + // ----- Epilogue (warps 0..3, 128 threads) ----- + constexpr uint32_t kNumBankGroupBytes = 16; + constexpr uint32_t kNumElemsPerBankGroup = kNumBankGroupBytes / sizeof(float); + static_assert(BLOCK_N % kNumElemsPerBankGroup == 0, "Invalid swizzling"); + + tmem_full_barrier->wait(0); + tcgen05_after_thread_sync(); + +#pragma unroll + for (uint32_t i = 0; i < BLOCK_N / kNumElemsPerBankGroup; ++i) + { + uint32_t tmem_addr = BLOCK_K * kNumCastStages + i * kNumElemsPerBankGroup; + auto smem_ptr = reinterpret_cast(smem_cd) + warp_idx * BLOCK_M / 4 * kSwizzleCDMode + + get_swizzled_smem_offset(i, lane_idx); + uint32_t values[kNumElemsPerBankGroup]; + static_assert(kNumElemsPerBankGroup == 4, "Invalid type"); + cute::SM100_TMEM_LOAD_32dp32b4x::copy(tmem_addr, values[0], values[1], values[2], values[3]); + cutlass::arch::fence_view_async_tmem_load(); + if (BLOCK_M == 128 or (BLOCK_M == 64 and lane_idx < 16)) + deep_gemm::ptx::st_shared(smem_ptr, values[0], values[1], values[2], values[3]); + if constexpr (BLOCK_M == 64) + __syncwarp(); + } + cutlass::arch::NamedBarrier::sync(kNumMMAThreads, 0); + + constexpr uint32_t kTotalOut = BLOCK_M * SHAPE_N; + const uint32_t tid = threadIdx.x; +#pragma unroll + for (uint32_t k = tid; k < kTotalOut; k += kNumMMAThreads) + { + uint32_t m = k / SHAPE_N; + uint32_t n = k - m * SHAPE_N; + uint32_t gm = m_block_idx * BLOCK_M + m; + if (gm < shape_m) + { + uint32_t col_group = n >> 2; + uint32_t in_group = n & 3; + uint32_t phys_col_grp = col_group ^ (m & 7); + uint32_t byte = m * kSwizzleCDMode + phys_col_grp * kNumBankGroupBytes + in_group * sizeof(float); + float val = *reinterpret_cast(reinterpret_cast(smem_cd) + byte); + if constexpr (kNumSplits == 1) + { + // Single-CTA-per-(m,n) at KS=1 → safe to store directly. + // Caller no longer needs to pre-zero D (see launcher). + D[gm * SHAPE_N + n] = val; + } + else + { + atomicAdd(&D[gm * SHAPE_N + n], val); + } + } + } + + if (warp_idx == 1) + cute::TMEM::Allocator1Sm().free(0, kNumTmemCols); + } + else + { + // ----- Pmap warp group (warps 4..7, 128 threads) ----- + const uint32_t sub_warp_idx = warp_idx - kNumMMAThreads / 32; + const uint32_t upper_row = sub_warp_idx * 16 + lane_idx / 4; + const uint32_t lower_row = upper_row + 8; + const uint32_t col_lane = lane_idx % 4; + + float pm_u[HC_MULT], pm_l[HC_MULT]; + float cm_u[HC_MULT][HC_MULT], cm_l[HC_MULT][HC_MULT]; +#pragma unroll + for (uint32_t hc = 0; hc < HC_MULT; ++hc) + { + pm_u[hc] = smem_post[upper_row * HC_MULT + hc]; + pm_l[hc] = smem_post[lower_row * HC_MULT + hc]; + } +#pragma unroll + for (uint32_t j = 0; j < HC_MULT; ++j) + { +#pragma unroll + for (uint32_t hc = 0; hc < HC_MULT; ++hc) + { + cm_u[j][hc] = smem_comb[upper_row * HC_MULT * HC_MULT + j * HC_MULT + hc]; + cm_l[j][hc] = smem_comb[lower_row * HC_MULT * HC_MULT + j * HC_MULT + hc]; + } + } + + float sqr_u = 0.f, sqr_l = 0.f; + constexpr uint32_t kNumBankGroupBytes = 16; + constexpr uint32_t kNumElemsPerBankGroup = kNumBankGroupBytes / sizeof(nv_bfloat16); + constexpr uint32_t kNumLoads = BLOCK_K / kNumElemsPerBankGroup; + constexpr uint32_t BLOCK_M_PER_WARP = BLOCK_M / 4; + static_assert(BLOCK_K * sizeof(nv_bfloat16) == kSwizzleAMode, "BLOCK_K must match swizzle A mode"); + static_assert(kNumLoads % 2 == 0, "kNumLoads must be even for LDSM.x4"); + + uint32_t s = 0; + for (uint32_t ht = 0; ht < H_TILES_PER_SPLIT; ++ht) + { + const uint32_t i_stage = ht % N_INPUT_STAGES; + full_input[i_stage]->wait((ht / N_INPUT_STAGES) & 1); + + uint32_t x_vals[2][kNumLoads]; + { + uint8_t const* x_base + = reinterpret_cast(smem_x_stg[i_stage]) + sub_warp_idx * BLOCK_M_PER_WARP * kSwizzleXMode; +#pragma unroll + for (uint32_t i = 0; i < kNumLoads; i += 2) + { + auto smem_ptr = x_base + get_swizzled_smem_offset(i + lane_idx / 16, lane_idx % 16); + deep_gemm::sm90::SM90_U32x4_LDSM_N::copy(x_vals[0][i + 0], x_vals[1][i + 0], x_vals[0][i + 1], + x_vals[1][i + 1], const_cast(smem_ptr)); + } + } + + uint32_t r_vals[HC_MULT][2][kNumLoads]; +#pragma unroll + for (uint32_t j = 0; j < HC_MULT; ++j) + { + uint8_t const* r_base = reinterpret_cast(smem_res[i_stage]) + + j * BLOCK_M * BLOCK_K * sizeof(nv_bfloat16) + sub_warp_idx * BLOCK_M_PER_WARP * kSwizzleResMode; +#pragma unroll + for (uint32_t i = 0; i < kNumLoads; i += 2) + { + auto smem_ptr + = r_base + get_swizzled_smem_offset(i + lane_idx / 16, lane_idx % 16); + deep_gemm::sm90::SM90_U32x4_LDSM_N::copy(r_vals[j][0][i + 0], r_vals[j][1][i + 0], + r_vals[j][0][i + 1], r_vals[j][1][i + 1], const_cast(smem_ptr)); + } + } + + float2 xf[2][kNumLoads]; +#pragma unroll + for (uint32_t u = 0; u < 2; ++u) + { +#pragma unroll + for (uint32_t i = 0; i < kNumLoads; ++i) + { + xf[u][i] = __bfloat1622float2(*reinterpret_cast(&x_vals[u][i])); + } + } + + // Wait for previous ht's residual_out TMA_STOREs to drain before we + // overwrite single-buffered smem_rc with new hc values. + if (ht > 0) + { + cute::tma_store_wait<0>(); + } + +#pragma unroll + for (uint32_t hc = 0; hc < HC_MULT; ++hc) + { + const uint32_t cast_stage_idx = s % kNumCastStages; + empty_cast[cast_stage_idx]->wait(((s / kNumCastStages) & 1) ^ 1); + + uint32_t rc_u_buf[kNumLoads], rc_l_buf[kNumLoads]; +#pragma unroll + for (uint32_t i = 0; i < kNumLoads; ++i) + { + float2 nu{pm_u[hc] * xf[0][i].x, pm_u[hc] * xf[0][i].y}; + float2 nl{pm_l[hc] * xf[1][i].x, pm_l[hc] * xf[1][i].y}; +#pragma unroll + for (uint32_t j = 0; j < HC_MULT; ++j) + { + float2 ruj = __bfloat1622float2(*reinterpret_cast(&r_vals[j][0][i])); + float2 rlj = __bfloat1622float2(*reinterpret_cast(&r_vals[j][1][i])); + nu.x = fmaf(cm_u[j][hc], ruj.x, nu.x); + nu.y = fmaf(cm_u[j][hc], ruj.y, nu.y); + nl.x = fmaf(cm_l[j][hc], rlj.x, nl.x); + nl.y = fmaf(cm_l[j][hc], rlj.y, nl.y); + } + nv_bfloat162 b_up = __float22bfloat162_rn(nu); + nv_bfloat162 b_lo = __float22bfloat162_rn(nl); + uint32_t b_up_bits = *reinterpret_cast(&b_up); + uint32_t b_lo_bits = *reinterpret_cast(&b_lo); + rc_u_buf[i] = b_up_bits; + rc_l_buf[i] = b_lo_bits; + float2 ru = __bfloat1622float2(b_up); + float2 rl = __bfloat1622float2(b_lo); + sqr_u = fmaf(ru.x, ru.x, sqr_u); + sqr_u = fmaf(ru.y, ru.y, sqr_u); + sqr_l = fmaf(rl.x, rl.x, sqr_l); + sqr_l = fmaf(rl.y, rl.y, sqr_l); + cute::SM100_TMEM_STORE_16dp256b1x::copy(*reinterpret_cast(&ru.x), + *reinterpret_cast(&ru.y), *reinterpret_cast(&rl.x), + *reinterpret_cast(&rl.y), cast_stage_idx * BLOCK_K + i * 8); + } + cutlass::arch::fence_view_async_tmem_store(); + tcgen05_before_thread_sync(); + full_cast[cast_stage_idx]->arrive(); + ++s; + + // STSM bf16 new_r values into smem_rc[hc] sub-region for this warp. + uint8_t* rc_base = reinterpret_cast(smem_rc) + hc * SMEM_RC_PER_HC + + sub_warp_idx * BLOCK_M_PER_WARP * kSwizzleRoutMode; +#pragma unroll + for (uint32_t i = 0; i < kNumLoads; i += 2) + { + auto smem_ptr + = rc_base + get_swizzled_smem_offset(i + lane_idx / 16, lane_idx % 16); + stsm_x4_b16_rout(smem_ptr, rc_u_buf[i + 0], rc_l_buf[i + 0], rc_u_buf[i + 1], rc_l_buf[i + 1]); + } + } + empty_input[i_stage]->arrive(); + + // Emit HC_MULT TMA_STOREs of residual_cur: one per hc slice, per-warp rows. + cute::tma_store_fence(); + if (cute::elect_one_sync()) + { + const uint32_t h_idx = (h_tile_start + ht) * BLOCK_K; +#pragma unroll + for (uint32_t hc = 0; hc < HC_MULT; ++hc) + { + uint8_t* rc_base = reinterpret_cast(smem_rc) + hc * SMEM_RC_PER_HC + + sub_warp_idx * BLOCK_M_PER_WARP * kSwizzleRoutMode; + cute::SM90_TMA_STORE_2D::copy(&tensor_map_residual_out, rc_base, hc * HIDDEN + h_idx, + m_offset + sub_warp_idx * BLOCK_M_PER_WARP); + cute::tma_store_arrive(); + } + } + } + + // Drain any in-flight residual_out TMA stores before exit. + cute::tma_store_wait<0>(); + + // Warp-reduce sqr across 4 col_lanes then atomicAdd to global. + sqr_u += __shfl_xor_sync(0xffffffff, sqr_u, 1); + sqr_u += __shfl_xor_sync(0xffffffff, sqr_u, 2); + sqr_l += __shfl_xor_sync(0xffffffff, sqr_l, 1); + sqr_l += __shfl_xor_sync(0xffffffff, sqr_l, 2); + if (col_lane == 0) + { + uint32_t gm_u = m_block_idx * BLOCK_M + upper_row; + uint32_t gm_l = m_block_idx * BLOCK_M + lower_row; + if constexpr (kNumSplits == 1) + { + // KS=1 → only this CTA writes (gm_u, gm_l), no race possible. + if (gm_u < shape_m) + sqr_sum[gm_u] = sqr_u; + if (gm_l < shape_m) + sqr_sum[gm_l] = sqr_l; + } + else + { + if (gm_u < shape_m) + atomicAdd(&sqr_sum[gm_u], sqr_u); + if (gm_l < shape_m) + atomicAdd(&sqr_sum[gm_l], sqr_l); + } + } + } + + // ======================================================================== + // Phase 3: cross-split barrier. + // For kNumSplits == 1 we only need a block-scope fence + __syncthreads. + // For kNumSplits > 1 ALL splits participate in Phase 4: each CTA + // increments done_counter, then spin-waits until the counter reaches + // kNumSplits. Phase 4 work is then partitioned across CTAs by + // k_split_idx — CTA i processes tokens + // [i * TOKS_PER_CTA, (i+1) * TOKS_PER_CTA) + // where TOKS_PER_CTA = BLOCK_M / kNumSplits. This replaces the old + // "last-home CTA does all Phase 4 work serially" design that bottlenecked + // Path D at BLOCK_M=64 (1 CTA's Phase 4 = ~28 µs regardless of M). + // ======================================================================== + if constexpr (kNumSplits == 1) + { + __threadfence_block(); + __syncthreads(); + } + else + { + __threadfence(); + __syncthreads(); + if (threadIdx.x == 0) + { + atomicAdd(&done_counter[m_block_idx], 1); + // Spin-wait until all kNumSplits CTAs finish Phase 2. The + // atomicAdd(..., 0) is a zero-increment load with full device + // coherence — cheap on B200 and avoids an extra flag allocation. + while (atomicAdd(&done_counter[m_block_idx], 0) < static_cast(kNumSplits)) + { + /* spin */ + } + } + __syncthreads(); + } + + // ======================================================================== + // Phase 4: inline bigFuse for this CTA's subset of BLOCK_M tokens. + // y_acc[tok, 0..HC_MULT) -> pre_mix (sigmoid) + // y_acc[tok, HC_MULT..2*HC_MULT) -> post_mix_out + // y_acc[tok, 2*HC_MULT..HC_MULT3) -> comb_mix_out (sinkhorn) + // layer_input[tok, h] = sum_j pre_mix[tok, j] * residual_cur[tok, j, h] + // + // Token subset: CTA k_split_idx handles tokens + // [k_split_idx * TOKS_PER_CTA, (k_split_idx + 1) * TOKS_PER_CTA), + // where TOKS_PER_CTA = ceil(BLOCK_M / kNumSplits). At kNumSplits == 1 + // this is the whole m_block (64 tokens); at kNumSplits == 8 it's 8 tokens + // spread over 8 CTAs running Phase 4 concurrently — ~8× the Phase-4 + // throughput of the old single-last-home-CTA design. + // + // Within a CTA, tokens are parallelized across warps: each warp handles + // max(1, TOKS_PER_CTA / NUM_WARPS_BF) tokens serially. Within a warp, + // lanes 0..HC_MULT-1 compute rmsnorm/sigmoid/sinkhorn; pre_mix is + // broadcast via __shfl_sync so all 32 lanes run the layer_input dot + // product across HIDDEN=4096. + // ======================================================================== + constexpr uint32_t BLOCK_SIZE_BF = kNumMMAThreads + kNumPmapThreads; // 256 + constexpr uint32_t WARP_SIZE_BF = 32; + constexpr uint32_t NUM_WARPS_BF = BLOCK_SIZE_BF / WARP_SIZE_BF; // 8 + constexpr uint32_t TOKS_PER_CTA = (BLOCK_M + kNumSplits - 1) / kNumSplits; + // When TOKS_PER_CTA < NUM_WARPS_BF (large kNumSplits / small BLOCK_M case), + // have WARPS_PER_TOK warps cooperate on the HIDDEN-stride layer_input loop + // so all 8 warps stay active. Example at BLOCK_M=64, kNumSplits=16: + // TOKS_PER_CTA=4, WARPS_PER_TOK=2, TOKS_PER_PASS=4, TOKEN_PASSES=1 — + // 2 warps per token each sweep HIDDEN/2, zero idle warps. + constexpr uint32_t WARPS_PER_TOK = (NUM_WARPS_BF > TOKS_PER_CTA) ? (NUM_WARPS_BF / TOKS_PER_CTA) : 1u; + constexpr uint32_t TOKS_PER_PASS = NUM_WARPS_BF / WARPS_PER_TOK; + constexpr uint32_t TOKEN_PASSES = (TOKS_PER_CTA + TOKS_PER_PASS - 1) / TOKS_PER_PASS; + constexpr uint32_t BF16_VEC_LI = 8; + // The vectorized loop covers floor(HIDDEN / H_STRIDE) * H_STRIDE elements. + // Anything past that is handled by a scalar-vec tail loop below — relax the + // old static_assert so HIDDEN values like 7168 (which 8-warp teams cannot + // cleanly span) are also valid. + static_assert(HIDDEN % BF16_VEC_LI == 0, "HIDDEN must be a multiple of BF16_VEC_LI=8"); + const uint32_t tid_bf = threadIdx.x; + const uint32_t lane_bf = tid_bf % WARP_SIZE_BF; + const uint32_t warp_bf = tid_bf / WARP_SIZE_BF; + const uint32_t warp_tok_pos = warp_bf / WARPS_PER_TOK; // which token in a pass + const uint32_t warp_in_team = warp_bf % WARPS_PER_TOK; // which warp inside team + const uint32_t cta_tok_base = k_split_idx * TOKS_PER_CTA; + +#pragma unroll 1 + for (uint32_t pass = 0; pass < TOKEN_PASSES; ++pass) + { + const uint32_t t_in_cta = pass * TOKS_PER_PASS + warp_tok_pos; + if (t_in_cta >= TOKS_PER_CTA) + continue; + const uint32_t t = cta_tok_base + t_in_cta; + if (t >= BLOCK_M) + continue; + const uint32_t tok = m_offset + t; + if (tok >= shape_m) + continue; + + // Lanes 0..HC_MULT-1 compute rmsnorm / sigmoid / sinkhorn; pre_mix is + // held in `pre_mix_local` on lanes 0..HC_MULT-1 and later broadcast to + // all 32 lanes via __shfl_sync. All warps in a team redundantly run + // these ~tens of FLOPs (cheap) to avoid a cross-warp SMEM sync; only + // warp_in_team==0 writes comb_mix_out / post_mix_out to GMEM. + float pre_mix_local = 0.f; + if (lane_bf < HC_MULT) + { + float const r_val = sqr_sum[tok]; + float y_local[HC_MULT3]; + float const* y_row = D + static_cast(tok) * SHAPE_N; +#pragma unroll + for (uint32_t c = 0; c < HC_MULT3; ++c) + y_local[c] = y_row[c]; + + float const rstd = rsqrtf(r_val / static_cast(HC_MULT * HIDDEN) + rms_eps); + float const s0 = hc_scale[0]; + float const s1 = hc_scale[1]; + float const s2 = hc_scale[2]; + + float v = y_local[lane_bf] * rstd * s0 + hc_base[lane_bf]; + pre_mix_local = 1.0f / (1.0f + __expf(-v)) + hc_pre_eps; + + v = y_local[HC_MULT + lane_bf] * rstd * s1 + hc_base[HC_MULT + lane_bf]; + float post_val = 1.0f / (1.0f + __expf(-v)) * hc_post_mult_value; + if (warp_in_team == 0) + { + post_mix_out[tok * HC_MULT + lane_bf] = post_val; + } + + float cm_vals[HC_MULT]; +#pragma unroll + for (uint32_t k = 0; k < HC_MULT; ++k) + cm_vals[k] = y_local[2 * HC_MULT + lane_bf * HC_MULT + k] * rstd * s2 + + hc_base[2 * HC_MULT + lane_bf * HC_MULT + k]; + + constexpr unsigned LANE_MASK = (1u << HC_MULT) - 1; + float const rowMax = fmaxf(fmaxf(cm_vals[0], cm_vals[1]), fmaxf(cm_vals[2], cm_vals[3])); +#pragma unroll + for (uint32_t k = 0; k < HC_MULT; ++k) + cm_vals[k] = __expf(cm_vals[k] - rowMax); + // Reciprocal-multiply for sinkhorn: 1 fdiv + 4 fmul instead of 4 + // fdivs per row-normalize. Equivalent under fp32 round-off. + float inv_rs = 1.0f / (cm_vals[0] + cm_vals[1] + cm_vals[2] + cm_vals[3]); +#pragma unroll + for (uint32_t k = 0; k < HC_MULT; ++k) + cm_vals[k] = cm_vals[k] * inv_rs + hc_sinkhorn_eps; +#pragma unroll + for (uint32_t k = 0; k < HC_MULT; ++k) + { + float cs = cm_vals[k]; + cs += __shfl_xor_sync(LANE_MASK, cs, 1); + cs += __shfl_xor_sync(LANE_MASK, cs, 2); + cm_vals[k] *= 1.0f / (cs + hc_sinkhorn_eps); + } + for (uint32_t it = 1; it < sinkhorn_repeat; ++it) + { + inv_rs = 1.0f / (cm_vals[0] + cm_vals[1] + cm_vals[2] + cm_vals[3] + hc_sinkhorn_eps); +#pragma unroll + for (uint32_t k = 0; k < HC_MULT; ++k) + cm_vals[k] *= inv_rs; +#pragma unroll + for (uint32_t k = 0; k < HC_MULT; ++k) + { + float cs = cm_vals[k]; + cs += __shfl_xor_sync(LANE_MASK, cs, 1); + cs += __shfl_xor_sync(LANE_MASK, cs, 2); + cm_vals[k] *= 1.0f / (cs + hc_sinkhorn_eps); + } + } + if (warp_in_team == 0) + { + float* cm_out_ptr = comb_mix_out + tok * HC_MULT2; +#pragma unroll + for (uint32_t k = 0; k < HC_MULT; ++k) + cm_out_ptr[lane_bf * HC_MULT + k] = cm_vals[k]; + } + } + + // Broadcast pre_mix[j] from lane j to all 32 lanes (intra-warp shfl). + float pm[HC_MULT]; +#pragma unroll + for (uint32_t j = 0; j < HC_MULT; ++j) + pm[j] = __shfl_sync(0xffffffff, pre_mix_local, j); + + // Layer_input[tok, h] = sum_j pm[j] * residual_cur[tok, j, h]. + // When WARPS_PER_TOK>1, warp_in_team 0..WARPS_PER_TOK-1 together cover + // HIDDEN in strides of WARPS_PER_TOK * 32 * 8. When WARPS_PER_TOK==1, + // each warp sweeps HIDDEN alone (same as the original single-warp case). + __nv_bfloat16 const* rbase = residual_cur_ptr + static_cast(tok) * HC_MULT * HIDDEN; + __nv_bfloat16* obase = layer_input_out + static_cast(tok) * HIDDEN; + + constexpr uint32_t H_STRIDE = WARPS_PER_TOK * WARP_SIZE_BF * BF16_VEC_LI; + // Largest multiple of H_STRIDE that fits in HIDDEN — after this comes + // the scalar-vec tail (only relevant when WARPS_PER_TOK*32*8 > HIDDEN + // residue, e.g. KS=112 / HIDDEN=7168 / H_STRIDE=2048 → tail = 1024). + constexpr uint32_t H_VEC_END = (HIDDEN / H_STRIDE) * H_STRIDE; + const uint32_t h_start = warp_in_team * WARP_SIZE_BF * BF16_VEC_LI + lane_bf * BF16_VEC_LI; + + if constexpr (!kFuseNorm) + { +#pragma unroll + for (uint32_t h = h_start; h < H_VEC_END; h += H_STRIDE) + { + // Issue all HC_MULT=4 residual_cur reads first so their L2 latency + // is hidden by the bf16→fp32 arithmetic that follows. The compiler + // schedules the 4 independent __ldg's in parallel. + uint4 raws[HC_MULT]; +#pragma unroll + for (uint32_t j = 0; j < HC_MULT; ++j) + { + raws[j] = __ldg(reinterpret_cast(&rbase[j * HIDDEN + h])); + } + float acc_li[BF16_VEC_LI] = {}; +#pragma unroll + for (uint32_t j = 0; j < HC_MULT; ++j) + { + __nv_bfloat162 const* pairs = reinterpret_cast<__nv_bfloat162 const*>(&raws[j]); +#pragma unroll + for (uint32_t v = 0; v < BF16_VEC_LI / 2; ++v) + { + float2 f = __bfloat1622float2(pairs[v]); + acc_li[2 * v + 0] += pm[j] * f.x; + acc_li[2 * v + 1] += pm[j] * f.y; + } + } + uint4 out_raw; + __nv_bfloat162* opairs = reinterpret_cast<__nv_bfloat162*>(&out_raw); +#pragma unroll + for (uint32_t v = 0; v < BF16_VEC_LI / 2; ++v) + opairs[v] = __float22bfloat162_rn(make_float2(acc_li[2 * v], acc_li[2 * v + 1])); + *reinterpret_cast(&obase[h]) = out_raw; + } + + // Scalar-vec tail: hidden residue [H_VEC_END, HIDDEN) is shorter than a + // full team stride. Distribute leftover BF16_VEC_LI-sized chunks across + // the team's threads in lane-major order (skip threads whose chunk + // falls past HIDDEN). Each chunk is still a single uint4 LDG/STG, so + // tail throughput matches the main loop's per-thread bandwidth — the + // only loss is that some lanes/warps in the team idle. + if constexpr (H_VEC_END < HIDDEN) + { + constexpr uint32_t TAIL_CHUNKS = (HIDDEN - H_VEC_END) / BF16_VEC_LI; + const uint32_t my_chunk = warp_in_team * WARP_SIZE_BF + lane_bf; + if (my_chunk < TAIL_CHUNKS) + { + const uint32_t h = H_VEC_END + my_chunk * BF16_VEC_LI; + uint4 raws[HC_MULT]; +#pragma unroll + for (uint32_t j = 0; j < HC_MULT; ++j) + { + raws[j] = __ldg(reinterpret_cast(&rbase[j * HIDDEN + h])); + } + float acc_li[BF16_VEC_LI] = {}; +#pragma unroll + for (uint32_t j = 0; j < HC_MULT; ++j) + { + __nv_bfloat162 const* pairs = reinterpret_cast<__nv_bfloat162 const*>(&raws[j]); +#pragma unroll + for (uint32_t v = 0; v < BF16_VEC_LI / 2; ++v) + { + float2 f = __bfloat1622float2(pairs[v]); + acc_li[2 * v + 0] += pm[j] * f.x; + acc_li[2 * v + 1] += pm[j] * f.y; + } + } + uint4 out_raw; + __nv_bfloat162* opairs = reinterpret_cast<__nv_bfloat162*>(&out_raw); +#pragma unroll + for (uint32_t v = 0; v < BF16_VEC_LI / 2; ++v) + opairs[v] = __float22bfloat162_rn(make_float2(acc_li[2 * v], acc_li[2 * v + 1])); + *reinterpret_cast(&obase[h]) = out_raw; + } + } + } + else + { + // ----------------------------------------------------------------- + // Fused RMSNorm path: layer_input[t,h] = bf16(li[t,h] * rsqrt( + // mean(li²)+norm_eps) * norm_weight[h]). + // + // Pass 1: identical to the un-fused path — compute li per chunk + // and STG to layer_input_out as bf16. The bf16 store lands in L2; + // pass 2 re-reads from L2 (no extra HBM read in the steady case). + // In parallel we accumulate per-thread `sum_sq_local`. + // Reduce: intra-warp __shfl_xor; cross-warp via SMEM + per-team + // named PTX barrier when WARPS_PER_TOK > 1 (KS ≥ 16 instances). + // Inactive teams `continue`'d above and never reach this barrier. + // Pass 2: re-LDG layer_input_out from L2, multiply by + // rsqrt * norm_weight, STG normalized bf16 back to the same + // address. Avoids the FMA recompute that doubling pass 1 would + // require (Path D Phase 4 is already FMA-heavy). + // + // Saves the separate RMSNorm kernel launch + its HBM round-trip + // (~14 KB read + ~14 KB write per token) vs the un-fused path. + // ----------------------------------------------------------------- + float sum_sq_local = 0.f; + +#pragma unroll + for (uint32_t h = h_start; h < H_VEC_END; h += H_STRIDE) + { + uint4 raws[HC_MULT]; +#pragma unroll + for (uint32_t j = 0; j < HC_MULT; ++j) + raws[j] = __ldg(reinterpret_cast(&rbase[j * HIDDEN + h])); + float acc_li[BF16_VEC_LI] = {}; +#pragma unroll + for (uint32_t j = 0; j < HC_MULT; ++j) + { + __nv_bfloat162 const* pairs = reinterpret_cast<__nv_bfloat162 const*>(&raws[j]); +#pragma unroll + for (uint32_t v = 0; v < BF16_VEC_LI / 2; ++v) + { + float2 f = __bfloat1622float2(pairs[v]); + acc_li[2 * v + 0] += pm[j] * f.x; + acc_li[2 * v + 1] += pm[j] * f.y; + } + } + // Round-to-bf16 *before* squaring so sum_sq matches the value + // we actually store (a separate RMSNorm kernel would also + // compute sum_sq from the bf16-rounded layer_input). + uint4 out_raw; + __nv_bfloat162* opairs = reinterpret_cast<__nv_bfloat162*>(&out_raw); +#pragma unroll + for (uint32_t v = 0; v < BF16_VEC_LI / 2; ++v) + opairs[v] = __float22bfloat162_rn(make_float2(acc_li[2 * v], acc_li[2 * v + 1])); + *reinterpret_cast(&obase[h]) = out_raw; +#pragma unroll + for (uint32_t v = 0; v < BF16_VEC_LI / 2; ++v) + { + float2 b = __bfloat1622float2(opairs[v]); + sum_sq_local += b.x * b.x + b.y * b.y; + } + } + if constexpr (H_VEC_END < HIDDEN) + { + constexpr uint32_t TAIL_CHUNKS = (HIDDEN - H_VEC_END) / BF16_VEC_LI; + const uint32_t my_chunk = warp_in_team * WARP_SIZE_BF + lane_bf; + if (my_chunk < TAIL_CHUNKS) + { + const uint32_t h = H_VEC_END + my_chunk * BF16_VEC_LI; + uint4 raws[HC_MULT]; +#pragma unroll + for (uint32_t j = 0; j < HC_MULT; ++j) + raws[j] = __ldg(reinterpret_cast(&rbase[j * HIDDEN + h])); + float acc_li[BF16_VEC_LI] = {}; +#pragma unroll + for (uint32_t j = 0; j < HC_MULT; ++j) + { + __nv_bfloat162 const* pairs = reinterpret_cast<__nv_bfloat162 const*>(&raws[j]); +#pragma unroll + for (uint32_t v = 0; v < BF16_VEC_LI / 2; ++v) + { + float2 f = __bfloat1622float2(pairs[v]); + acc_li[2 * v + 0] += pm[j] * f.x; + acc_li[2 * v + 1] += pm[j] * f.y; + } + } + uint4 out_raw; + __nv_bfloat162* opairs = reinterpret_cast<__nv_bfloat162*>(&out_raw); +#pragma unroll + for (uint32_t v = 0; v < BF16_VEC_LI / 2; ++v) + opairs[v] = __float22bfloat162_rn(make_float2(acc_li[2 * v], acc_li[2 * v + 1])); + *reinterpret_cast(&obase[h]) = out_raw; +#pragma unroll + for (uint32_t v = 0; v < BF16_VEC_LI / 2; ++v) + { + float2 b = __bfloat1622float2(opairs[v]); + sum_sq_local += b.x * b.x + b.y * b.y; + } + } + } + + // Intra-warp reduce sum_sq → one value broadcast to all 32 lanes. + sum_sq_local += __shfl_xor_sync(0xffffffff, sum_sq_local, 16); + sum_sq_local += __shfl_xor_sync(0xffffffff, sum_sq_local, 8); + sum_sq_local += __shfl_xor_sync(0xffffffff, sum_sq_local, 4); + sum_sq_local += __shfl_xor_sync(0xffffffff, sum_sq_local, 2); + sum_sq_local += __shfl_xor_sync(0xffffffff, sum_sq_local, 1); + + // Cross-warp reduce when WARPS_PER_TOK > 1. Use a per-team named + // PTX barrier (bar.sync %0, %1) instead of __syncthreads so that + // inactive teams (which `continue`'d above) don't deadlock. + // Barrier IDs 1..TOKS_PER_PASS are private to each team. + if constexpr (WARPS_PER_TOK > 1) + { + __shared__ float team_sumsq[TOKS_PER_PASS][WARPS_PER_TOK]; + if (lane_bf == 0) + team_sumsq[warp_tok_pos][warp_in_team] = sum_sq_local; + asm volatile("bar.sync %0, %1;" ::"r"(warp_tok_pos + 1u), + "n"(static_cast(WARPS_PER_TOK * WARP_SIZE_BF))); + float team_sum = 0.f; +#pragma unroll + for (uint32_t w = 0; w < WARPS_PER_TOK; ++w) + team_sum += team_sumsq[warp_tok_pos][w]; + sum_sq_local = team_sum; + } + + float const rsqrt_val = rsqrtf(sum_sq_local / static_cast(HIDDEN) + norm_eps); + + // Pass 2: re-LDG the un-normalized layer_input we just wrote + // (L2-hot), LDG norm_weight, normalize, STG back. No FMA recompute. + __nv_bfloat16 const* nbase = norm_weight; +#pragma unroll + for (uint32_t h = h_start; h < H_VEC_END; h += H_STRIDE) + { + uint4 li_raw = __ldg(reinterpret_cast(&obase[h])); + uint4 nw_raw = __ldg(reinterpret_cast(&nbase[h])); + __nv_bfloat162 const* li_pairs = reinterpret_cast<__nv_bfloat162 const*>(&li_raw); + __nv_bfloat162 const* nw_pairs = reinterpret_cast<__nv_bfloat162 const*>(&nw_raw); + uint4 out_raw; + __nv_bfloat162* opairs = reinterpret_cast<__nv_bfloat162*>(&out_raw); +#pragma unroll + for (uint32_t v = 0; v < BF16_VEC_LI / 2; ++v) + { + float2 li_f = __bfloat1622float2(li_pairs[v]); + float2 nw_f = __bfloat1622float2(nw_pairs[v]); + opairs[v] + = __float22bfloat162_rn(make_float2(li_f.x * rsqrt_val * nw_f.x, li_f.y * rsqrt_val * nw_f.y)); + } + *reinterpret_cast(&obase[h]) = out_raw; + } + if constexpr (H_VEC_END < HIDDEN) + { + constexpr uint32_t TAIL_CHUNKS = (HIDDEN - H_VEC_END) / BF16_VEC_LI; + const uint32_t my_chunk = warp_in_team * WARP_SIZE_BF + lane_bf; + if (my_chunk < TAIL_CHUNKS) + { + const uint32_t h = H_VEC_END + my_chunk * BF16_VEC_LI; + uint4 li_raw = __ldg(reinterpret_cast(&obase[h])); + uint4 nw_raw = __ldg(reinterpret_cast(&nbase[h])); + __nv_bfloat162 const* li_pairs = reinterpret_cast<__nv_bfloat162 const*>(&li_raw); + __nv_bfloat162 const* nw_pairs = reinterpret_cast<__nv_bfloat162 const*>(&nw_raw); + uint4 out_raw; + __nv_bfloat162* opairs = reinterpret_cast<__nv_bfloat162*>(&out_raw); +#pragma unroll + for (uint32_t v = 0; v < BF16_VEC_LI / 2; ++v) + { + float2 li_f = __bfloat1622float2(li_pairs[v]); + float2 nw_f = __bfloat1622float2(nw_pairs[v]); + opairs[v] = __float22bfloat162_rn( + make_float2(li_f.x * rsqrt_val * nw_f.x, li_f.y * rsqrt_val * nw_f.y)); + } + *reinterpret_cast(&obase[h]) = out_raw; + } + } + } + } +#else + if (blockIdx.x == 0 and threadIdx.x == 0) + DG_DEVICE_ASSERT(false and "This kernel only supports sm_100a"); +#endif +} + +} // namespace fused_mhc + +#pragma clang diagnostic pop diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu b/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu new file mode 100644 index 000000000000..3be14448766b --- /dev/null +++ b/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu @@ -0,0 +1,865 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * 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. + */ + +// Single-launch fused hyper-connection boundary op. +// +// This implementation wraps the SM100 tcgen05-based residual-out variant +// (fused_tf32_pmap_gemm_rout_atomic_impl) to produce residual_cur, D, and +// sqr_sum in a single kernel launch; the big-fuse postlogue kernel then +// consumes (D, sqr_sum, residual_cur) to emit (post_mix_cur, comb_mix_cur, +// layer_input_cur). Semantically identical to +// +// residual_cur = prev_mHC.post_mapping(x_prev, residual_prev, ...) +// post_mix_cur, comb_mix_cur, layer_input_cur = self.pre_mapping(residual_cur) +// +// but exposed as a single entry point. + +#ifdef TRTLLM_MHC_ENABLE_FUSED_HC +#include "fused_tf32_pmap_gemm.cuh" +#endif +#include "mhcKernels.h" +#include "mhc_fused_fma.cuh" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/envUtils.h" + +#include +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::mhc +{ + +// ---- Single-launch workspace zero kernel ----------------------------------- +// +// Replaces 2-3 separate cudaMemsetAsync calls for the atomic accumulator +// workspaces (y_acc, r_acc, optional done_counter). Avoids per-memset launch +// latency that is visible at small M / high-frequency inference. +namespace +{ + +__global__ void fhcZeroWorkspacesKernel(float* __restrict__ y_acc, uint32_t y_elems, float* __restrict__ r_acc, + uint32_t r_elems, int* __restrict__ done_counter, uint32_t done_elems) +{ + uint32_t const tid = blockIdx.x * blockDim.x + threadIdx.x; + uint32_t const stride = gridDim.x * blockDim.x; + for (uint32_t i = tid; i < y_elems; i += stride) + { + y_acc[i] = 0.0f; + } + for (uint32_t i = tid; i < r_elems; i += stride) + { + r_acc[i] = 0.0f; + } + if (done_counter != nullptr) + { + for (uint32_t i = tid; i < done_elems; i += stride) + { + done_counter[i] = 0; + } + } +} + +inline void fhcZeroWorkspaces(float* y_acc, uint32_t y_elems, float* r_acc, uint32_t r_elems, int* done_counter, + uint32_t done_elems, cudaStream_t stream) +{ + uint32_t const total = y_elems + r_elems + (done_counter != nullptr ? done_elems : 0u); + if (total == 0u) + { + return; + } + constexpr uint32_t kBlock = 256; + // Cap grid so we don't over-launch for small workspaces; 148 SMs on B200. + uint32_t const num_blocks = min(static_cast((total + kBlock - 1) / kBlock), 148u * 8u); + fhcZeroWorkspacesKernel<<>>( + y_acc, y_elems, r_acc, r_elems, done_counter, done_elems); +} + +} // namespace + +// ---- mHC fused kernel shape constants (mirrors the Python module) ---- +// HC_MULT * (2 + HC_MULT) = 4 * 6 = 24. +static constexpr uint32_t FHC_SHAPE_N = 24; +static constexpr uint32_t FHC_HC_MULT = 4; +static constexpr uint32_t FHC_BLOCK_K = 64; + +#ifdef TRTLLM_MHC_ENABLE_FUSED_HC +static constexpr uint32_t FHC_HIDDEN_FLASH = 4096; +static constexpr uint32_t FHC_HIDDEN_PRO = 7168; +static constexpr uint32_t FHC_BLOCK_M = 64; +static constexpr uint32_t FHC_BLOCK_N = 32; +static constexpr uint32_t FHC_SWIZZLE_CD = 128; +// Rebalanced from N_B=12 / N_INPUT=2: SASS PC-sampling on M=4096 KS=2 showed +// ~25% of all stalls landed on a single NANOSLEEP.SYNCS (mbarrier wait) — the +// pmap warp was input-buffer-starved with only 2 TMA stages for 56 h_tiles. +// Trade 5 B-stages (-40 KiB) for +1 input stage (+40 KiB), keeping total SMEM +// at 226 KiB. N_B=7 still leaves >= HC_MULT slack for the MMA pipeline. +static constexpr uint32_t FHC_N_B_STAGES = 7; +static constexpr uint32_t FHC_N_INPUT_STG = 3; +static constexpr uint32_t FHC_NUM_MMA_TH = 128; +static constexpr uint32_t FHC_NUM_PMAP_TH = 128; + +template