diff --git a/3rdparty/fetch_content.json b/3rdparty/fetch_content.json index 6a4679db5262..6c3356cfea27 100644 --- a/3rdparty/fetch_content.json +++ b/3rdparty/fetch_content.json @@ -32,7 +32,7 @@ { "name": "deepgemm", "git_repository": "https://github.com/deepseek-ai/DeepGEMM", - "git_tag": "c491439ed5966833d56883ca302b6f72e74f8105", + "git_tag": "67fc64863d43521080bf2005e6528d0fceee9510", "git_submodules_recurse": true, "source_subdir": "dont-add-this-project-with-add-subdirectory" }, diff --git a/cpp/include/tensorrt_llm/batch_manager/llmRequest.h b/cpp/include/tensorrt_llm/batch_manager/llmRequest.h index 644813e381b7..32558cdc78b3 100644 --- a/cpp/include/tensorrt_llm/batch_manager/llmRequest.h +++ b/cpp/include/tensorrt_llm/batch_manager/llmRequest.h @@ -1908,6 +1908,15 @@ class GenericLlmRequest return mPerfMetrics.kvCacheMetrics.numNewAllocatedBlocks; } + void updateKvCachePerfMetrics( + SizeType32 allocTotalBlocks, SizeType32 allocNewBlocks, SizeType32 reusedBlocks, SizeType32 missedBlocks) + { + updateAllocTotalBlocksPerRequest(allocTotalBlocks); + updateAllocNewBlocksPerRequest(allocNewBlocks); + updateReusedBlocksPerRequest(reusedBlocks); + updateMissedBlocksPerRequest(missedBlocks); + } + void updateReusedBlocksPerRequest(SizeType32 reusedBlocksPerRequest) { mPerfMetrics.kvCacheMetrics.numReusedBlocks += reusedBlocksPerRequest; diff --git a/cpp/tensorrt_llm/CMakeLists.txt b/cpp/tensorrt_llm/CMakeLists.txt index 06aefb3886ba..91beff7cf5bd 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 + mhcKernels_src + compressorKernels_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/IndexerTopK.h b/cpp/tensorrt_llm/kernels/IndexerTopK.h index 324e597dc779..ac9164771f6c 100644 --- a/cpp/tensorrt_llm/kernels/IndexerTopK.h +++ b/cpp/tensorrt_llm/kernels/IndexerTopK.h @@ -27,47 +27,63 @@ TRTLLM_NAMESPACE_BEGIN namespace kernels { -/// Indexer TopK decode. Three tiers: -/// - GVR Heuristic (preIdx provided, K in {512,1024,2048}, numColumns in -/// [kSeqSmall, splitWorkThreshold), numRows below the -/// architecture-derived wave/L2 bound). -/// - Single-block (numColumns < split-work threshold) -/// - Multi-pass radix (numColumns >= split-work threshold; requires -/// `scratch` sized via indexerTopKDecodeScratchBytes, -/// zero-init on first call and may be reused). -/// -/// `is_prefill = true` forces single-block (split-work suppressed). -void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indices, int const splitWorkThreshold, - int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, - int const topK = 2048, int const* preIdx = nullptr, int const preIdxStride = 0, int const preIdxCount = 0, - float* heuristicScratch = nullptr, cudaStream_t const stream = 0, void* scratch = nullptr, size_t scratchBytes = 0, - bool is_prefill = false); +// Number of blocks-per-row used by the multi-block split + merge dispatch path of +// invokeIndexerTopKDecode. Returns 1 when the single-block path is preferred. +// Callers that allocate aux buffers must use this same helper to size them, and +// must pass the same splitWorkThreshold they will pass to invokeIndexerTopKDecode +// (a value <= 0 selects the internal default). +int computeIndexerTopKDecodeBlocksPerRow(int numRows, int numColumns, int splitWorkThreshold = 0); -/// Size of the multi-pass radix `scratch` buffer for these shapes. -size_t indexerTopKDecodeScratchBytes(int numRows, int numColumns, int topK); +/// fp32 indexer TopK decode — L2-aware BS-threshold dispatcher with four +/// fallback tiers: +/// - GVR Heuristic (preIdx provided, kSeqSmall ≤ N < splitWork, BS < kBsLarge, K ∈ {512,1024,2048}) +/// - Insertion sort (N < kSortingAlgorithmThreshold) +/// - Radix sort (kSortingAlgorithmThreshold ≤ N < splitWork) +/// - Radix split-work (N ≥ splitWork — uses outLogitsAux / outIndicesAux) +void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indices, float* outLogitsAux, + int* outIndicesAux, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, + int const stride1, int const next_n, int const topK = 2048, int const* preIdx = nullptr, int const preIdxStride = 0, + int const preIdxCount = 0, float* heuristicScratch = nullptr, int const compressRatio = 1, + cudaStream_t const stream = 0); -/// bf16 overload; same contract. +/// bf16 indexer TopK decode — same dispatch axes as the fp32 entry, except +/// kBsL2 uses sizeof(__nv_bfloat16) bytes/elem (L2 footprint is half) and +/// the split-work tier is unsupported (the bf16/fp16 entry does not expose +/// the float aux buffers required for split-work). Insertion + radix tiers +/// share topKPerRowDecode with fp32 — histogram and sort run on float keys +/// after a static_cast(InputT) at HBM-read sites. +/// +/// Aborts with TLLM_CHECK if numColumns ≥ splitWorkThreshold; callers in +/// that regime must use the fp32 entry. void invokeIndexerTopKDecode(__nv_bfloat16 const* logits, int const* seqLens, int* indices, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, int const topK = 2048, int const* preIdx = nullptr, int const preIdxStride = 0, - int const preIdxCount = 0, __nv_bfloat16* heuristicScratch = nullptr, cudaStream_t const stream = 0, - void* scratch = nullptr, size_t scratchBytes = 0, bool is_prefill = false); + int const preIdxCount = 0, __nv_bfloat16* heuristicScratch = nullptr, int const compressRatio = 1, + cudaStream_t const stream = 0); -/// fp16 overload; same contract. +/// fp16 indexer TopK decode — see bf16 overload for dispatcher contract. void invokeIndexerTopKDecode(__half const* logits, int const* seqLens, int* indices, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, int const topK = 2048, int const* preIdx = nullptr, int const preIdxStride = 0, int const preIdxCount = 0, - __half* heuristicScratch = nullptr, cudaStream_t const stream = 0, void* scratch = nullptr, size_t scratchBytes = 0, - bool is_prefill = false); + __half* heuristicScratch = nullptr, int const compressRatio = 1, cudaStream_t const stream = 0); void invokeIndexerTopKPrefill(float const* logits, int const* rowStarts, int const* rowEnds, int* indices, int const numRows, int const numColumns, int const stride0, int const stride1, int const topK = 2048, cudaStream_t const stream = 0); -/// True iff invokeIndexerTopKDecode would pick the GVR tier for this shape: -/// K in {512,1024,2048}, numColumns in [kSeqSmall, splitWorkThreshold), and -/// numRows below the architecture-derived wave/L2 bound. Lets callers -/// provision preIdx / heuristicScratch only when needed. +/// Returns true iff invokeIndexerTopKDecode would route to the GVR Heuristic +/// kernel for this (numRows, numColumns, topK) triple, assuming valid preIdx +/// is provided and stride1 == 1. Useful for callers that need to provision a +/// preIdx tensor or heuristicScratch buffer only when GVR will be selected. +/// +/// Mirrors the gating logic of the dispatcher: K ∈ {512, 1024, 2048}, +/// numColumns ∈ [kSeqSmall, splitWorkThreshold), numRows < kBsLarge, where +/// kBsLarge = min(kBsWave, kBsL2) and kBsL2 scales with bytesPerElem. +/// +/// @param numRows logits rows (batch · next_n) +/// @param numColumns logits columns (max sequence length) +/// @param topK requested output size +/// @param bytesPerElem element size of logits (4 for fp32, 2 for bf16/fp16) bool canIndexerTopKDecodeUseGvr(int numRows, int numColumns, int topK, int bytesPerElem = 4); } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/compressorKernels/CMakeLists.txt b/cpp/tensorrt_llm/kernels/compressorKernels/CMakeLists.txt new file mode 100644 index 000000000000..160dbced4b19 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/compressorKernels/CMakeLists.txt @@ -0,0 +1,25 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 1993-2024 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..14fd7d19b38b --- /dev/null +++ b/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu @@ -0,0 +1,1763 @@ +/* + * Copyright (c) 2022-2025, 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..ca8424f64fc8 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.h @@ -0,0 +1,99 @@ +/* + * Copyright (c) 2022-2025, 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/customMoeRoutingKernels.cu b/cpp/tensorrt_llm/kernels/customMoeRoutingKernels.cu index 58ccb4ac8ea8..f76db261e5a3 100644 --- a/cpp/tensorrt_llm/kernels/customMoeRoutingKernels.cu +++ b/cpp/tensorrt_llm/kernels/customMoeRoutingKernels.cu @@ -15,6 +15,7 @@ */ #include "moeTopKFuncs.cuh" +#include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/config.h" #include "tensorrt_llm/common/cudaTypeUtils.cuh" #include "tensorrt_llm/common/envUtils.h" @@ -309,6 +310,176 @@ INSTANTIATE_RENORM_MOE_ROUTING(half, __nv_bfloat16, int32_t, true); INSTANTIATE_RENORM_MOE_ROUTING(__nv_bfloat16, __nv_bfloat16, int32_t, true); #endif +static constexpr int kTOPK = 6; + +// CUDA kernel for gate forward +// Input: pre-computed scores from linear(x, weight) done outside kernel +// Template parameters: +// nExperts: number of experts +// topK: number of top experts to select +// hash: true for hash mode, false for topk mode +// One warp per row (batch element) +template +__global__ void gate_forward_kernel( + float const* __restrict__ scores_in, // [batch_size, nExperts] - pre-computed from linear(x, weight) + float const* __restrict__ bias, // [nExperts] (only used when hash=false) + int const* __restrict__ input_ids, // [batch_size] (only used when hash=true) + int const* __restrict__ tid2eid, // [vocab_size, topK] (only used when hash=true) + float* __restrict__ out_weights, // [batch_size, topK] + int* __restrict__ out_indices, // [batch_size, topK] + int batch_size, float route_scale) +{ + // Compile-time constants + constexpr int kExpertsPerThread = nExperts / WARP_SIZE; + constexpr int kWarpsPerBlock = 4; // Adjust based on occupancy needs + + // Shared memory for original scores (one array per warp in the block) + __shared__ float smem_scores[kWarpsPerBlock][nExperts]; + + // One warp per batch element + int const global_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) / WARP_SIZE; + int const local_warp_id = (threadIdx.x / WARP_SIZE) % kWarpsPerBlock; + int const lane_id = threadIdx.x % WARP_SIZE; + + if (global_warp_id >= batch_size) + return; + + auto warp = cg::tiled_partition(cg::this_thread_block()); + + // Pointer to this warp's shared memory and input scores + float* my_smem = smem_scores[local_warp_id]; + float const* scores_row = scores_in + global_warp_id * nExperts; + +// Load scores, apply score function (softplus + sqrt), and store to shared memory +#pragma unroll + for (int e = 0; e < kExpertsPerThread; ++e) + { + int expert_id = lane_id + e * WARP_SIZE; + float s = scores_row[expert_id]; + float sp = log1pf(expf(s)); + float score = sqrtf(sp); + my_smem[expert_id] = score; // Store original score to shared memory + } + __syncwarp(); // Ensure all scores are written before reading + + // Output: each of first K lanes holds one value + float my_topk_value = 0.0f; + int my_topk_index = 0; + + if constexpr (hash) + { + // Hash mode: directly read from shared memory + int token_id = input_ids[global_warp_id]; + int const* expert_ids = tid2eid + token_id * topK; + + if (lane_id < topK) + { + int expert_id = expert_ids[lane_id]; + my_topk_index = expert_id; + my_topk_value = my_smem[expert_id]; // Direct lookup from shared memory + } + } + else + { + // Topk mode: load from shared memory, add bias to registers for topk + float scores[kExpertsPerThread]; + int indices[kExpertsPerThread]; + +#pragma unroll + for (int e = 0; e < kExpertsPerThread; ++e) + { + int expert_id = lane_id + e * WARP_SIZE; + indices[e] = expert_id; + scores[e] = my_smem[expert_id] + bias[expert_id]; // Add bias for topk selection + } + + // Use reduceTopK to find top-k experts + float topk_values[topK]; + int32_t topk_indices[topK]; + constexpr float minValue = -1e30f; + reduce_topk::reduceTopK( + warp, topk_values, topk_indices, scores, indices, minValue, topK); + + // Gather original weights (without bias) from shared memory + if (lane_id < topK) + { + int expert_id = topk_indices[lane_id]; + my_topk_index = expert_id; + my_topk_value = my_smem[expert_id]; // Read original score (no bias) + } + } + + // Reduce to get sum (first K lanes have values, others have 0) + float weight_sum = cg::reduce(warp, my_topk_value, cg::plus{}); + + // Normalize weights and write output (first K lanes) + if (lane_id < topK) + { + out_weights[global_warp_id * topK + lane_id] = (my_topk_value / weight_sum) * route_scale; + out_indices[global_warp_id * topK + lane_id] = my_topk_index; + } +} + +// C++ wrapper function (output tensors passed as parameters) +// All tensors are float32 +template +void launch_gate_forward_kernel(float* scores_in, float* bias, int* input_ids, int* tid2eid, float* out_weights, + int* out_indices, int batch_size, float route_scale, cudaStream_t stream) +{ + constexpr int warps_per_block = 4; + constexpr int threads_per_block = warps_per_block * WARP_SIZE; + int const blocks = (batch_size + warps_per_block - 1) / warps_per_block; + + gate_forward_kernel<<>>( + scores_in, bias, input_ids, tid2eid, out_weights, out_indices, batch_size, route_scale); +} + +void gate_forward(void* scores_in, // [batch_size, nExperts] - pre-computed from linear(x, weight) + void* bias, // nullptr if hash mode + void* input_ids, // nullptr if non-hash mode + void* tid2eid, // nullptr if non-hash mode + void* out_weights, // [batch_size, topK] - pre-allocated + void* out_indices, // [batch_size, topK] - pre-allocated + int batch_size, int n_experts, float route_scale, bool is_hash, cudaStream_t stream) +{ + auto* scores = static_cast(scores_in); + auto* bias_ptr = static_cast(bias); + auto* input_ids_ptr = static_cast(input_ids); + auto* tid2eid_ptr = static_cast(tid2eid); + auto* weights = static_cast(out_weights); + auto* indices = static_cast(out_indices); + + switch (n_experts) + { + case 256: + if (is_hash) + { + launch_gate_forward_kernel<256, true>( + scores, nullptr, input_ids_ptr, tid2eid_ptr, weights, indices, batch_size, route_scale, stream); + } + else + { + launch_gate_forward_kernel<256, false>( + scores, bias_ptr, nullptr, nullptr, weights, indices, batch_size, route_scale, stream); + } + break; + case 384: + if (is_hash) + { + launch_gate_forward_kernel<384, true>( + scores, nullptr, input_ids_ptr, tid2eid_ptr, weights, indices, batch_size, route_scale, stream); + } + else + { + launch_gate_forward_kernel<384, false>( + scores, bias_ptr, nullptr, nullptr, weights, indices, batch_size, route_scale, stream); + } + break; + default: TLLM_CHECK_WITH_INFO(false, "gate_forward only supports n_experts 256 or 384"); + } + sync_check_cuda_error(stream); +} + } // namespace kernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/customMoeRoutingKernels.h b/cpp/tensorrt_llm/kernels/customMoeRoutingKernels.h index f8240b436393..367676fcc334 100644 --- a/cpp/tensorrt_llm/kernels/customMoeRoutingKernels.h +++ b/cpp/tensorrt_llm/kernels/customMoeRoutingKernels.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2025-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. @@ -29,6 +29,16 @@ namespace kernels template void invokeCustomMoeRouting(InputT* routerLogits, OutputT* topkValues, IdxT* topkIndices, int64_t const numTokens, int64_t const numExperts, int64_t const topK, cudaStream_t const stream); + +// Gate forward function for custom MoE routing +// All tensors are expected to be float32 for scores/weights, int32 for indices +void gate_forward(void* scores_in, // [batch_size, nExperts] - pre-computed from linear(x, weight) + void* bias, // nullptr if hash mode + void* input_ids, // nullptr if non-hash mode + void* tid2eid, // nullptr if non-hash mode + void* out_weights, // [batch_size, topK] - pre-allocated + void* out_indices, // [batch_size, topK] - pre-allocated + int batch_size, int n_experts, float route_scale, bool is_hash, cudaStream_t stream); } // namespace kernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp b/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp index b0ea6333a88c..86407533c21d 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp @@ -380,10 +380,15 @@ std::vector get_candidate_configs_sm100_dynamic_cluster_shape std::vector candidate_configs; if ((config & CutlassGemmConfig::FP4_ONLY) != 0) { - if (sm == 100) + if (sm == 100 || sm == 103) { if (schedule != EpilogueScheduleType::TMA) + { return {}; + } + } + if (sm == 100) + { candidate_configs.push_back(CutlassGemmConfig{CutlassTileConfigSM100::CtaShape128x64x128B, MainloopScheduleType::AUTO, schedule, cluster1sm, dynamic_cluster_shape, fallback_cluster_shape, sm}); if (supports_2sm) diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.cu b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.cu new file mode 100644 index 000000000000..b70d517feea5 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.cu @@ -0,0 +1,210 @@ +/* + * 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 1x128 FP8 quantize + UE8M0 scale packing. +// +// Replaces the (scale_1x128_kernel + pack_fp32_into_ue8m0) two-kernel sequence +// used by SM100 deep_gemm fp8 block-scale GEMMs. Adapted from the SM120 MoE +// in-kernel packing pattern (`scale_1x128_kernel_sm120` in +// sm120_blockwise_gemm/sm120_fp8_moe_gemm_1d1d.cuh), specialised for the +// non-MoE case (single contiguous batch, no token offsets). + +#include "fp8_blockscale_quant_packed.h" + +#include "tensorrt_llm/common/config.h" +#include "tensorrt_llm/common/envUtils.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::fp8_blockscale_gemm +{ + +namespace +{ + +__device__ __forceinline__ float reciprocal_approximate_ftz_local(float a) +{ + float b; + asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a)); + return b; +} + +// Each warp consumes one row × 4 quantization blocks (4 × 128 = 512 K elems). +// 32 lanes split into 4 lane-groups of 8: each group covers 1 quant block +// (8 lanes × 16 BF16 elems = 128 elems). After per-block amax, lanes +// 0/8/16/24 each hold one UE8M0 scale byte; lane 0 packs them into a uint32 +// and stores in the deep_gemm-expected MN-major layout. +template +__global__ void fp8_quantize_1x128_packed_kernel_impl(__nv_fp8_e4m3* __restrict__ fp8_output, + int32_t* __restrict__ packed_scale_output, __nv_bfloat16 const* __restrict__ input, int const m, int const k, + int const scale_leading_dim_uint32) +{ + int const packed_sf_k_idx = static_cast(blockIdx.x); + int const warp_id = static_cast(threadIdx.x) >> 5; + int const lane_id = static_cast(threadIdx.x) & 31; + int const m_idx = static_cast(blockIdx.y) * WarpsPerBlock + warp_id; + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaGridDependencySynchronize(); +#endif + + if (m_idx >= m) + { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaTriggerProgrammaticLaunchCompletion(); +#endif + return; + } + + int const k_base = packed_sf_k_idx * 512 + lane_id * 16; + + // ---- 1. Load 16 BF16 elements per lane. ---- + auto const* in_ptr = reinterpret_cast(input + static_cast(m_idx) * k + k_base); + constexpr int kLoadNumElems = sizeof(double4) / sizeof(__nv_bfloat16); // 16 + + union LoadTrick + { + double4 pack; + __nv_bfloat16 v[kLoadNumElems]; + }; + + LoadTrick load_trick; + bool const k_in_range = (k_base < k); + load_trick.pack = k_in_range ? in_ptr[0] : double4{}; + + if (k_in_range && k_base + kLoadNumElems > k) + { + int const valid = k - k_base; +#pragma unroll + for (int i = 0; i < kLoadNumElems; ++i) + { + if (i >= valid) + { + load_trick.v[i] = __nv_bfloat16(0.0f); + } + } + } + + // ---- 2. Per-block amax (lanes 0..7 / 8..15 / 16..23 / 24..31 = 4 quant blocks). ---- + __nv_bfloat16 max_elem = __nv_bfloat16(0.0f); +#pragma unroll + for (int i = 0; i < kLoadNumElems; ++i) + { + max_elem = __hmax(max_elem, __habs(load_trick.v[i])); + } + float amax = static_cast(max_elem); + amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFFu, amax, 4, 8)); + amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFFu, amax, 2, 8)); + amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFFu, amax, 1, 8)); + amax = fmaxf(amax, 1e-10f); + + // ---- 3. UE8M0 dequant scale. ---- + float const dequant_scale_raw = amax * reciprocal_approximate_ftz_local(448.0f); + __nv_fp8_e8m0 ue8m0_scale; + ue8m0_scale.__x = __nv_cvt_float_to_e8m0(dequant_scale_raw, __NV_SATFINITE, cudaRoundPosInf); + + // Recover quant_scale = 1 / 2^(exp - 127) for fp8 conversion. + constexpr uint32_t FP32_EXPONENT_BIAS = 127u; + float const quant_scale = (ue8m0_scale.__x == 0) + ? 1.0f + : exp2f(static_cast(FP32_EXPONENT_BIAS) - static_cast(ue8m0_scale.__x)); + + // ---- 4. Quantize and store FP8 output. ---- + constexpr int kStoreNumElems = sizeof(float4) / sizeof(__nv_fp8_e4m3); // 16 + + union StoreTrick + { + float4 pack; + __nv_fp8_e4m3 v[kStoreNumElems]; + }; + + StoreTrick store_trick; + store_trick.pack = float4{}; +#pragma unroll + for (int i = 0; i < kStoreNumElems; ++i) + { + store_trick.v[i] = __nv_fp8_e4m3(static_cast(load_trick.v[i]) * quant_scale); + } + auto* out_ptr = reinterpret_cast(fp8_output + static_cast(m_idx) * k + k_base); + if (k_in_range) + { + if (k_base + kStoreNumElems > k) + { + int const valid = k - k_base; +#pragma unroll + for (int i = 0; i < kStoreNumElems; ++i) + { + if (i >= valid) + { + store_trick.v[i] = __nv_fp8_e4m3(0.0f); + } + } + } + out_ptr[0] = store_trick.pack; + } + + // ---- 5. Pack 4 UE8M0 scales (lanes 0/8/16/24) and store. ---- + uint32_t const s0 = __shfl_sync(0xFFFFFFFFu, static_cast(ue8m0_scale.__x), 0); + uint32_t const s1 = __shfl_sync(0xFFFFFFFFu, static_cast(ue8m0_scale.__x), 8); + uint32_t const s2 = __shfl_sync(0xFFFFFFFFu, static_cast(ue8m0_scale.__x), 16); + uint32_t const s3 = __shfl_sync(0xFFFFFFFFu, static_cast(ue8m0_scale.__x), 24); + if (lane_id == 0) + { + // Mask off scale bytes whose sf_k is past the actual K. + int const num_sf_k = (k + 127) / 128; + int const sf_k_base = packed_sf_k_idx * 4; + uint32_t packed = 0u; + if (sf_k_base + 0 < num_sf_k) + packed |= s0; + if (sf_k_base + 1 < num_sf_k) + packed |= (s1 << 8); + if (sf_k_base + 2 < num_sf_k) + packed |= (s2 << 16); + if (sf_k_base + 3 < num_sf_k) + packed |= (s3 << 24); + // Layout: packed_scale[packed_sf_k_idx, m_idx] + packed_scale_output[static_cast(packed_sf_k_idx) * scale_leading_dim_uint32 + m_idx] = packed; + } + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + +} // namespace + +void launch_fp8_quantize_1x128_packed_bf16_e4m3(__nv_fp8_e4m3* fp8_output, int32_t* packed_scale_output, + __nv_bfloat16 const* input, int m, int k, int scale_leading_dim_uint32, cudaStream_t stream) +{ + constexpr int kWarpsPerBlock = 4; + int const num_packed_sf_k = (((k + 127) / 128) + 3) / 4; + int const m_blocks = (m + kWarpsPerBlock - 1) / kWarpsPerBlock; + dim3 const grid(num_packed_sf_k, m_blocks, 1); + dim3 const block(kWarpsPerBlock * 32, 1, 1); + + tensorrt_llm::common::launchWithPdlWhenEnabled("fp8_quantize_1x128_packed_kernel_impl", + fp8_quantize_1x128_packed_kernel_impl, grid, block, 0, stream, fp8_output, packed_scale_output, + input, m, k, scale_leading_dim_uint32); +} + +} // namespace kernels::fp8_blockscale_gemm + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.h new file mode 100644 index 000000000000..a4079cd9b054 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.h @@ -0,0 +1,52 @@ +/* + * 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. + */ + +// Host-callable launcher for the fused FP8 1x128 quantize + UE8M0-pack kernel. +// Kernel implementation lives in fp8_blockscale_quant_packed.cu and is built +// by nvcc; this header is safe to include from .cpp files compiled by g++. + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::fp8_blockscale_gemm +{ + +// Launches the fused 1x128 FP8 quant + UE8M0 pack kernel. +// +// Inputs: +// input : BF16 [m, k] row-major contiguous +// Outputs: +// fp8_output : E4M3 [m, k] row-major contiguous +// packed_scale_output : uint32 [packed_sf_k, scale_leading_dim_uint32] +// where packed_sf_k = ceil(ceil(k/128)/4) +// +// `scale_leading_dim_uint32` is the (uint32) stride between consecutive +// packed_sf_k rows of the scale tensor; caller is responsible for choosing +// it (typically aligned to 4 uint32 = 16 bytes for TMA alignment). +void launch_fp8_quantize_1x128_packed_bf16_e4m3(__nv_fp8_e4m3* fp8_output, int32_t* packed_scale_output, + __nv_bfloat16 const* input, int m, int k, int scale_leading_dim_uint32, cudaStream_t stream); + +} // namespace kernels::fp8_blockscale_gemm + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.cu b/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.cu new file mode 100644 index 000000000000..3dbbb7a6ac33 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.cu @@ -0,0 +1,157 @@ +/* + * 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. + */ + +#include "tensorrt_llm/kernels/deepseekV4QNormKernel.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace +{ + +constexpr int kWarpSize = 32; +constexpr int kWarpsPerBlock = 4; +constexpr int kThreadsPerBlock = kWarpSize * kWarpsPerBlock; + +template +struct Vec2Traits; + +template <> +struct Vec2Traits +{ + using Type = half2; + + __device__ static float2 toFloat2(Type value) + { + return __half22float2(value); + } + + __device__ static Type fromFloat2(float2 value) + { + return __floats2half2_rn(value.x, value.y); + } +}; + +template <> +struct Vec2Traits<__nv_bfloat16> +{ + using Type = __nv_bfloat162; + + __device__ static float2 toFloat2(Type value) + { + return __bfloat1622float2(value); + } + + __device__ static Type fromFloat2(float2 value) + { + return __floats2bfloat162_rn(value.x, value.y); + } +}; + +__device__ __forceinline__ float warpReduceSum(float value) +{ + for (int mask = kWarpSize / 2; mask > 0; mask >>= 1) + { + value += __shfl_xor_sync(0xFFFFFFFF, value, mask); + } + return value; +} + +template +__global__ void deepseekV4QNormKernel(T const* input, T* output, int totalRows, float eps) +{ + static_assert(kHeadDim % (2 * kWarpSize) == 0); + constexpr int kPairsPerRow = kHeadDim / 2; + constexpr int kPairsPerLane = kPairsPerRow / kWarpSize; + + using Vec2 = typename Vec2Traits::Type; + + int const warpId = threadIdx.x / kWarpSize; + int const laneId = threadIdx.x % kWarpSize; + int const row = blockIdx.x * kWarpsPerBlock + warpId; + + if (row >= totalRows) + { + return; + } + + auto const* inputPair = reinterpret_cast(input + static_cast(row) * kHeadDim); + auto* outputPair = reinterpret_cast(output + static_cast(row) * kHeadDim); + + float2 values[kPairsPerLane]; + float sumSquares = 0.0F; + +#pragma unroll + for (int i = 0; i < kPairsPerLane; ++i) + { + int const pairIdx = i * kWarpSize + laneId; + values[i] = Vec2Traits::toFloat2(inputPair[pairIdx]); + sumSquares += values[i].x * values[i].x + values[i].y * values[i].y; + } + + sumSquares = warpReduceSum(sumSquares); + float const scale = rsqrtf(sumSquares / static_cast(kHeadDim) + eps); + +#pragma unroll + for (int i = 0; i < kPairsPerLane; ++i) + { + int const pairIdx = i * kWarpSize + laneId; + float2 value{values[i].x * scale, values[i].y * scale}; + outputPair[pairIdx] = Vec2Traits::fromFloat2(value); + } +} + +template +void dispatchDeepseekV4QNorm( + void const* input, void* output, int totalRows, int headDim, float eps, cudaStream_t stream) +{ + assert(headDim == 512); + + dim3 const block(kThreadsPerBlock); + dim3 const grid((totalRows + kWarpsPerBlock - 1) / kWarpsPerBlock); + deepseekV4QNormKernel + <<>>(static_cast(input), static_cast(output), totalRows, eps); +} + +} // namespace + +void invokeDeepseekV4QNorm( + void const* input, void* output, int totalRows, int headDim, bool isBfloat16, float eps, cudaStream_t stream) +{ + if (totalRows == 0) + { + return; + } + + if (isBfloat16) + { + dispatchDeepseekV4QNorm<__nv_bfloat16>(input, output, totalRows, headDim, eps, stream); + } + else + { + dispatchDeepseekV4QNorm(input, output, totalRows, headDim, eps, stream); + } +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.h b/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.h new file mode 100644 index 000000000000..13b9a623d0f9 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.h @@ -0,0 +1,33 @@ +/* + * 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 "tensorrt_llm/common/config.h" + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +void invokeDeepseekV4QNorm( + void const* input, void* output, int totalRows, int headDim, bool isBfloat16, float eps, cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu b/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu index e33f46e1e152..102866de363f 100644 --- a/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu +++ b/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu @@ -291,6 +291,55 @@ template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__n template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 16, 256, 6144>( float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); +// hidden_dim=4096 instantiations (DeepSeek-V4). +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 1, 256, 4096>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 2, 256, 4096>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 3, 256, 4096>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 4, 256, 4096>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 5, 256, 4096>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 6, 256, 4096>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 7, 256, 4096>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 8, 256, 4096>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 9, 256, 4096>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 10, 256, 4096>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 11, 256, 4096>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 12, 256, 4096>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 13, 256, 4096>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 14, 256, 4096>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 15, 256, 4096>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 16, 256, 4096>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + } // namespace kernels::dsv3MinLatencyKernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/heuristicTopKDecode.cu b/cpp/tensorrt_llm/kernels/heuristicTopKDecode.cu index 25b0de3ef0a5..839ae5483c05 100644 --- a/cpp/tensorrt_llm/kernels/heuristicTopKDecode.cu +++ b/cpp/tensorrt_llm/kernels/heuristicTopKDecode.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2019-2025, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2019-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. @@ -47,15 +47,19 @@ using heuristic_topk::KernelSmemTplK; // same kernel template. Smem layout is derived from GvrParams // at compile time. template -__global__ void __launch_bounds__(BLOCK_SIZE) heuristicTopKMultiRowKernel(float const* __restrict__ logits, - int const* __restrict__ seqLens, int const* __restrict__ preIdx, float* __restrict__ scratchValues, - int* __restrict__ outIndices, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount) +__global__ void __launch_bounds__(BLOCK_SIZE) + heuristicTopKMultiRowKernel(float const* __restrict__ logits, int const* __restrict__ seqLens, + int const* __restrict__ preIdx, float* __restrict__ scratchValues, int* __restrict__ outIndices, int stride0, + int next_n, int topK, int preIdxStride, int preIdxCount, int compressRatio) { using SmemT = KernelSmemTplK::kC, GvrParams::kNumBins>; int const rowIdx = blockIdx.x; int const seq_len = seqLens[rowIdx / next_n]; - int const N = seq_len - next_n + (rowIdx % next_n) + 1; + // seqLens is in uncompressed token space; the logits/preIdx live in + // compressed-index space when compressRatio > 1 (DSv4 indexer). + int const actual_kv_len = seq_len - next_n + (rowIdx % next_n) + 1; + int const N = actual_kv_len / compressRatio; float const* __restrict__ input = logits + static_cast(rowIdx) * stride0; int const* __restrict__ rowPreIdx = preIdx + static_cast(rowIdx / next_n) * preIdxStride; @@ -84,9 +88,19 @@ __global__ void __launch_bounds__(BLOCK_SIZE) heuristicTopKMultiRowKernel(float return; } - // +1 accounts for the temporal shift: prev_topk indices were computed at - // seq_len-1, but the current step has one additional KV token appended. - int const preIdxOffset = (rowIdx % next_n) + 1; + // Temporal-shift offset to map prev-step's top-K indices into this step's + // KV index space. + // compressRatio == 1 (DSv3.2): +1 — KV grew by exactly 1 token per + // decode step; prev indices were at seq_len-1 so a uniform +1 maps + // them to the equivalent positions under the indexer's "newest-first" + // layout. The (rowIdx % next_n) addend extends this to MTP windows. + // compressRatio == 4 (DSv4): 0 — in compressed-index space new + // compressed entries are appended at the end; prev indices in + // [0, c_prev-1] remain valid as-is. Per-row Δc varies (0 or 1) with + // prev kv_len mod 4 alignment, but a uniform offset of 0 stays + // within-bounds for all rows and preserves the temporal-correlation + // hint (vertical top-K consistency validated offline). + int const preIdxOffset = (compressRatio == 1) ? ((rowIdx % next_n) + 1) : 0; gvrTopKJob(input, N, rowPreIdx, preIdxCount, topK, outputValues, outputIndices, smem, preIdxOffset); #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaTriggerProgrammaticLaunchCompletion(); @@ -103,16 +117,18 @@ __global__ void __launch_bounds__(BLOCK_SIZE) heuristicTopKMultiRowKernel(float // Templated on (InputT, TopK). Smem layout is derived from // GvrParams. template -__global__ void __launch_bounds__(BLOCK_SIZE) heuristicTopKMultiRowKernelDtype(InputT const* __restrict__ logits, - int const* __restrict__ seqLens, int const* __restrict__ preIdx, InputT* __restrict__ scratchValues, - int* __restrict__ outIndices, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount) +__global__ void __launch_bounds__(BLOCK_SIZE) + heuristicTopKMultiRowKernelDtype(InputT const* __restrict__ logits, int const* __restrict__ seqLens, + int const* __restrict__ preIdx, InputT* __restrict__ scratchValues, int* __restrict__ outIndices, int stride0, + int next_n, int topK, int preIdxStride, int preIdxCount, int compressRatio) { // dtype path uses fp32 keys[] in smem (down-conversion deferred to writeback). using SmemT = KernelSmemTplK::kC, GvrParams::kNumBins>; int const rowIdx = blockIdx.x; int const seq_len = seqLens[rowIdx / next_n]; - int const N = seq_len - next_n + (rowIdx % next_n) + 1; + int const actual_kv_len = seq_len - next_n + (rowIdx % next_n) + 1; + int const N = actual_kv_len / compressRatio; InputT const* __restrict__ input = logits + static_cast(rowIdx) * stride0; int const* __restrict__ rowPreIdx = preIdx + static_cast(rowIdx / next_n) * preIdxStride; @@ -142,7 +158,8 @@ __global__ void __launch_bounds__(BLOCK_SIZE) heuristicTopKMultiRowKernelDtype(I return; } - int const preIdxOffset = (rowIdx % next_n) + 1; + // See fp32 path: cr==1 → (rowIdx % next_n)+1; cr!=1 (DSv4) → 0. + int const preIdxOffset = (compressRatio == 1) ? ((rowIdx % next_n) + 1) : 0; gvrTopKJobDtype( input, N, rowPreIdx, preIdxCount, topK, outputValues, outputIndices, smem, preIdxOffset); #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) @@ -152,24 +169,25 @@ __global__ void __launch_bounds__(BLOCK_SIZE) heuristicTopKMultiRowKernelDtype(I // Explicit instantiations — 6 (dtype × K) combos. Launchers dispatch on // runtime topK via switch, so all 6 must be available at link time. +// Trailing `int` is the compressRatio parameter (1 = V3.2, 4 = V4 indexer). template __global__ void heuristicTopKMultiRowKernelDtype<__nv_bfloat16, 512>( - __nv_bfloat16 const*, int const*, int const*, __nv_bfloat16*, int*, int, int, int, int, int); + __nv_bfloat16 const*, int const*, int const*, __nv_bfloat16*, int*, int, int, int, int, int, int); template __global__ void heuristicTopKMultiRowKernelDtype<__nv_bfloat16, 1024>( - __nv_bfloat16 const*, int const*, int const*, __nv_bfloat16*, int*, int, int, int, int, int); + __nv_bfloat16 const*, int const*, int const*, __nv_bfloat16*, int*, int, int, int, int, int, int); template __global__ void heuristicTopKMultiRowKernelDtype<__nv_bfloat16, 2048>( - __nv_bfloat16 const*, int const*, int const*, __nv_bfloat16*, int*, int, int, int, int, int); + __nv_bfloat16 const*, int const*, int const*, __nv_bfloat16*, int*, int, int, int, int, int, int); template __global__ void heuristicTopKMultiRowKernelDtype<__half, 512>( - __half const*, int const*, int const*, __half*, int*, int, int, int, int, int); + __half const*, int const*, int const*, __half*, int*, int, int, int, int, int, int); template __global__ void heuristicTopKMultiRowKernelDtype<__half, 1024>( - __half const*, int const*, int const*, __half*, int*, int, int, int, int, int); + __half const*, int const*, int const*, __half*, int*, int, int, int, int, int, int); template __global__ void heuristicTopKMultiRowKernelDtype<__half, 2048>( - __half const*, int const*, int const*, __half*, int*, int, int, int, int, int); + __half const*, int const*, int const*, __half*, int*, int, int, int, int, int, int); template __global__ void heuristicTopKMultiRowKernel<512>( - float const*, int const*, int const*, float*, int*, int, int, int, int, int); + float const*, int const*, int const*, float*, int*, int, int, int, int, int, int); template __global__ void heuristicTopKMultiRowKernel<1024>( - float const*, int const*, int const*, float*, int*, int, int, int, int, int); + float const*, int const*, int const*, float*, int*, int, int, int, int, int, int); template __global__ void heuristicTopKMultiRowKernel<2048>( - float const*, int const*, int const*, float*, int*, int, int, int, int, int); + float const*, int const*, int const*, float*, int*, int, int, int, int, int, int); // Dispatch on topK at runtime — each TopK-instantiation gets its own smem // size (driven by GvrParams::kC/kNumBins) and own kfn pointer @@ -184,7 +202,7 @@ template __global__ void heuristicTopKMultiRowKernel<2048>( template void launchHeuristicTopKDecodeImpl(InputT const* logits, int const* seqLens, int const* preIdx, int* outIndices, InputT* scratchValues, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount, int numRows, - cudaStream_t stream) + int compressRatio, cudaStream_t stream) { TLLM_CHECK_WITH_INFO( topK == 512 || topK == 1024 || topK == 2048, "heuristicTopKDecode requires topK ∈ {512, 1024, 2048}"); @@ -224,7 +242,7 @@ void launchHeuristicTopKDecodeImpl(InputT const* logits, int const* seqLens, int config.attrs = attrs; cudaLaunchKernelEx(&config, kfn, logits, seqLens, preIdx, scratchValues, outIndices, stride0, next_n, topK, - preIdxStride, preIdxCount); + preIdxStride, preIdxCount, compressRatio); }; switch (topK) @@ -240,26 +258,26 @@ void launchHeuristicTopKDecodeImpl(InputT const* logits, int const* seqLens, int void launchHeuristicTopKDecode(float const* logits, int const* seqLens, int const* preIdx, int* outIndices, float* scratchValues, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount, int numRows, - cudaStream_t stream) + int compressRatio, cudaStream_t stream) { launchHeuristicTopKDecodeImpl(logits, seqLens, preIdx, outIndices, scratchValues, stride0, next_n, topK, - preIdxStride, preIdxCount, numRows, stream); + preIdxStride, preIdxCount, numRows, compressRatio, stream); } void launchHeuristicTopKDecode(__nv_bfloat16 const* logits, int const* seqLens, int const* preIdx, int* outIndices, __nv_bfloat16* scratchValues, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount, int numRows, - cudaStream_t stream) + int compressRatio, cudaStream_t stream) { launchHeuristicTopKDecodeImpl<__nv_bfloat16>(logits, seqLens, preIdx, outIndices, scratchValues, stride0, next_n, - topK, preIdxStride, preIdxCount, numRows, stream); + topK, preIdxStride, preIdxCount, numRows, compressRatio, stream); } void launchHeuristicTopKDecode(__half const* logits, int const* seqLens, int const* preIdx, int* outIndices, __half* scratchValues, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount, int numRows, - cudaStream_t stream) + int compressRatio, cudaStream_t stream) { launchHeuristicTopKDecodeImpl<__half>(logits, seqLens, preIdx, outIndices, scratchValues, stride0, next_n, topK, - preIdxStride, preIdxCount, numRows, stream); + preIdxStride, preIdxCount, numRows, compressRatio, stream); } } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/heuristicTopKDecode.h b/cpp/tensorrt_llm/kernels/heuristicTopKDecode.h index 5e38c64cba06..0d2330f76545 100644 --- a/cpp/tensorrt_llm/kernels/heuristicTopKDecode.h +++ b/cpp/tensorrt_llm/kernels/heuristicTopKDecode.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2019-2025, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2019-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. @@ -33,21 +33,28 @@ inline constexpr int kHeuristicSize = 2048; /// Launch heuristic TopK decode kernel — fp32 input. /// @param scratchValues Caller-owned buffer of size [numRows * topK] floats. /// Required for CUDA Graph compatibility — must have a stable device address. +/// @param compressRatio KV compression ratio (1 = V3.2 indexer; 4 = V4 indexer +/// whose logits/preIdx live in compressed-token-index space). For +/// compressRatio != 1, preIdxOffset is forced to 0 (append-at-end in +/// compressed space → prev-step indices remain valid as-is); the +/// existing (rowIdx % next_n)+1 shift is used only when compressRatio==1. void launchHeuristicTopKDecode(float const* logits, int const* seqLens, int const* preIdx, int* outIndices, float* scratchValues, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount, int numRows, - cudaStream_t stream); + int compressRatio, cudaStream_t stream); /// Launch heuristic TopK decode kernel — bf16 input. /// scratchValues is [numRows * topK] of bf16 (matches input dtype). +/// @param compressRatio See fp32 overload. void launchHeuristicTopKDecode(__nv_bfloat16 const* logits, int const* seqLens, int const* preIdx, int* outIndices, __nv_bfloat16* scratchValues, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount, int numRows, - cudaStream_t stream); + int compressRatio, cudaStream_t stream); /// Launch heuristic TopK decode kernel — fp16 input. /// scratchValues is [numRows * topK] of fp16 (matches input dtype). +/// @param compressRatio See fp32 overload. void launchHeuristicTopKDecode(__half const* logits, int const* seqLens, int const* preIdx, int* outIndices, __half* scratchValues, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount, int numRows, - cudaStream_t stream); + int compressRatio, cudaStream_t stream); } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/indexerTopK.cu b/cpp/tensorrt_llm/kernels/indexerTopK.cu index ad1e728b3298..b35531a480a2 100644 --- a/cpp/tensorrt_llm/kernels/indexerTopK.cu +++ b/cpp/tensorrt_llm/kernels/indexerTopK.cu @@ -39,6 +39,13 @@ TRTLLM_NAMESPACE_BEGIN namespace kernels { +// Radix histogram bin count, used by topKPerRowJob's per-step distribution +// pass. 2048 = 2^11 maps directly to the fp16 fast path's 11-bit bin index +// (sign + exponent + top mantissa bits, mask 0x7FF). All file-scope dispatch +// thresholds below derive from this so that a future change to the histogram +// width does not require re-tuning the heuristics. +constexpr int kNumBins = 2048; + namespace { @@ -169,6 +176,20 @@ __device__ bool processHistogramStep(int const* indices, InputT const* logits, i int& thresholdBinIdx, SmemOutputType& smemOutput, int* smemThresholdBinIdx, int* smemFinalDstIdx, int* smemFinalBinSize, int* smemFoundTopKValues, SmemFinalType& smemFinal, int stride1, int rowStart, int topK) { + // Step 0 is the fp16 fast path; if it could not resolve top-K (threshold bin + // exceeded kNumFinalItems) we restart the fp32 radix from scratch in steps 1-3. + // Discard any candidates step 0 wrote into smemOutput so step 1 doesn't + // double-count valid entries that fall under both fp16 and fp32 thresholds + // (this is what produced 2x duplicated indices when many -FLT_MAX padding + // entries dominated the threshold bin in the multi-block merge path). + if constexpr (step == 1) + { + if (threadIdx.x == 0) + { + smemFoundTopKValues[0] = 0; + smemFinalDstIdx[0] = 0; + } + } // Clear the histogram. #pragma unroll for (int idx = threadIdx.x; idx < kNumBins; idx += kNumThreadsPerBlock) @@ -271,18 +292,13 @@ __device__ bool processHistogramStep(int const* indices, InputT const* logits, i // The threshold bin. thresholdBinIdx = smemThresholdBinIdx[0]; - // Skip auto-promote at step 0 when we'll continue: half-precision bins - // don't align with step 2's full-precision bit-pattern filter, so a - // step-0 promote would be double-counted at step 2. - bool const step0WillContinue = (step == 0) && (smemFinalBinSize[0] > kNumFinalItems); - auto processBins = [&](InputT logitIn, int idx) { float const logit = static_cast(logitIn); if (isPartialMatch(logit, logitPattern)) { uint32_t binIdx = extractBinIdx(logit); - if (binIdx < thresholdBinIdx && !step0WillContinue) + if (binIdx < thresholdBinIdx) { // The element is part of the top-k selection int dstIdx = atomicAdd(&smemFoundTopKValues[0], 1); @@ -368,21 +384,28 @@ __device__ bool processHistogramStep(int const* indices, InputT const* logits, i return smemFinalBinSize[0] > kNumFinalItems; } -// Smem holder for topKPerRowJob's final-sort. Always reserves -// BlockRadixSort::TempStorage so the sort algorithm can be picked at runtime -// (the union's size is dominated by FinalItems = 16 KB at our shapes). -template -struct TopKSmem +// Follows half - 11 - 11 - 10 bit iterations +template +static __device__ void topKPerRowJob(int const* indices, InputT const* logits, int rowStart, int rowEnd, + int* outIndices, float* outLogits, int stride1, int topK) { + // The number of slots for the final pass. static constexpr int kNumFinalItems = 2048; + // The number of elements per thread for the final sort. static constexpr int kNumFinalItemsPerThread = kNumFinalItems / kNumThreadsPerBlock; + // The class to sort the elements during the final pass. using FinalSort = cub::BlockRadixSort; - using FinalSortTempStorage = typename FinalSort::TempStorage; + using FinalSortTempStorage = std::conditional_t; + // The class to compute the inclusive prefix-sum over the histogram. using Scan = cub::BlockScan; + // The structure to store the final items (for the final pass). struct FinalItems { + // Shared memory to store the indices for the final pass. int indices[kNumFinalItems]; + // Shared memory to store the logits for the final pass. float logits[kNumFinalItems]; }; @@ -392,47 +415,34 @@ struct TopKSmem int data[kNumBins]; }; - union Final + // Shared memory to compute the block sort. + __shared__ union { FinalItems items; FinalSortTempStorage finalSort; Histogram histo; - }; - - Final smemFinal; - int smemThresholdBinIdx[1]; - int smemFinalDstIdx[1]; - int smemFinalBinSize[1]; - int smemFoundTopKValues[1]; -}; - -// Follows half - 11 - 11 - 10 bit iterations -template -static __device__ void topKPerRowJob(int const* indices, InputT const* logits, int rowStart, int rowEnd, - int* outIndices, float* outLogits, int stride1, int topK, TopKSmem& smem) -{ - static constexpr int kNumFinalItems = TopKSmem::kNumFinalItems; - static constexpr int kNumFinalItemsPerThread = TopKSmem::kNumFinalItemsPerThread; - using FinalSort = typename TopKSmem::FinalSort; - - auto& smemFinal = smem.smemFinal; - int* smemThresholdBinIdx = smem.smemThresholdBinIdx; - int* smemFinalDstIdx = smem.smemFinalDstIdx; - int* smemFinalBinSize = smem.smemFinalBinSize; - int* smemFoundTopKValues = smem.smemFoundTopKValues; + } smemFinal; // Shared memory to store the selected indices. // If we are processing using multiple blocks, we need to store the logits and // indices. extern __shared__ int32_t smemOutput[]; + // Shared memory to store the threshold bin. + __shared__ int smemThresholdBinIdx[1]; + // Shared memory counter to register the candidates for the final phase. + __shared__ int smemFinalDstIdx[1]; + // Shared memory to determine if the threshold bin fits in the final items. + __shared__ int smemFinalBinSize[1]; + // Shared memory to keep track of the top-k values found so far by the + // previous iterations + __shared__ int smemFoundTopKValues[1]; + // The length of the row. int rowLen = rowEnd - rowStart; // Shortcut if the length of the row is smaller than Top-K. Indices are not - // sorted by their corresponding logit. Unreachable when mergeBlocks=true: - // both merge callers pass rowLen = numBlocksPerRow * topK > topK. + // sorted by their corresponding logit. if (rowLen <= topK) { for (int rowIt = threadIdx.x; rowIt < rowLen; rowIt += kNumThreadsPerBlock) @@ -502,12 +512,10 @@ static __device__ void topKPerRowJob(int const* indices, InputT const* logits, i if (!continueToNextStep) { - // Sort the threshold-bin candidates. Insertion sort wins below ~512 - // items (O(n^2/T) with no fixed cost); BlockRadixSort wins above - // (constant cost padded to kNumFinalItems = 2048). - constexpr int kInsertionSortBranchThreshold = 512; - int const finalCount = smemFinalDstIdx[0]; - if (finalCount > kInsertionSortBranchThreshold) + // The histogram did not proceed to the final 10 bits, therefore we need to + // sort the final items The logits of the elements to be sorted in the final + // pass. + if constexpr (useRadixSort) { // Sorting with radix sort float finalLogits[kNumFinalItemsPerThread]; @@ -518,6 +526,12 @@ static __device__ void topKPerRowJob(int const* indices, InputT const* logits, i for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) { finalLogits[ii] = -FLT_MAX; + // Indices must be paired with the -FLT_MAX sentinel logit so unused + // slots survive SortDescendingBlockedToStriped and the post-sort copy + // as -1 rather than garbage. Without this, when the threshold bin is + // dominated by -FLT_MAX padding (multi-block merge path) the trailing + // top-K slots receive arbitrary register contents. + finalIndices[ii] = -1; } // Read the elements from SMEM. @@ -525,7 +539,7 @@ static __device__ void topKPerRowJob(int const* indices, InputT const* logits, i for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) { int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x; - if (srcIdx < finalCount) + if (srcIdx < smemFinalDstIdx[0]) { finalLogits[ii] = smemFinal.items.logits[srcIdx]; finalIndices[ii] = smemFinal.items.indices[srcIdx]; @@ -558,13 +572,13 @@ static __device__ void topKPerRowJob(int const* indices, InputT const* logits, i } else { - // Sorting with insertion sort. + // Sorting with insertion sort auto baseIdx = smemFoundTopKValues[0]; - for (int i = threadIdx.x; i < finalCount; i += kNumThreadsPerBlock) + for (int i = threadIdx.x; i < smemFinalDstIdx[0]; i += kNumThreadsPerBlock) { int outIndex = 0; auto logit = smemFinal.items.logits[i]; - for (int j = 0; j < finalCount; j++) + for (int j = 0; j < smemFinalDstIdx[0]; j++) { auto otherLogit = smemFinal.items.logits[j]; if (logit < otherLogit || (logit == otherLogit && i < j)) @@ -572,6 +586,7 @@ static __device__ void topKPerRowJob(int const* indices, InputT const* logits, i outIndex++; } } + // Store if outIndex is in bounds if (outIndex + baseIdx < topK) { smemOutput[outIndex + baseIdx] = smemFinal.items.indices[i]; @@ -593,10 +608,6 @@ static __device__ void topKPerRowJob(int const* indices, InputT const* logits, i outIndices[i] = smemOutput[i]; outLogits[i] = reinterpret_cast(smemOutput + topK)[i]; } - else if constexpr (mergeBlocks) - { - outIndices[i] = smemOutput[i]; - } else { if (stride1 == 1) @@ -613,7 +624,7 @@ static __device__ void topKPerRowJob(int const* indices, InputT const* logits, i } } // namespace -template +template static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowPrefill(InputT const* logits, int const* rowStarts, int const* rowEnds, int* outIndices, int stride0, int stride1, int const topK, int const offsetIndex) @@ -621,9 +632,6 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowPrefill( #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaGridDependencySynchronize(); #endif - // The number of bins in the histogram. - static constexpr int kNumBins = 2048; - // The row computed by this block. int rowIdx = blockIdx.x + offsetIndex; @@ -635,32 +643,30 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowPrefill( outIndices += static_cast(rowIdx) * topK; logits += static_cast(rowIdx) * stride0; - __shared__ TopKSmem smem; - topKPerRowJob( - nullptr, logits, rowStart, rowEnd, outIndices, nullptr, stride1, topK, smem); + topKPerRowJob( + nullptr, logits, rowStart, rowEnd, outIndices, nullptr, stride1, topK); #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaTriggerProgrammaticLaunchCompletion(); #endif } -template +template static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode(InputT const* logits, int const* seqLens, - int* outIndices, int stride0, int stride1, int const topK, int next_n, float* outLogits = nullptr, - int const numBlocksToMerge = 0, int const* indices = nullptr) + int* outIndices, int stride0, int stride1, int const topK, int next_n, int compressRatio, + float* outLogits = nullptr, int const numBlocksToMerge = 0, int const* indices = nullptr) { #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaGridDependencySynchronize(); #endif - // The number of bins in the histogram. - static constexpr int kNumBins = 2048; - // The row computed by this block. int rowIdx = blockIdx.x; // The range of logits within the row. int rowStart = 0; int seq_len = seqLens[rowIdx / next_n]; - int rowEnd = seq_len - next_n + (rowIdx % next_n) + 1; + int actual_kv_len = seq_len - next_n + (rowIdx % next_n) + 1; + int rowEnd = actual_kv_len / compressRatio; // Local pointers to this block if constexpr (!multipleBlocksPerRow && !mergeBlocks) @@ -683,9 +689,8 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode(I } logits += static_cast(rowIdx) * stride0; - __shared__ TopKSmem smem; - topKPerRowJob( - indices, logits, rowStart, rowEnd, outIndices, outLogits, stride1, topK, smem); + topKPerRowJob( + indices, logits, rowStart, rowEnd, outIndices, outLogits, stride1, topK); #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaTriggerProgrammaticLaunchCompletion(); #endif @@ -694,415 +699,29 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode(I namespace { -// Multi-pass radix: 4 launches (half-precision top-11 bits, then float bits -// [21..32), [10..21), [0..10)). Per-row state lives in DRAM scratch; candidate -// data uses two ping-pong buffers. Pass 3's last block emits the final top-K. -static constexpr int kRadixBins = 2048; - -// Per-row scratch state. The last block of each pass (picked via -// `finishedBlocks`) scans the global histogram and writes the next pass's -// threshold. -struct alignas(64) RadixState -{ - int candCount; // candidates entering current pass - int outIdx; // running outIndices write position - int kRemaining; // topK - outIdx - int filterCnt; // running candidate-buf write position - int thresholdBin; // prior pass's threshold; overwritten this pass's last block - int finishedBlocks; // last-block atomic, reset between passes - int thresholdLess; // count of bins below thresholdBin; pass-3 emit routes - // ties (bin == threshold) into a disjoint slot range. - int padding[1]; -}; - -// Common last-block trailer: prefix-scan the merged global histogram, locate -// the bin where the running count crosses kRemaining, stash to state for the -// next pass. step < 3 also resets the per-row global histogram. In step 1 -// the trailer writes st.candCount = full row length for pass 2 to consume -// (pass 1 reads the length inline from seqLens, so no init kernel is needed). -template -__device__ __forceinline__ void radixLastBlockTrailer(int* gHist, RadixState& st, int topK, int rowFullLen) -{ - using Scan = cub::BlockScan; - __shared__ typename Scan::TempStorage scanStorage; - __shared__ int s_thresholdBin; - __shared__ int s_runningBefore; - __shared__ int s_thresholdCount; - - if (threadIdx.x == 0) - { - s_thresholdBin = -1; - s_runningBefore = 0; - s_thresholdCount = 0; - } - __syncthreads(); - - // kRemaining for THIS pass: - // step 1: topK (no auto-promotes have happened yet). - // step 2/3: topK - outIdx (where outIdx was atomic-incremented during - // the filter loop). The histogram for the next pass's - // threshold pick must target this fresh value, not the - // stale state.kRemaining left over from the previous pass. - int const kRem = (step == 1) ? topK : (topK - st.outIdx); - if (kRem <= 0) - { - // All top-k slots already filled by auto-promotes in this pass — - // no need for a next pass to emit anything. - if (threadIdx.x == 0) - { - st.thresholdBin = -1; - st.candCount = 0; - st.kRemaining = 0; - st.finishedBlocks = 0; - st.filterCnt = 0; - (void) rowFullLen; - } - if constexpr (step < 3) - { - __syncthreads(); - for (int i = threadIdx.x; i < kRadixBins; i += kThreads) - gHist[i] = 0; - } - return; - } - constexpr int kRoundsPerScan = kRadixBins / kThreads; - int running = 0; - for (int r = 0; r < kRoundsPerScan; ++r) - { - int bin = r * kThreads + threadIdx.x; - int c = gHist[bin]; - int prefix, total; - Scan(scanStorage).ExclusiveSum(c, prefix, total); - prefix += running; - int next = prefix + c; - if (prefix < kRem && next >= kRem && s_thresholdBin == -1) - { - atomicCAS(&s_thresholdBin, -1, bin); - if (s_thresholdBin == bin) - { - s_runningBefore = prefix; - s_thresholdCount = c; - } - } - running += total; - __syncthreads(); - if (s_thresholdBin != -1) - break; - } - __syncthreads(); - - if (threadIdx.x == 0) - { - // If the cumsum over the whole histogram never reached kRem the row - // has fewer items than we still need (e.g. decode rows shorter than - // topK). Set thresholdBin to a sentinel above any valid bin so the - // next pass's `bin < thresholdBin` test accepts every surviving - // candidate as auto-promote. - st.thresholdBin = (s_thresholdBin == -1) ? kRadixBins : s_thresholdBin; - // s_runningBefore is the count of histogram items in bins < threshold. - // In the sentinel case it stays at its init 0; that's fine because - // pass-3's inline final-emit only uses thresholdLess to position the - // ties (bin == threshold) write base, and the sentinel branch has - // no items in the threshold bin (the sentinel is above all valid bins). - st.thresholdLess = s_runningBefore; - st.finishedBlocks = 0; - if constexpr (step == 1) - { - // Pass 2 also scans the full row from `logits`, so it reads - // st.candCount = full row length. Pass 1 did not write to it - // and the state struct started at zero (cudaMemsetAsync). - st.candCount = rowFullLen; - (void) s_thresholdCount; - (void) kRem; - } - else - { - st.candCount = st.filterCnt; - int newKRem = topK - st.outIdx; - if (newKRem < 0) - newKRem = 0; - st.kRemaining = newKRem; - st.filterCnt = 0; - } - } - if constexpr (step < 3) - { - __syncthreads(); - for (int i = threadIdx.x; i < kRadixBins; i += kThreads) - gHist[i] = 0; - } -} - -// Fused histogram + filter pass kernel. -// -// step == 1: pass 1, no filter, just histogram top-11 bits of the row. -// step == 2: pass 2 reads `logits` (pass 1 didn't write a candidate buffer), -// for each item: -// bin1 < thresholdBin1 → write to outIndices (auto-promote) -// bin1 == thresholdBin1 → append to candBufOut, count its bin2 -// in the histogram for pass 2 -// else → drop -// step == 3: same as step 2 but reads `candBufIn` (pass 2's output) and -// uses extractBinIdx<2> for the prior-bits check, extractBinIdx<3> -// for the histogram. -// -// Last block of every pass runs `radixLastBlockTrailer` to compute the -// next pass's threshold and reset cross-pass state. -template -static __global__ __launch_bounds__(kThreads) void radixPassKernel(InputT const* logits, int const* seqLens, - int* outIndices, int const* candBufIn, int* candBufOut, int* histograms, RadixState* state, int stride0, int next_n, - int topK) -{ - int rowIdx = blockIdx.y; - int blockInRow = blockIdx.x; - int blocksPerRow = gridDim.x; - - RadixState& st = state[rowIdx]; - int* gHist = histograms + static_cast(rowIdx) * kRadixBins; - - __shared__ int sHist[kRadixBins]; - for (int i = threadIdx.x; i < kRadixBins; i += kThreads) - sHist[i] = 0; - __syncthreads(); - - if constexpr (step == 1) - { - // Read seqLen inline so pass 1 does not depend on st.candCount being - // pre-initialised by a separate init kernel; the cudaMemsetAsync that - // zeroes state+histograms together is enough. Pass-1 trailer below - // writes st.candCount = seqLens[rowIdx] for pass 2 to consume. - int const rowEnd = seqLens[rowIdx / next_n] - next_n + (rowIdx % next_n) + 1; - InputT const* in = logits + static_cast(rowIdx) * stride0; - size_t threadRank = static_cast(blockInRow) * kThreads + threadIdx.x; - size_t numThreads = static_cast(blocksPerRow) * kThreads; - auto f = [&](InputT vIn, size_t /*idx*/) - { - float const v = static_cast(vIn); - uint32_t bin = extractBinIdx(v); - atomicAdd(&sHist[bin], 1); - }; - vectorized_process(threadRank, numThreads, in, static_cast(rowEnd), f); - } - else if constexpr (step == 2) - { - int const rowEnd = st.candCount; - InputT const* in = logits + static_cast(rowIdx) * stride0; - int* outIdxArr = outIndices + static_cast(rowIdx) * topK; - int* candArr = candBufOut + static_cast(rowIdx) * stride0; - int const prevThresh = st.thresholdBin; - size_t threadRank = static_cast(blockInRow) * kThreads + threadIdx.x; - size_t numThreads = static_cast(blocksPerRow) * kThreads; - auto f = [&](InputT vIn, size_t i) - { - float const v = static_cast(vIn); - int bin1 = static_cast(extractBinIdx<1>(v)); - if (bin1 < prevThresh) - { - int pos = atomicAdd(&st.outIdx, 1); - if (pos < topK) - outIdxArr[pos] = static_cast(i); - } - else if (bin1 == prevThresh) - { - int pos = atomicAdd(&st.filterCnt, 1); - candArr[pos] = static_cast(i); - uint32_t bin2 = extractBinIdx(v); - atomicAdd(&sHist[bin2], 1); - } - }; - vectorized_process(threadRank, numThreads, in, static_cast(rowEnd), f); - } - else // step == 3 - { - int const candCnt = st.candCount; - int const* candArrIn = candBufIn + static_cast(rowIdx) * stride0; - InputT const* in = logits + static_cast(rowIdx) * stride0; - int* outIdxArr = outIndices + static_cast(rowIdx) * topK; - int* candArrOut = candBufOut + static_cast(rowIdx) * stride0; - int const prevThresh = st.thresholdBin; - for (int i = blockInRow * kThreads + threadIdx.x; i < candCnt; i += blocksPerRow * kThreads) - { - int srcIdx = candArrIn[i]; - float v = static_cast(in[srcIdx]); - int bin2 = static_cast(extractBinIdx<2>(v)); - if (bin2 < prevThresh) - { - int pos = atomicAdd(&st.outIdx, 1); - if (pos < topK) - outIdxArr[pos] = srcIdx; - } - else if (bin2 == prevThresh) - { - int pos = atomicAdd(&st.filterCnt, 1); - candArrOut[pos] = srcIdx; - uint32_t bin3 = extractBinIdx(v); - atomicAdd(&sHist[bin3], 1); - } - } - } - __syncthreads(); - - for (int i = threadIdx.x; i < kRadixBins; i += kThreads) - { - int c = sHist[i]; - if (c) - atomicAdd(&gHist[i], c); - } - - __threadfence(); - __shared__ int isLast; - if (threadIdx.x == 0) - { - int prev = atomicAdd(&st.finishedBlocks, 1); - isLast = (prev == blocksPerRow - 1) ? 1 : 0; - } - __syncthreads(); - if (!isLast) - return; - - int const rowFullLen = (step == 1) ? (seqLens[rowIdx / next_n] - next_n + (rowIdx % next_n) + 1) : 0; - radixLastBlockTrailer(gHist, st, topK, rowFullLen); - - if constexpr (step == 3) - { - // Final emit, folded into the last block of pass 3. Scan candBufOut - // (top 22 bits == thresholdBin2) and route items into outIndices by - // bin3 vs thresholdBin3. Two-counter scheme so threshold-bin ties - // don't race definite top-k items on the same atomic: - // bin3 < thresh3 → slots [ltBase, ltBase + prefix3) - // bin3 == thresh3 → slots [ltBase + prefix3, topK) - __syncthreads(); - int const filterCnt = st.candCount; - int const thresh3 = st.thresholdBin; - int const prefix3 = st.thresholdLess; - int const* candArr = candBufOut + static_cast(rowIdx) * stride0; - InputT const* in = logits + static_cast(rowIdx) * stride0; - int* outIdxArr = outIndices + static_cast(rowIdx) * topK; - int const ltBase = st.outIdx; // already at outBase here - int const eqBase = ltBase + prefix3; - int const eqCap = topK - eqBase; // ≥ 0 by trailer invariant - __shared__ int sEqEmitted; - if (threadIdx.x == 0) - sEqEmitted = 0; - __syncthreads(); - for (int i = threadIdx.x; i < filterCnt; i += kThreads) - { - int srcIdx = candArr[i]; - float v = static_cast(in[srcIdx]); - int bin3 = static_cast(extractBinIdx<3>(v)); - if (bin3 < thresh3) - { - // atomicAdd on st.outIdx is safe: by construction exactly - // prefix3 items fall in this branch, so pos stays in - // [ltBase, eqBase) which is strictly inside [0, topK). - int pos = atomicAdd(&st.outIdx, 1); - outIdxArr[pos] = srcIdx; - } - else if (bin3 == thresh3) - { - int pos = atomicAdd(&sEqEmitted, 1); - if (pos < eqCap) - outIdxArr[eqBase + pos] = srcIdx; - } - } - __syncthreads(); - if (threadIdx.x == 0) - { - int eq = sEqEmitted < eqCap ? sEqEmitted : eqCap; - int filled = eqBase + eq; - if (filled > topK) - filled = topK; - for (int i = filled; i < topK; ++i) - outIdxArr[i] = -1; - } - // Reset the per-row global histogram and st.outIdx so the next call - // sees a clean state without a per-call cudaMemsetAsync. Caller must - // zero-initialize the scratch buffer before the first call. - __syncthreads(); - for (int i = threadIdx.x; i < kRadixBins; i += kThreads) - gHist[i] = 0; - if (threadIdx.x == 0) - st.outIdx = 0; - } -} - -// Scratch layout (uint8 buffer, 64-byte aligned regions): -// RadixState[numRows] -// int histograms[numRows * kRadixBins] (zeroed on first call by -// torch::zeros allocator; pass-3 -// trailer zeroes for subsequent -// calls) -// int candBuf1[numRows * stride0] (pass 2 → pass 3 input) -// int candBuf2[numRows * stride0] (pass 3 → fused final filter) -static size_t radixScratchBytes(int numRows, int numColumns) -{ - auto roundUp = [](size_t x) { return (x + 63) & ~size_t(63); }; - size_t s = 0; - s += roundUp(sizeof(RadixState) * numRows); - s += roundUp(sizeof(int) * static_cast(numRows) * kRadixBins); - s += roundUp(sizeof(int) * static_cast(numRows) * numColumns); - s += roundUp(sizeof(int) * static_cast(numRows) * numColumns); - return s; -} - -template -static void launchMultiPassRadix(void* scratch, InputT const* logits, int const* seqLens, int* outIndices, int numRows, - int numColumns, int topK, int stride0, int next_n, cudaLaunchAttribute const* attrs, cudaStream_t stream) -{ - auto roundUp = [](size_t x) { return (x + 63) & ~size_t(63); }; - char* base = static_cast(scratch); - RadixState* state = reinterpret_cast(base); - base += roundUp(sizeof(RadixState) * numRows); - int* histograms = reinterpret_cast(base); - base += roundUp(sizeof(int) * static_cast(numRows) * kRadixBins); - int* candBuf1 = reinterpret_cast(base); - base += roundUp(sizeof(int) * static_cast(numRows) * numColumns); - int* candBuf2 = reinterpret_cast(base); - - int sm_cnt = 132; - { - int dev = 0; - cudaGetDevice(&dev); - cudaDeviceGetAttribute(&sm_cnt, cudaDevAttrMultiProcessorCount, dev); - } - // Block fan-out heuristic: target ~4 active blocks/SM (one wave at the - // achievable occupancy of radixPassKernel<512, 1>), with a per-block work - // floor of 2048 items (4 items/thread at 512-wide). - int targetTotalBlocks = sm_cnt * 4; - int numBlocksPerRow = (targetTotalBlocks + numRows - 1) / numRows; - int maxByCols = numColumns / 2048; - if (numBlocksPerRow > maxByCols) - numBlocksPerRow = maxByCols; - if (numBlocksPerRow < 1) - numBlocksPerRow = 1; - - constexpr int kPassThreads = 512; - - auto launchPass = [&](void const* kernel, int const* candIn, int* candOut) - { - cudaLaunchConfig_t cfg{}; - cfg.gridDim = dim3(numBlocksPerRow, numRows); - cfg.blockDim = kPassThreads; - cfg.dynamicSmemBytes = 0; - cfg.stream = stream; - cfg.numAttrs = 1; - cfg.attrs = const_cast(attrs); - void* args[] = {(void*) &logits, (void*) &seqLens, (void*) &outIndices, (void*) &candIn, (void*) &candOut, - (void*) &histograms, (void*) &state, (void*) &stride0, (void*) &next_n, (void*) &topK}; - cudaLaunchKernelExC(&cfg, kernel, args); - }; - - launchPass( - reinterpret_cast(&radixPassKernel), (int const*) nullptr, (int*) nullptr); - launchPass( - reinterpret_cast(&radixPassKernel), (int const*) nullptr, candBuf1); - // Pass 3 emits the final top-K inline in its last-block trailer (see - // radixPassKernel) instead of requiring a separate filter launch. - launchPass( - reinterpret_cast(&radixPassKernel), (int const*) candBuf1, candBuf2); -} - -// Architecture-derived GVR eligibility bounds (cached per-process). +// Insertion vs radix crossover for the single-block path. Radix always pays a +// histogram-clear + per-step scan over kNumBins bins (~4 refinement passes); +// insertion only maintains a topK SMEM array. The crossover is where insertion's +// O(numColumns * lg topK) catches up to radix's O(numColumns + kNumBins) per +// pass — measured empirically around 6 histograms of work. +constexpr int kSortingAlgorithmThreshold = 6 * kNumBins; +// Force the multi-block split-and-merge path above this column count: per-block +// work would otherwise dwarf the merge-pass cost. Callers may override via the +// splitWorkThreshold parameter. Tuned empirically; below this width the merge +// launch overhead isn't worth saving the per-block radix cost. +constexpr int kDefaultSplitWorkThreshold = 200 * 1000; +// Cap blocks-per-row in the multi-block path. Bounds the aux-buffer size +// (numRows * blocksPerRow * topK * sizeof(int32)) and the merge-step input width +// so the second-pass merge kernel stays SMEM-resident. +constexpr int kMaxBlocksPerRowDecode = 10; +// Each sub-block must amortize the radix histogram overhead, i.e. cover at least +// one full histogram pass worth of columns. +constexpr int kDecodeMinColsPerSubBlock = kNumBins; + +// Scheme X bound calculator — shared between fp32 and bf16/fp16 dispatchers. +// Caches hardware attrs (SM count, L2 capacity) and the small-N threshold +// once per process via std::call_once. Per-call cost is just two reads +// from cached static variables plus a small arithmetic block, no syscalls. struct SchemeXBounds { int smCount; @@ -1151,33 +770,148 @@ inline SchemeXBounds getSchemeXBounds(int numColumns, int bytesPerElem) return b; } -// Unified dispatcher (fp32 / bf16 / fp16). Each tier's kernel is templated on -// InputT and casts to float at HBM-read sites; the scratch buffer is uint8. -template -void invokeIndexerTopKDecodeImpl(InputT const* logits, int const* seqLens, int* indices, int const splitWorkThreshold, - int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, int const topK, - int const* preIdx, int const preIdxStride, int const preIdxCount, InputT* heuristicScratch, - cudaStream_t const stream, void* scratch, size_t scratchBytes, bool is_prefill) +} // namespace + +int computeIndexerTopKDecodeBlocksPerRow(int numRows, int numColumns, int splitWorkThreshold) +{ + if (numRows <= 0) + { + return 1; + } + + int const forceSplitThreshold = splitWorkThreshold > 0 ? splitWorkThreshold : kDefaultSplitWorkThreshold; + + // Preserve original behavior for very long sequences. + if (numColumns >= forceSplitThreshold) + { + return kMaxBlocksPerRowDecode; + } + + // Query the actual SM count from the driver so the dispatch tracks the + // hardware rather than a baked-in target (H100=132, B200=148, …). + auto const bounds = getSchemeXBounds(numColumns, /*bytesPerElem=*/4); + TLLM_CHECK_WITH_INFO(bounds.smCount > 0, "indexerTopK: failed to query device SM count"); + int const smCount = bounds.smCount; + int const maxByCols = std::max(1, numColumns / kDecodeMinColsPerSubBlock); + int const maxBp = std::min(maxByCols, kMaxBlocksPerRowDecode); + + int blocksPerRow; + if (numRows < smCount / 2) + { + // Sub-half-wave band: bp=2 by itself leaves SMs idle (numRows × 2 < smCount), + // so sweep bp ∈ [2, maxBp] for the choice that minimizes waves(bp) / bp, + // where waves(bp) = ceil(numRows * bp / smCount). The wave-quantization-aware + // sweep avoids spilling one extra block per row across a wave boundary, which + // is what a naive ceil(smCount / numRows) target would do. + int bestBp = 1; + int bestWaves = 1; // numRows < smCount/2 → bp=1 always fits in a single wave + for (int bp = 2; bp <= maxBp; ++bp) + { + int const totalBlocks = numRows * bp; + int const waves = (totalBlocks + smCount - 1) / smCount; + // waves / bp < bestWaves / bestBp ⇔ waves * bestBp < bestWaves * bp + if (waves * bestBp < bestWaves * bp) + { + bestWaves = waves; + bestBp = bp; + } + } + blocksPerRow = bestBp; + } + else + { + // numRows >= smCount/2: bp=2 saturates SMs with one wave (numRows*2 >= smCount) + // and stays on the multi-block split+merge path. Crucially this path uses + // a different kernel instantiation than bp=1, and only the bp=1 single-block + // radix kernel pays the wave-scheduling cliff when gridDim.x approaches + // smCount (measured on B200, cols=196608, topK=2048: BS=131=125us, + // BS=132=312us, BS=148=390us — the cliff disappears entirely with bp=2). + // bp=2 is also the cheapest split (smallest merge input); larger bp piles + // on merge-pass overhead without proportional gain on the shapes measured. + // Falls back to 1 only when maxByCols caps it at 1 for very narrow rows + // (numColumns < kDecodeMinColsPerSubBlock). + blocksPerRow = std::min(2, maxByCols); + } + return std::max(1, blocksPerRow); +} + +void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indices, float* outLogitsAux, + int* outIndicesAux, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, + int const stride1, int const next_n, int const topK, int const* preIdx, int const preIdxStride, + int const preIdxCount, float* heuristicScratch, int const compressRatio, cudaStream_t const stream) { - // Split-work cutoff: matches main's 200k default. is_prefill forces - // single-block via a 1<<30 threshold no shape can reach: prefill chunks are - // bounded by max_num_tokens, well below the multi-pass radix crossover at - // any practical setting. - int const adaptiveSplitWorkThreshold = is_prefill ? (1 << 30) : 200 * 1000; - int const effectiveSplitWorkThreshold = splitWorkThreshold > 0 ? splitWorkThreshold : adaptiveSplitWorkThreshold; constexpr int kNumThreadsPerBlock = 512; + int const effectiveSplitWorkThreshold = splitWorkThreshold > 0 ? splitWorkThreshold : kDefaultSplitWorkThreshold; + + // ======================================================================== + // Small-N dispatch axis. + // + // GVR Heuristic Top-K has a *fixed* per-launch overhead from Phase-1 + // (preIdx stats reduction over M=2048) and Phase-4 (2048-bin histogram + // snap), totaling ~11 µs regardless of N. For small N (≤16K), this + // fixed cost dominates and the kernel loses to the existing + // insertion-sort/radix path. Empirically (random data, B200 BS=1): + // N=8192 : Heuristic 16.5 µs vs Radix 11.2 µs (radix 1.47× faster) + // N=16384 : Heuristic 21.9 µs vs Radix 22.0 µs (parity) + // N=32768 : Heuristic 26.1 µs vs Radix 32.9 µs (heuristic 1.26× faster) + // N=131072 : Heuristic 43.4 µs vs Radix 76.1 µs (heuristic 1.75× faster) + // + // Route N < kSeqSmall to the existing Radix/Insertion path (which itself + // splits at kSortingAlgorithmThreshold=12288). kSeqSmall is set at the + // empirical crossover point. + // + // ======================================================================== + // Architecture-derived BS-threshold dispatch — jointly bounded by + // occupancy AND L2 cache capacity. + // + // Two physical constraints bound when the per-row heuristic kernel + // remains faster than a radix streaming kernel: + // + // (A) Occupancy bound — 3·SM − SM/8 (wave geometry + setup margin) + // Each CTA uses ~58 KB SMEM (fixed, independent of N), so B200's + // 228 KB dynamic SMEM allows max 3 CTA/SM. Above 3·SM rows per + // launch, tail-wave imbalance causes stragglers. The -SM/8 margin + // (~1/8 wave) covers CTA setup + L2 ingestion overhead. + // On B200(148 SM): 3×148 − 18 = 426. + // + // (B) L2 cache bound — 0.9·L2 / (4·N) per-CTA logits fit + // Each CTA streams its row (N×4B) through L2 per Phase-2 iter. + // With num_concurrent_CTAs × N × 4B > L2, eviction dominates. + // On B200(126 MB L2) with N=70K: 0.9·126MB/(4·70690) ≈ 440, + // which is ~ equal to (A)=426 — the two constraints cross over + // near the SWE-Bench data point. + // For N > 73K the L2 bound tightens below (A) and must take + // over; e.g. N=128K → kBsL2=238, N=196K → kBsL2=155. + // + // Dispatch threshold = min(kBsWave, kBsL2), still data-agnostic (only + // queries hardware attrs). At N≈70K both bounds produce ~426, so the + // L2 axis is a no-op there; for larger N it auto-tightens the threshold. + // + // Small-N lower bound `kSeqSmall` (default 12288) lets the Heuristic + // axis take over wherever the original Radix-radix branch would have + // triggered. Random-data benchmarks suggest the crossover is 16384, + // but workloads with strongly preIdx-correlated logits make P1 stats + // accurate and P2 converge in 1-2 iterations, shifting the real + // crossover into the [12288, 16384] band. Configurable via + // TRTLLM_HEURISTIC_NMIN env (>=1024). + // ======================================================================== + auto const bounds = getSchemeXBounds(numColumns, /*bytesPerElem=*/4); + int const kBsWave = bounds.kBsWave; + int const kBsL2 = bounds.kBsL2; + int const kBsLarge = bounds.kBsLarge; + int const kSeqSmall = bounds.kSeqSmall; - // GVR eligibility (matches main's rule): supported K, stride1 contiguous, - // preIdx + scratch provided, numColumns in [kSeqSmall, splitWorkThreshold), - // and numRows below the architecture-derived wave/L2 bound. is_prefill - // suppresses GVR through effectiveSplitWorkThreshold being huge. - auto const bounds = getSchemeXBounds(numColumns, /*bytesPerElem=*/static_cast(sizeof(InputT))); bool const isSupportedTopK = (topK == 512 || topK == 1024 || topK == 2048); - bool const canUseHeuristic = preIdx != nullptr && stride1 == 1 && isSupportedTopK && preIdxCount == topK - && preIdxStride >= preIdxCount && heuristicScratch != nullptr && numColumns >= bounds.kSeqSmall - && numColumns < effectiveSplitWorkThreshold && numRows < bounds.kBsLarge; - - // Env-gated dispatch trace (TRTLLM_SCHEMEX_DEBUG=1). + // compressRatio == 1: DSv3.2 indexer (no compressor). + // compressRatio == 4: DSv4 indexer (overlap compressor); logits/preIdx in + // compressed-token-index space. Kernel handles N = actual_kv_len/cr and + // forces preIdxOffset=0 internally for cr != 1. + bool const compressRatioOk = (compressRatio == 1 || compressRatio == 4); + bool const canUseHeuristic = compressRatioOk && preIdx != nullptr && stride1 == 1 && isSupportedTopK + && preIdxCount == topK && preIdxStride >= preIdxCount && numColumns < effectiveSplitWorkThreshold + && numColumns >= kSeqSmall && heuristicScratch != nullptr && numRows < kBsLarge; + + // Optional env-gated dispatch trace (set TRTLLM_SCHEMEX_DEBUG=1 to enable) { static std::once_flag sDebugOnceFlag; static bool sDebug = false; @@ -1189,87 +923,204 @@ void invokeIndexerTopKDecodeImpl(InputT const* logits, int const* seqLens, int* }); if (sDebug) { - fprintf(stderr, "[Scheme X] numRows=%d numColumns=%d kBsLarge=%d kSeqSmall=%d -> %s path\n", numRows, - numColumns, bounds.kBsLarge, bounds.kSeqSmall, canUseHeuristic ? "Heuristic" : "Radix"); + fprintf(stderr, + "[Scheme X] numRows=%d numColumns=%d kBsWave=%d kBsL2=%d kBsLarge=%d kSeqSmall=%d smCount=%d " + "L2=%dMB -> %s path%s\n", + numRows, numColumns, kBsWave, kBsL2, kBsLarge, kSeqSmall, bounds.smCount, + bounds.l2Bytes / (1024 * 1024), canUseHeuristic ? "Heuristic" : "Radix", + (numColumns < kSeqSmall) ? " (small-N route)" : ""); } } if (canUseHeuristic) { launchHeuristicTopKDecode(logits, seqLens, preIdx, indices, heuristicScratch, stride0, next_n, topK, - preIdxStride, preIdxCount, numRows, stream); + preIdxStride, preIdxCount, numRows, compressRatio, stream); + sync_check_cuda_error(stream); + return; } - else if (numColumns < effectiveSplitWorkThreshold) + + int const blocksPerRow = computeIndexerTopKDecodeBlocksPerRow(numRows, numColumns, splitWorkThreshold); + + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); + + if (blocksPerRow == 1) { - // Single-block tier: one CTA per row. - auto* kernel_instance = &topKPerRowDecode; + // Single block per row. Below kSortingAlgorithmThreshold use insertion sort, + // above use the histogram-radix path. + bool const useRadixSort = numColumns >= kSortingAlgorithmThreshold; + auto* kernel_instance = useRadixSort ? &topKPerRowDecode + : &topKPerRowDecode; + cudaLaunchConfig_t config; config.gridDim = numRows; config.blockDim = kNumThreadsPerBlock; config.dynamicSmemBytes = topK * sizeof(int32_t); config.stream = stream; - cudaLaunchAttribute attrs[1]; - attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); config.numAttrs = 1; config.attrs = attrs; + cudaLaunchKernelEx(&config, kernel_instance, logits, seqLens, indices, stride0, stride1, topK, next_n, - /*outLogits=*/nullptr, /*numBlocksToMerge=*/0, /*indices=*/nullptr); + compressRatio, nullptr, 0, nullptr); } else { - // Multi-pass radix. radixPassKernel reads logits contiguously, so - // strided inputs would rank the wrong values — gate on stride1 == 1. - // (The single-block tier handles stride1 != 1 via topKPerRowJob's - // strided fallback.) - TLLM_CHECK_WITH_INFO(stride1 == 1, "indexer top-k split-work tier (multi-pass radix) requires stride1 == 1."); - TLLM_CHECK_WITH_INFO(scratch != nullptr && scratchBytes >= radixScratchBytes(numRows, numColumns), - "indexer top-k split-work tier: scratch buffer missing or too small."); - cudaLaunchAttribute radixAttrs[1]; - radixAttrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - radixAttrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); - launchMultiPassRadix( - scratch, logits, seqLens, indices, numRows, numColumns, topK, stride0, next_n, radixAttrs, stream); + // Split each row across `blocksPerRow` blocks, then merge with a second pass. + auto* kernel_instance_part1 = &topKPerRowDecode; + cudaLaunchConfig_t config_part1; + config_part1.gridDim = dim3(numRows, blocksPerRow); + config_part1.blockDim = kNumThreadsPerBlock; + config_part1.dynamicSmemBytes = 2 * topK * sizeof(int32_t); + config_part1.stream = stream; + config_part1.numAttrs = 1; + config_part1.attrs = attrs; + + cudaLaunchKernelEx(&config_part1, kernel_instance_part1, logits, seqLens, outIndicesAux, stride0, stride1, topK, + next_n, compressRatio, outLogitsAux, 0, nullptr); + + constexpr int kNumThreadsPerBlockMerge = 1024; + auto* kernel_instance_part2 = &topKPerRowDecode; + cudaLaunchConfig_t config_part2; + config_part2.gridDim = numRows; + config_part2.blockDim = kNumThreadsPerBlockMerge; + config_part2.dynamicSmemBytes = topK * sizeof(int32_t); + config_part2.stream = stream; + config_part2.numAttrs = 1; + config_part2.attrs = attrs; + + cudaLaunchKernelEx(&config_part2, kernel_instance_part2, outLogitsAux, seqLens, indices, blocksPerRow * topK, 1, + topK, next_n, 1, nullptr, blocksPerRow, outIndicesAux); } sync_check_cuda_error(stream); } -} // anonymous namespace +// ============================================================================ +// bf16 / fp16 dispatcher overloads +// ============================================================================ +// Reuses the BS-threshold + small-N dispatch axes (kBsLarge, kSeqSmall) from +// the fp32 dispatcher, except kBsL2 uses sizeof(InputT) bytes/element instead +// of 4 — L2 footprint is half, so bf16/fp16 path remains valid for larger BS +// than fp32 at the same N. +// +// Fallback chain when GVR-Heuristic preconditions are not met (preIdx +// missing, BS too large, or numColumns < kSeqSmall): +// numColumns < kSortingAlgorithmThreshold (12288) → insertion sort +// kSortingAlgorithmThreshold ≤ numColumns < splitWorkThreshold → radix sort +// numColumns ≥ splitWorkThreshold (200K default) → unsupported +// +// Insertion + radix tiers use the same topKPerRowDecode kernel as fp32 with +// InputT propagated through; the histogram and sort steps operate on float +// keys after a static_cast(InputT) at HBM-read sites, so accuracy is +// identical to casting input to fp32 before the kernel. +// +// The split-work tier requires float aux buffers (outLogitsAux / +// outIndicesAux) that the bf16/fp16 entry does not expose; callers in that +// regime must use the fp32 entry. -void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indices, int const splitWorkThreshold, - int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, int const topK, - int const* preIdx, int const preIdxStride, int const preIdxCount, float* heuristicScratch, - cudaStream_t const stream, void* scratch, size_t scratchBytes, bool is_prefill) +namespace { - invokeIndexerTopKDecodeImpl(logits, seqLens, indices, splitWorkThreshold, numRows, numColumns, stride0, - stride1, next_n, topK, preIdx, preIdxStride, preIdxCount, heuristicScratch, stream, scratch, scratchBytes, - is_prefill); -} -size_t indexerTopKDecodeScratchBytes(int numRows, int numColumns, int /*topK*/) +template +void invokeIndexerTopKDecodeDtype(InputT const* logits, int const* seqLens, int* indices, int const splitWorkThreshold, + int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, int const topK, + int const* preIdx, int const preIdxStride, int const preIdxCount, InputT* heuristicScratch, int const compressRatio, + cudaStream_t const stream) { - return radixScratchBytes(numRows, numColumns); + static_assert(std::is_same_v || std::is_same_v, + "invokeIndexerTopKDecodeDtype is for bf16/fp16 only"); + + constexpr int kNumThreadsPerBlock = 512; + int const effectiveSplitWorkThreshold = splitWorkThreshold > 0 ? splitWorkThreshold : kDefaultSplitWorkThreshold; + + // bf16/fp16: bytes_per_element = sizeof(InputT) = 2 → kBsL2 doubles vs fp32. + auto const bounds = getSchemeXBounds(numColumns, /*bytesPerElem=*/static_cast(sizeof(InputT))); + int const kBsLarge = bounds.kBsLarge; + int const kSeqSmall = bounds.kSeqSmall; + + bool const isSupportedTopK = (topK == 512 || topK == 1024 || topK == 2048); + // See fp32 path: cr==1 (V3.2) and cr==4 (V4 indexer) are both supported. + bool const compressRatioOk = (compressRatio == 1 || compressRatio == 4); + bool const canUseHeuristic = compressRatioOk && preIdx != nullptr && stride1 == 1 && isSupportedTopK + && preIdxCount == topK && preIdxStride >= preIdxCount && numColumns < effectiveSplitWorkThreshold + && numColumns >= kSeqSmall && heuristicScratch != nullptr && numRows < kBsLarge; + + if (canUseHeuristic) + { + launchHeuristicTopKDecode(logits, seqLens, preIdx, indices, heuristicScratch, stride0, next_n, topK, + preIdxStride, preIdxCount, numRows, compressRatio, stream); + } + else if (numColumns < kSortingAlgorithmThreshold) + { + // Insertion sort path — InputT propagated; histogram/sort run on float keys. + auto* kernel_instance = &topKPerRowDecode; + + cudaLaunchConfig_t config; + config.gridDim = numRows; + config.blockDim = kNumThreadsPerBlock; + config.dynamicSmemBytes = topK * sizeof(int32_t); + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); + config.numAttrs = 1; + config.attrs = attrs; + + cudaLaunchKernelEx(&config, kernel_instance, logits, seqLens, indices, stride0, stride1, topK, next_n, + compressRatio, nullptr, 0, nullptr); + } + else if (numColumns < effectiveSplitWorkThreshold) + { + // Radix sort path — InputT propagated; histogram/sort run on float keys. + auto* kernel_instance = &topKPerRowDecode; + + cudaLaunchConfig_t config; + config.gridDim = numRows; + config.blockDim = kNumThreadsPerBlock; + config.dynamicSmemBytes = topK * sizeof(int32_t); + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); + config.numAttrs = 1; + config.attrs = attrs; + + cudaLaunchKernelEx(&config, kernel_instance, logits, seqLens, indices, stride0, stride1, topK, next_n, + compressRatio, nullptr, 0, nullptr); + } + else + { + TLLM_CHECK_WITH_INFO(false, + "indexer_topk_decode bf16/fp16 path does not support numColumns >= splitWorkThreshold " + "(split-work path requires float aux buffers not exposed in the bf16/fp16 entry). " + "Got numColumns=%d splitWorkThreshold=%d. Use the fp32 entry for this regime.", + numColumns, effectiveSplitWorkThreshold); + } + + sync_check_cuda_error(stream); } +} // anonymous namespace + void invokeIndexerTopKDecode(__nv_bfloat16 const* logits, int const* seqLens, int* indices, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, int const topK, int const* preIdx, int const preIdxStride, int const preIdxCount, - __nv_bfloat16* heuristicScratch, cudaStream_t const stream, void* scratch, size_t scratchBytes, bool is_prefill) + __nv_bfloat16* heuristicScratch, int const compressRatio, cudaStream_t const stream) { - invokeIndexerTopKDecodeImpl<__nv_bfloat16>(logits, seqLens, indices, splitWorkThreshold, numRows, numColumns, - stride0, stride1, next_n, topK, preIdx, preIdxStride, preIdxCount, heuristicScratch, stream, scratch, - scratchBytes, is_prefill); + invokeIndexerTopKDecodeDtype<__nv_bfloat16>(logits, seqLens, indices, splitWorkThreshold, numRows, numColumns, + stride0, stride1, next_n, topK, preIdx, preIdxStride, preIdxCount, heuristicScratch, compressRatio, stream); } void invokeIndexerTopKDecode(__half const* logits, int const* seqLens, int* indices, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, int const topK, - int const* preIdx, int const preIdxStride, int const preIdxCount, __half* heuristicScratch, - cudaStream_t const stream, void* scratch, size_t scratchBytes, bool is_prefill) + int const* preIdx, int const preIdxStride, int const preIdxCount, __half* heuristicScratch, int const compressRatio, + cudaStream_t const stream) { - invokeIndexerTopKDecodeImpl<__half>(logits, seqLens, indices, splitWorkThreshold, numRows, numColumns, stride0, - stride1, next_n, topK, preIdx, preIdxStride, preIdxCount, heuristicScratch, stream, scratch, scratchBytes, - is_prefill); + invokeIndexerTopKDecodeDtype<__half>(logits, seqLens, indices, splitWorkThreshold, numRows, numColumns, stride0, + stride1, next_n, topK, preIdx, preIdxStride, preIdxCount, heuristicScratch, compressRatio, stream); } void invokeIndexerTopKPrefill(float const* logits, int const* rowStarts, int const* rowEnds, int* indices, @@ -1278,10 +1129,18 @@ void invokeIndexerTopKPrefill(float const* logits, int const* rowStarts, int con { constexpr int kNumThreadsPerBlock = 512; - // One launch over all rows; the per-row sort algorithm is picked at - // runtime inside topKPerRowJob. - topKPerRowPrefill<<>>( - logits, rowStarts, rowEnds, indices, stride0, stride1, topK, 0); + int numInsertionBlocks = std::min(numRows, kSortingAlgorithmThreshold); + topKPerRowPrefill + <<>>( + logits, rowStarts, rowEnds, indices, stride0, stride1, topK, 0); + + if (numRows > kSortingAlgorithmThreshold) + { + int numRadixBlocks = numRows - kSortingAlgorithmThreshold; + topKPerRowPrefill + <<>>( + logits, rowStarts, rowEnds, indices, stride0, stride1, topK, kSortingAlgorithmThreshold); + } sync_check_cuda_error(stream); } @@ -1293,8 +1152,6 @@ bool canIndexerTopKDecodeUseGvr(int numRows, int numColumns, int topK, int bytes { return false; } - // Mirrors the dispatcher's effectiveSplitWorkThreshold default. - constexpr int kDefaultSplitWorkThreshold = 200 * 1000; auto const bounds = getSchemeXBounds(numColumns, bytesPerElem); return numColumns >= bounds.kSeqSmall && numColumns < kDefaultSplitWorkThreshold && numRows < bounds.kBsLarge; } diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/CMakeLists.txt b/cpp/tensorrt_llm/kernels/mhcKernels/CMakeLists.txt new file mode 100644 index 000000000000..58dc760fcf9d --- /dev/null +++ b/cpp/tensorrt_llm/kernels/mhcKernels/CMakeLists.txt @@ -0,0 +1,35 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 1993-2024 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>) + +# 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) 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..9e74d4eb20f7 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh @@ -0,0 +1,1602 @@ +/* + * Copyright (c) 2022-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..86586edc099a --- /dev/null +++ b/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu @@ -0,0 +1,830 @@ +/* + * Copyright (c) 2022-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. + +#include "fused_tf32_pmap_gemm.cuh" +#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_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_BLOCK_K = 64; +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