diff --git a/cpp/tensorrt_llm/CMakeLists.txt b/cpp/tensorrt_llm/CMakeLists.txt index 35369f803d9b..d5da630a27a5 100644 --- a/cpp/tensorrt_llm/CMakeLists.txt +++ b/cpp/tensorrt_llm/CMakeLists.txt @@ -199,6 +199,7 @@ set(TRTLLM_LINK_LIBS mamba2_mtp_ssm_cache_src kda_decode_src kimi_k3_attn_res_src + glm_moe_dsa_kernels_src kimi_k3_model_src ws_layernorm_src fusedGatedRMSNormQuant_src diff --git a/cpp/tensorrt_llm/kernels/CMakeLists.txt b/cpp/tensorrt_llm/kernels/CMakeLists.txt index 5f2190da30bf..1074efd244c4 100644 --- a/cpp/tensorrt_llm/kernels/CMakeLists.txt +++ b/cpp/tensorrt_llm/kernels/CMakeLists.txt @@ -32,6 +32,7 @@ add_subdirectory(fusedGatedRMSNormQuant) add_subdirectory(mamba2MTPSSMCache) add_subdirectory(kdaDecode) add_subdirectory(kimiK3AttnRes) +add_subdirectory(glmMoeDsa) add_subdirectory(mhcKernels) add_subdirectory(compressorKernels) add_subdirectory(fusedSampling) @@ -78,6 +79,15 @@ list(FILTER SRC_CPP EXCLUDE REGEX "kdaDecode/.*") list(FILTER SRC_CU EXCLUDE REGEX "kdaDecode/.*") list(FILTER SRC_CPP EXCLUDE REGEX "kimiK3AttnRes/.*") list(FILTER SRC_CU EXCLUDE REGEX "kimiK3AttnRes/.*") +# The GLM-5.2 GB300 target's kernels build in glmMoeDsa/'s own library; the +# other glmMoeDsa/ kernels stay in this glob. +list( + FILTER + SRC_CU + EXCLUDE + REGEX + "glmMoeDsa/(crossRankTopk|crossRankTopkClusterMerge|crossRankTopkGlobalEmit|denseMlp|draftLiveRowHop|draftRoutedMoe|draftTail|indexerLayerRopeKvAppend|postAttentionHop|postAttentionHopFromPartials|postAttentionHopRouting|ropeKvAppend|sparseMlaDecode)\\.cu$" +) list(FILTER SRC_CPP EXCLUDE REGEX "mhcKernels/.*") list(FILTER SRC_CU EXCLUDE REGEX "mhcKernels/.*") list(FILTER SRC_CPP EXCLUDE REGEX "compressorKernels/.*") diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/CMakeLists.txt b/cpp/tensorrt_llm/kernels/glmMoeDsa/CMakeLists.txt new file mode 100644 index 000000000000..db750686e8db --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/CMakeLists.txt @@ -0,0 +1,42 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); you may not +# use this file except in compliance with the License. You may obtain a copy of +# the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +# License for the specific language governing permissions and limitations under +# the License. + +add_library( + glm_moe_dsa_kernels_src OBJECT + crossRankTopk.cu + crossRankTopkClusterMerge.cu + crossRankTopkGlobalEmit.cu + denseMlp.cu + draftLiveRowHop.cu + draftRoutedMoe.cu + draftTail.cu + indexerLayerRopeKvAppend.cu + postAttentionHop.cu + postAttentionHopFromPartials.cu + postAttentionHopRouting.cu + ropeKvAppend.cu + sparseMlaDecode.cu) + +# The GLM-MoE-DSA kernels are written for the SM100 (datacenter Blackwell) +# family (for example 32-byte global loads), so pin the target to that family +# instead of inheriting the global architecture list. On builds without any +# sm_100-family architecture the target still compiles: device code that needs +# newer instructions is guarded by __CUDA_ARCH__ and traps. +set_cuda_architectures(glm_moe_dsa_kernels_src 100f) + +set_property(TARGET glm_moe_dsa_kernels_src PROPERTY POSITION_INDEPENDENT_CODE + ON) +set_property(TARGET glm_moe_dsa_kernels_src PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS + ON) diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopk.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopk.cu new file mode 100644 index 000000000000..d859f89960b4 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopk.cu @@ -0,0 +1,1557 @@ +/* + * 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/glmMoeDsa/crossRankTopk.h" + +#include + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa::cross_rank_topk +{ +namespace +{ + +constexpr int32_t kBins = 2048; +static_assert(kBins >= kTopK, "the output row is staged in the histogram"); +constexpr int32_t kWarps = kThreads / 32; +//! Per-warp key statistics kept in shared memory (see KeyStats). +constexpr int32_t kNumKeyStats = 5; +//! Tokens of each page one source holds. +constexpr int32_t kBlockTokens = kPageTokens / kSourcesPerRank; +//! Static shared memory of the kernel (round-2 segment tables), kept out of the dynamic budget. +constexpr size_t kStaticSmemBytes = 1024; + +// Scalar slots in shared memory. +enum ScalarSlot : int32_t +{ + kSlotBin = 0, + kSlotAbove, + kSlotBinCount, + kSlotTotal, + kSlotCounter2, + kSlotSrcCount, // [kMaxSources] + kSlotSrcBound = kSlotSrcCount + kMaxSources, // [kMaxSources] + kSlotSrcLen = kSlotSrcBound + kMaxSources, // [kMaxSources] + kSlotSrcOffset = kSlotSrcLen + kMaxSources, // [kMaxSources + 1] + kNumScalarSlots = kSlotSrcOffset + kMaxSources + 1 +}; + +__host__ __device__ inline int32_t numSources(int32_t world) +{ + return world * kSourcesPerRank; +} + +// Even share of kTopK per source. +__host__ __device__ inline int32_t evenShare(int32_t world) +{ + return (kTopK + numSources(world) - 1) / numSources(world); +} + +// Most keys a source sends in round 1 (an even number). +__host__ __device__ inline int32_t roundCap(int32_t world) +{ + return min(kSourceCapacity, 4 * evenShare(world)); +} + +// Keys one source can hold: one block of every local page. +__host__ __device__ inline int32_t chunkCapacity(int32_t localCapacity) +{ + int32_t const pages = (localCapacity + kPageTokens - 1) / kPageTokens; + return pages * kBlockTokens; +} + +__host__ __device__ inline int32_t mergeCapacity(int32_t world) +{ + return numSources(world) * roundCap(world); +} + +__host__ __device__ inline size_t smemBytes(int32_t localCapacity, int32_t world) +{ + size_t words = 2 * size_t{kBins} + 2 * size_t{kWarps} + size_t{kNumKeyStats} * kWarps + size_t{kNumScalarSlots}; + words = (words + 3) / 4 * 4; + words += static_cast(chunkCapacity(localCapacity)) + 2 * static_cast(mergeCapacity(world)); + return words * sizeof(uint32_t); +} + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 900) +// The kernel targets SM90+ (multicast Lamport exchange). +__global__ void __launch_bounds__(kThreads, 1) cpTopkKernel(CpTopkParams p) +{ + __trap(); +} +#else + +//! Bits of one radix-select digit (kBins bins). +constexpr int32_t kDigitBits = 11; +//! Order-preserving key of +0.0: the keys of non-negative values are at least this, those of negative values below. +//! Real indexer logits are mostly negative, and the cut of a row (and of many sources) often lies among small +//! negative values while the largest keys are positive: one key range across zero spans ~2^31 key units, where a +//! radix digit has no resolution. The selects therefore work on one sign at a time: with fewer than `need` keys at +//! or above kPositiveKey these are all taken and the cut is searched among the negative keys only, else among the +//! non-negative keys only; each range is narrow enough for one or two passes. +constexpr uint32_t kPositiveKey = 0x80000000u; +//! Independent 16-byte candidate loads one merge thread issues before it waits on any of them. +constexpr int32_t kPollBatch = 4; +//! Polls of a still-empty Lamport word between two reads of the global timer. +constexpr int32_t kPollsPerTimerRead = 256; + +// Keys a saturated source aims to send in round 1 (kSourceBudget at 8 ranks). +__device__ __forceinline__ int32_t roundNeed(int32_t world) +{ + return min(kSourceCapacity, 3 * evenShare(world)); +} + +// Tokens of `rank` among the global tokens [0, n). +__host__ __device__ __forceinline__ int32_t localLength(int32_t n, int32_t rank, int32_t world) +{ + if (n <= 0) + { + return 0; + } + int32_t const fullPages = n / kPageTokens; + int32_t const rem = n % kPageTokens; + int32_t const ownedFull = (fullPages - rank + world - 1) / world; + int32_t len = ownedFull * kPageTokens; + if (rem > 0 && fullPages % world == rank) + { + len += rem; + } + return len; +} + +// Source `src` of a rank holds one block of kBlockTokens consecutive tokens of every local page: block +// (src + page) % kSourcesPerRank. A fully selected page is thus spread over all sources of its rank. +// Keys of source `src` among `localLen` local columns (only the last local page can be partial). +__host__ __device__ __forceinline__ int32_t sourceLength(int32_t localLen, int32_t src) +{ + int32_t const pages = (localLen + kPageTokens - 1) / kPageTokens; + if (pages == 0) + { + return 0; + } + int32_t const last = pages - 1; + int32_t const tail = localLen - last * kPageTokens - ((src + last) % kSourcesPerRank) * kBlockTokens; + return last * kBlockTokens + min(max(tail, 0), kBlockTokens); +} + +// Global token of a local column under the page-interleaved partition. +__device__ __forceinline__ uint32_t localToGlobal(int32_t col, int32_t rank, int32_t world) +{ + int32_t const page = col / kPageTokens; + return static_cast((page * world + rank) * kPageTokens + (col % kPageTokens)); +} + +struct SmemLayout +{ + uint32_t* hist; // [2 * kBins] radix histograms of alternate passes; the first one stages the output row + uint32_t* warpSums; // [2 * kWarps] + uint32_t* stats; // [kNumKeyStats * kWarps] per-warp key statistics + uint32_t* scal; // [kNumScalarSlots] + uint32_t* chunk; // [chunkCapacity] order-preserving keys of this source + uint32_t* mkey; // [mergeCapacity] staged round-1 pairs, then the merge candidates (merge CTA) + uint32_t* midx; // [mergeCapacity] +}; + +__device__ inline SmemLayout carveSmem(uint32_t* base, int32_t localCapacity, int32_t world) +{ + SmemLayout s; + uint32_t* p = base; + s.hist = p; + p += 2 * kBins; + s.warpSums = p; + p += 2 * kWarps; + s.stats = p; + p += kNumKeyStats * kWarps; + s.scal = p; + p += kNumScalarSlots; + p = base + ((p - base) + 3) / 4 * 4; + s.chunk = p; + p += chunkCapacity(localCapacity); + s.mkey = p; + p += mergeCapacity(world); + s.midx = p; + return s; +} + +__device__ __forceinline__ uint32_t orderedKey(float v) +{ + uint32_t const b = __float_as_uint(v); + uint32_t const k = (b & 0x80000000u) ? ~b : (b | 0x80000000u); + // NaN keys above +inf are clamped so that no key, and no exclusive bound key + 1, equals the empty marker. + return k < kEmptyWord - 1u ? k : kEmptyWord - 2u; +} + +// Bits of x (0 for 0): the digit range of keys in [base, base + x]. +__device__ __forceinline__ int32_t bitLength(uint32_t x) +{ + return x == 0u ? 0 : 32 - __clz(x); +} + +// Statistics of a set of keys: extremes overall and per sign, and the count of non-negative keys. +struct KeyStats +{ + uint32_t maxKey = 0u; // largest key (0 when none) + uint32_t minKey = kEmptyWord; // smallest key (kEmptyWord when none) + uint32_t positives = 0u; // keys >= kPositiveKey + uint32_t maxNeg = 0u; // largest key below kPositiveKey (0 when none) + uint32_t minPos = kEmptyWord; // smallest key >= kPositiveKey (kEmptyWord when none) + + __device__ __forceinline__ void add(uint32_t k) + { + bool const pos = k >= kPositiveKey; + maxKey = max(maxKey, k); + minKey = min(minKey, k); + positives += pos ? 1u : 0u; + maxNeg = pos ? maxNeg : max(maxNeg, k); + minPos = pos ? min(minPos, k) : minPos; + } +}; + +// Reduce a thread's statistics over its warp; lane 0 stores them in the per-warp array (read with readStats after +// the next block barrier). +__device__ __forceinline__ void storeWarpStats(KeyStats const& t, uint32_t* stats) +{ + uint32_t const v0 = __reduce_max_sync(0xFFFFFFFFu, t.maxKey); + uint32_t const v1 = __reduce_min_sync(0xFFFFFFFFu, t.minKey); + uint32_t const v2 = __reduce_add_sync(0xFFFFFFFFu, t.positives); + uint32_t const v3 = __reduce_max_sync(0xFFFFFFFFu, t.maxNeg); + uint32_t const v4 = __reduce_min_sync(0xFFFFFFFFu, t.minPos); + if (threadIdx.x % 32 == 0) + { + int32_t const w = threadIdx.x / 32; + stats[w] = v0; + stats[kWarps + w] = v1; + stats[2 * kWarps + w] = v2; + stats[3 * kWarps + w] = v3; + stats[4 * kWarps + w] = v4; + } +} + +// Block statistics from the per-warp array (every warp reduces it; no barrier). +__device__ __forceinline__ KeyStats readStats(uint32_t const* stats) +{ + int32_t const lane = threadIdx.x % 32; + KeyStats b; + b.maxKey = __reduce_max_sync(0xFFFFFFFFu, stats[lane]); + b.minKey = __reduce_min_sync(0xFFFFFFFFu, stats[kWarps + lane]); + b.positives = __reduce_add_sync(0xFFFFFFFFu, stats[2 * kWarps + lane]); + b.maxNeg = __reduce_max_sync(0xFFFFFFFFu, stats[3 * kWarps + lane]); + b.minPos = __reduce_min_sync(0xFFFFFFFFu, stats[4 * kWarps + lane]); + return b; +} + +// Digit setup of a select of the need-th largest key: over the non-negative keys when there are at least `need` of +// them, else over the negative keys with every non-negative key counted as taken (`aboveHi`). +struct SignRange +{ + uint32_t lo; // smallest key of the range (also the digit base) + uint64_t hi; // exclusive end of the histogrammed range + int32_t bits; // digit range of the range + uint32_t aboveHi; // keys at or above hi, all taken + bool negative; // the cut is among the negative keys +}; + +__device__ __forceinline__ SignRange signRange(KeyStats const& st, uint32_t need) +{ + if (st.positives >= need || st.positives == 0u) + { + uint32_t const lo = st.positives >= need ? st.minPos : st.minKey; + return SignRange{lo, uint64_t{kEmptyWord}, bitLength(st.maxKey - lo), 0u, st.positives == 0u}; + } + return SignRange{st.minKey, uint64_t{kPositiveKey}, bitLength(st.maxNeg - st.minKey), st.positives, true}; +} + +__device__ __forceinline__ int64_t round1Offset(int32_t buf, int32_t src, int32_t row) +{ + return int64_t{buf} * kBufferWords + (int64_t{src} * kMaxRows + row) * kRound1SlotWords; +} + +__device__ __forceinline__ int64_t round2Offset(int32_t buf, int32_t src, int32_t row) +{ + return int64_t{buf} * kBufferWords + int64_t{kMaxSources} * kMaxRows * kRound1SlotWords + + (int64_t{src} * kMaxRows + row) * kRound2SlotWords; +} + +//! Word offset of the candidates of a slot (after the 16-byte header). +constexpr int32_t kSlotHeaderWords = 4; + +__device__ __forceinline__ uint2 loadVolatile2(uint32_t const* ptr) +{ + uint2 v; + asm volatile("ld.volatile.global.v2.u32 {%0, %1}, [%2];" : "=r"(v.x), "=r"(v.y) : "l"(ptr) : "memory"); + return v; +} + +__device__ __forceinline__ uint4 loadVolatile4(uint32_t const* ptr) +{ + uint4 v; + asm volatile("ld.volatile.global.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(ptr) + : "memory"); + return v; +} + +__device__ __forceinline__ void store2(uint32_t* ptr, uint32_t a, uint32_t b) +{ + asm volatile("st.global.v2.u32 [%0], {%1, %2};" ::"l"(ptr), "r"(a), "r"(b) : "memory"); +} + +__device__ __forceinline__ void store4(uint32_t* ptr, uint4 v) +{ + asm volatile("st.global.v4.u32 [%0], {%1, %2, %3, %4};" ::"l"(ptr), "r"(v.x), "r"(v.y), "r"(v.z), "r"(v.w) + : "memory"); +} + +__device__ __forceinline__ bool isEmpty(uint2 v) +{ + return v.x == kEmptyWord || v.y == kEmptyWord; +} + +// Whether the first pair of a 16-byte candidate vector, and the second one when `second`, were written. +__device__ __forceinline__ bool vectorReady(uint4 v, bool second) +{ + return v.x != kEmptyWord && v.y != kEmptyWord && (!second || (v.z != kEmptyWord && v.w != kEmptyWord)); +} + +__device__ __forceinline__ uint64_t globalTimerNs() +{ + uint64_t t; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t)); + return t; +} + +// What a polling thread needs to give up: the error word, the timeout and this rank (for the report). +struct WaitCtx +{ + int32_t* state; + uint64_t timeoutNs; + int32_t rank; +}; + +// A peer did not deliver within the timeout. The selection cannot be completed correctly any more, and the peer's +// late writes could land in a buffer this rank resets or reuses, so record which wait failed and abort the kernel: +// the resulting CUDA error fails the run instead of letting any rank continue with a partial selection. +__device__ __noinline__ void abortOnTimeout(WaitCtx const& w, int32_t what) +{ + bool const first = atomicExch(&w.state[3], what) == 0; + __threadfence_system(); + if (first) + { + printf( + "[glm_moe_dsa] crossRankTopk: rank %d, row %d: wait %d for a peer exceeded %llu ms; aborting the kernel\n", + w.rank, static_cast(blockIdx.y), what, static_cast(w.timeoutNs / 1000000ull)); + } + __trap(); +} + +// Spin until both words of a Lamport pair were written; aborts the kernel after the timeout. +__device__ uint2 pollPair(uint32_t const* ptr, WaitCtx const& w, int32_t what) +{ + uint2 v = loadVolatile2(ptr); + uint64_t start = 0; + for (int32_t polls = 1; isEmpty(v); ++polls) + { + if (polls % kPollsPerTimerRead == 0) + { + uint64_t const now = globalTimerNs(); + if (start == 0) + { + start = now; + } + else if (now - start > w.timeoutNs) + { + abortOnTimeout(w, what); + } + } + v = loadVolatile2(ptr); + } + return v; +} + +// Spin until a 16-byte candidate vector was written (see vectorReady); aborts the kernel after the timeout. +__device__ uint4 pollVector(uint32_t const* ptr, bool second, WaitCtx const& w, int32_t what) +{ + uint4 v = loadVolatile4(ptr); + uint64_t start = 0; + for (int32_t polls = 1; !vectorReady(v, second); ++polls) + { + if (polls % kPollsPerTimerRead == 0) + { + uint64_t const now = globalTimerNs(); + if (start == 0) + { + start = now; + } + else if (now - start > w.timeoutNs) + { + abortOnTimeout(w, what); + } + } + v = loadVolatile4(ptr); + } + return v; +} + +// Block-wide inclusive scan of one value per thread. Returns the inclusive prefix; `total` gets the block sum. +__device__ uint32_t blockInclusiveScan(uint32_t v, uint32_t* warpSums, uint32_t& total) +{ + int32_t const lane = threadIdx.x % 32; + int32_t const warp = threadIdx.x / 32; + uint32_t x = v; +#pragma unroll + for (int32_t o = 1; o < 32; o <<= 1) + { + uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, x, o); + if (lane >= o) + { + x += y; + } + } + if (lane == 31) + { + warpSums[warp] = x; + } + __syncthreads(); + if (warp == 0) + { + uint32_t w = warpSums[lane]; +#pragma unroll + for (int32_t o = 1; o < 32; o <<= 1) + { + uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, w, o); + if (lane >= o) + { + w += y; + } + } + warpSums[kWarps + lane] = w; + } + __syncthreads(); + total = warpSums[2 * kWarps - 1]; + uint32_t const before = warp > 0 ? warpSums[kWarps + warp - 1] : 0u; + return x + before; +} + +// Exclusive block-wide scan with a single barrier: every warp scans the per-warp sums itself. Returns this +// thread's exclusive prefix; `total` gets the block sum. warpSums must not be rewritten before the next barrier. +__device__ uint32_t blockExclusiveScan(uint32_t v, uint32_t* warpSums, uint32_t& total) +{ + int32_t const lane = threadIdx.x % 32; + int32_t const warp = threadIdx.x / 32; + uint32_t x = v; +#pragma unroll + for (int32_t o = 1; o < 32; o <<= 1) + { + uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, x, o); + if (lane >= o) + { + x += y; + } + } + if (lane == 31) + { + warpSums[warp] = x; + } + __syncthreads(); + uint32_t const w = warpSums[lane]; + uint32_t y = w; +#pragma unroll + for (int32_t o = 1; o < 32; o <<= 1) + { + uint32_t const z = __shfl_up_sync(0xFFFFFFFFu, y, o); + if (lane >= o) + { + y += z; + } + } + total = __shfl_sync(0xFFFFFFFFu, y, 31); + return __shfl_sync(0xFFFFFFFFu, y - w, warp) + x - v; +} + +// Find, counting from the top bin, the bin of `hist` that holds the target-th largest key, with two barriers: +// per-warp sums, then each warp scans them and only the warp holding the target scans its own bins. Returns the +// total count (every thread); when target <= total, the bin, the count above it and its own count go to the +// scalar slots. +__device__ uint32_t findBin(SmemLayout const& s, uint32_t const* hist, int32_t nbins, uint32_t target) +{ + int32_t const lane = threadIdx.x % 32; + int32_t const warp = threadIdx.x / 32; + int32_t const perThread = (nbins + kThreads - 1) / kThreads; + int32_t const top = nbins - 1 - perThread * static_cast(threadIdx.x); + uint32_t mine = 0; + for (int32_t i = 0; i < perThread && top - i >= 0; ++i) + { + mine += hist[top - i]; + } + uint32_t const warpTotal = __reduce_add_sync(0xFFFFFFFFu, mine); + if (lane == 0) + { + s.warpSums[warp] = warpTotal; + } + __syncthreads(); + uint32_t const w = s.warpSums[lane]; + uint32_t x = w; +#pragma unroll + for (int32_t o = 1; o < 32; o <<= 1) + { + uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, x, o); + if (lane >= o) + { + x += y; + } + } + uint32_t const total = __shfl_sync(0xFFFFFFFFu, x, 31); + uint32_t const warpBefore = __shfl_sync(0xFFFFFFFFu, x - w, warp); + if (warpBefore < target && target <= warpBefore + warpTotal) + { + uint32_t t = mine; +#pragma unroll + for (int32_t o = 1; o < 32; o <<= 1) + { + uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, t, o); + if (lane >= o) + { + t += y; + } + } + uint32_t acc = warpBefore + t - mine; + if (acc < target && target <= acc + mine) + { + for (int32_t i = 0; i < perThread && top - i >= 0; ++i) + { + uint32_t const c = hist[top - i]; + if (acc + c >= target) + { + s.scal[kSlotBin] = static_cast(top - i); + s.scal[kSlotAbove] = acc; + s.scal[kSlotBinCount] = c; + break; + } + acc += c; + } + } + } + __syncthreads(); + return total; +} + +struct SelectOut +{ + uint32_t edge; // lower edge of the selected bin (the exact key after the last digit) + uint32_t above; // predicate keys strictly above the selected bin + uint32_t binCount; // predicate keys in the selected bin + uint32_t total; // predicate keys + bool all; // take every predicate key (fewer than needed, or all fit the capacity) + bool full; // exact mode: the keys at or above `edge` are exactly the `need` largest + int32_t passes; // radix passes run +}; + +// Candidate keys and token indices held in shared memory (16-byte aligned, readable in groups of four up to the +// next multiple of four). +struct SmemKeys +{ + static constexpr bool kVector = true; + uint32_t const* keys; + uint32_t const* idx; + + __device__ uint4 keys4(int32_t j) const + { + return *reinterpret_cast(keys + j); + } + + __device__ uint4 index4(int32_t j) const + { + return *reinterpret_cast(idx + j); + } +}; + +// Merge candidates in shared memory followed by round-2 candidates read from their global Lamport slots. +struct MergeKeys +{ + static constexpr bool kVector = false; + uint32_t const* smemKeys; + uint32_t const* smemIdx; + int32_t smemCount; + uint32_t const* const* slots; // round-2 slot of each segment + uint32_t const* segBegin; // [numSegments + 1] candidate offsets + int32_t numSegments; + + __device__ uint32_t const* pair(int32_t k) const + { + // Last segment whose begin is <= k. + int32_t lo = 0; + int32_t hi = numSegments - 1; + while (lo < hi) + { + int32_t const mid = (lo + hi + 1) / 2; + if (static_cast(segBegin[mid]) <= k) + { + lo = mid; + } + else + { + hi = mid - 1; + } + } + return slots[lo] + kSlotHeaderWords + 2 * (k - static_cast(segBegin[lo])); + } + + __device__ uint32_t operator()(int32_t j) const + { + return j < smemCount ? smemKeys[j] : pair(j - smemCount)[0]; + } + + __device__ uint32_t index(int32_t j) const + { + return j < smemCount ? smemIdx[j] : pair(j - smemCount)[1]; + } +}; + +// Debug slots of this CTA starting at `offset` (nullptr without a debug buffer). +__device__ __forceinline__ int64_t* debugSlots(CpTopkParams const& p, int32_t offset) +{ + if (p.debugTimes == nullptr) + { + return nullptr; + } + return p.debugTimes + (static_cast(blockIdx.y) * kSourcesPerRank + blockIdx.x) * kDebugSlots + offset; +} + +// clock64 of thread 0 into a debug slot (no-op without a debug buffer). +__device__ __forceinline__ void stamp(int64_t* dbg, int32_t slot) +{ + if (dbg != nullptr && threadIdx.x == 0) + { + dbg[slot] = clock64(); + } +} + +// A value into a debug slot (thread 0; no-op without a debug buffer). +__device__ __forceinline__ void putDebug(int64_t* dbg, int32_t slot, int64_t value) +{ + if (dbg != nullptr && threadIdx.x == 0) + { + dbg[slot] = value; + } +} + +// Phase timestamp for diagnosis (thread 0 of the CTA; no-op without a debug buffer). +__device__ __forceinline__ void markPhase(CpTopkParams const& p, int32_t phase) +{ + if (p.debugTimes != nullptr && threadIdx.x == 0) + { + uint64_t t; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t)); + int64_t* d = p.debugTimes + (static_cast(blockIdx.y) * kSourcesPerRank + blockIdx.x) * kDebugSlots; + d[phase] = static_cast(t); + d[16 + phase] = static_cast(clock64()); + } +} + +// Histogram update of one key of a radix pass. +__device__ __forceinline__ void histogramKey(uint32_t* hist, bool inRange, uint32_t k, uint32_t lo, uint64_t hi, + uint32_t base, uint32_t prefMask, uint32_t prefVal, int32_t shift, uint32_t binMask) +{ + uint32_t const d = k - base; + if (inRange && k >= lo && static_cast(k) < hi && (d & prefMask) == prefVal) + { + atomicAdd(&hist[(d >> shift) & binMask], 1u); + } +} + +// Block radix select of the need-th largest of the n keys at or above lo, of which `aboveHi` (< need) lie at or +// above hi and are counted as taken without being histogrammed. Every key in [lo, hi) is at least `base`, and +// key - base < 2^topBits; digits of up to kDigitBits bits of key - base are refined from the top. With cap > 0 the +// select stops at the first bin edge whose keys at or above it number at most `cap`. With cap = 0 it runs to the +// exact key, or stops early when the selected bin is wholly taken (`full`). Alternate passes use alternate +// histograms, so the next one is cleared while the current one is searched. Counts in the result include aboveHi. +template +__device__ SelectOut blockSelect(SmemLayout const& s, KeyAt const& keyAt, int32_t n, uint32_t lo, uint64_t hi, + uint32_t need, uint32_t cap, uint32_t base, int32_t topBits, int64_t* dbg = nullptr, uint32_t aboveHi = 0u) +{ + if (dbg != nullptr && threadIdx.x == 0) + { + dbg[0] = n; + } + SelectOut out{lo, 0u, 0u, 0u, false, false, 0}; + uint32_t prefMask = 0; + uint32_t prefVal = 0; + uint32_t above = aboveHi; + int32_t remaining = max(topBits, 1); + uint32_t* hist = s.hist; + uint32_t* histNext = s.hist + kBins; + for (int32_t b = threadIdx.x; b < (1 << min(kDigitBits, remaining)); b += kThreads) + { + hist[b] = 0; + } + __syncthreads(); +#pragma unroll 1 + for (int32_t pass = 0; remaining > 0; ++pass) + { + int32_t const width = min(kDigitBits, remaining); + int32_t const shift = remaining - width; + int32_t const nbins = 1 << width; + uint32_t const binMask = static_cast(nbins - 1); + out.passes = pass + 1; + if (dbg != nullptr && threadIdx.x == 0 && pass < 3) + { + dbg[1 + 3 * pass] = clock64(); + } + if constexpr (KeyAt::kVector) + { + for (int32_t j0 = 4 * static_cast(threadIdx.x); j0 < n; j0 += 4 * kThreads) + { + uint4 const q = keyAt.keys4(j0); + histogramKey(hist, true, q.x, lo, hi, base, prefMask, prefVal, shift, binMask); + histogramKey(hist, j0 + 1 < n, q.y, lo, hi, base, prefMask, prefVal, shift, binMask); + histogramKey(hist, j0 + 2 < n, q.z, lo, hi, base, prefMask, prefVal, shift, binMask); + histogramKey(hist, j0 + 3 < n, q.w, lo, hi, base, prefMask, prefVal, shift, binMask); + } + } + else + { + for (int32_t j = threadIdx.x; j < n; j += kThreads) + { + histogramKey(hist, true, keyAt(j), lo, hi, base, prefMask, prefVal, shift, binMask); + } + } + __syncthreads(); + if (dbg != nullptr && threadIdx.x == 0 && pass < 3) + { + dbg[2 + 3 * pass] = clock64(); + } + if (shift > 0) + { + for (int32_t b = threadIdx.x; b < (1 << min(kDigitBits, shift)); b += kThreads) + { + histNext[b] = 0; + } + } + // Find the bin holding the (need - above)-th largest key, counting from the top bin. + uint32_t const target = need - above; + uint32_t const total = findBin(s, hist, nbins, target); + if (dbg != nullptr && threadIdx.x == 0 && pass < 3) + { + dbg[3 + 3 * pass] = clock64(); + } + if (pass == 0) + { + out.total = total + aboveHi; + if (total + aboveHi <= cap || target > total) + { + out.all = true; + out.above = total + aboveHi; + return out; + } + } + // The scalars are next written after two more barriers (the next histogram and scan). + uint32_t const bin = s.scal[kSlotBin]; + above += s.scal[kSlotAbove]; + out.binCount = s.scal[kSlotBinCount]; + prefVal |= bin << shift; + prefMask |= binMask << shift; + out.edge = base + prefVal; + out.above = above; + remaining = shift; + uint32_t* const t = hist; + hist = histNext; + histNext = t; + if (cap > 0 && above + out.binCount <= cap) + { + return out; + } + if (cap == 0 && remaining > 0 && above + out.binCount == need) + { + out.full = true; + return out; + } + } + return out; +} + +// Send one source's selected keys (keys in [lo, hi) of the chunk) to its slot; returns the count sent. +// Takes every predicate key at or above the edge when they fit `cap`. Otherwise (more than `cap` keys tie at the +// exact edge) it takes every key above the edge, plus, with kPartialTies, the first keys equal to it in source +// order up to `cap` (any subset of equal keys is a valid tie-aware choice). Round 1 sends no partial ties, so the +// keys a second round requests from a source (at most the edge) were never sent before. +// The header carries the count and an exclusive upper bound of the keys left unsent (0 when none). +// With `stage` (round 1, at most kSourceCapacity keys), the pairs are gathered in shared memory first and written +// with coalesced 16-byte stores; otherwise each pair is stored where it is found. +template +__device__ uint32_t sendSource(SmemLayout const& s, SelectOut const& sel, int32_t n, uint32_t lo, uint64_t hi, + uint32_t cap, uint32_t* mcSlot, int32_t src, int32_t rank, int32_t world, uint2* stage, int64_t* dbg = nullptr) +{ + bool const strict = !sel.all && sel.above + sel.binCount > cap; + uint32_t const threshold = sel.all ? lo : sel.edge; + // Keys equal to the edge that fit after the keys above it (sel.above < need <= cap when strict). + uint32_t const tieBudget = kPartialTies && strict ? cap - sel.above : 0u; + uint32_t const sent = sel.all ? sel.total : (strict ? sel.above + tieBudget : sel.above + sel.binCount); + uint32_t const bound = sel.all ? lo : (strict ? sel.edge + 1u : sel.edge); + stamp(dbg, 42); + // One page block (kBlockTokens keys) per thread; positions from a block scan, so the slot holds the keys in + // source order: first those taken outright, then (strict only) the admitted keys equal to the edge. + int32_t const groups = (n + kBlockTokens - 1) / kBlockTokens; + uint32_t takenMain = 0; + uint32_t takenTies = 0; + uint32_t const tieBase = strict ? sel.above : 0u; + for (int32_t g0 = 0; g0 < groups; g0 += kThreads) + { + int32_t const g = g0 + static_cast(threadIdx.x); + uint32_t k[kBlockTokens]; + uint32_t mask = 0; + uint32_t tieMask = 0; + if (g < groups) + { + uint4 const a = *reinterpret_cast(s.chunk + kBlockTokens * g); + uint4 const b = *reinterpret_cast(s.chunk + kBlockTokens * g + 4); + k[0] = a.x; + k[1] = a.y; + k[2] = a.z; + k[3] = a.w; + k[4] = b.x; + k[5] = b.y; + k[6] = b.z; + k[7] = b.w; +#pragma unroll + for (int32_t i = 0; i < kBlockTokens; ++i) + { + bool const in = kBlockTokens * g + i < n && k[i] >= lo && static_cast(k[i]) < hi; + bool const take = in && (strict ? k[i] > threshold : k[i] >= threshold); + bool const tie = kPartialTies && in && strict && k[i] == threshold; + mask |= (take ? 1u : 0u) << i; + tieMask |= (tie ? 1u : 0u) << i; + } + } + // At most kThreads * kBlockTokens keys per pass: both counts fit 16 bits. + uint32_t const c = static_cast(__popc(mask)) | (static_cast(__popc(tieMask)) << 16); + uint32_t chunkTaken; + uint32_t const before = blockExclusiveScan(c, s.warpSums, chunkTaken); + uint32_t pos = takenMain + (before & 0xFFFFu); + uint32_t tiePos = takenTies + (before >> 16); + takenMain += chunkTaken & 0xFFFFu; + takenTies += chunkTaken >> 16; + if ((mask | tieMask) != 0u) + { + int32_t const colBase = g * kPageTokens + ((src + g) % kSourcesPerRank) * kBlockTokens; +#pragma unroll + for (int32_t i = 0; i < kBlockTokens; ++i) + { + uint32_t slot = kEmptyWord; + if ((mask >> i) & 1u) + { + slot = pos++; + } + else if constexpr (kPartialTies) + { + if ((tieMask >> i) & 1u) + { + slot = tiePos < tieBudget ? tieBase + tiePos : kEmptyWord; + ++tiePos; + } + } + if (slot != kEmptyWord) + { + uint32_t const token = localToGlobal(colBase + i, rank, world); + if (stage != nullptr) + { + stage[slot] = make_uint2(k[i], token); + } + else + { + store2(mcSlot + kSlotHeaderWords + 2 * static_cast(slot), k[i], token); + } + } + } + } + // The staged pairs are complete, and the scan's warp sums may be rewritten. + __syncthreads(); + } + stamp(dbg, 43); + if (stage != nullptr) + { + // Two pairs per 16-byte store; an odd count leaves the last half of the last vector empty. + int32_t const vectors = static_cast(sent + 1u) / 2; + for (int32_t v = threadIdx.x; v < vectors; v += kThreads) + { + uint2 const a = stage[2 * v]; + uint2 const b + = 2 * v + 1 < static_cast(sent) ? stage[2 * v + 1] : make_uint2(kEmptyWord, kEmptyWord); + store4(mcSlot + kSlotHeaderWords + 4 * static_cast(v), make_uint4(a.x, a.y, b.x, b.y)); + } + } + // Receivers check every candidate word, so the header needs no ordering against the candidates. + if (threadIdx.x == kThreads - 1) + { + store2(mcSlot, sent, bound); + } + stamp(dbg, 44); + return sent; +} + +// Load one source's keys (one block of kBlockTokens tokens of every local page) from the local logits into the +// chunk; returns the number of keys. Consecutive threads read consecutive keys (whole 32-byte blocks per eight +// threads), and each thread issues kLoadBatch loads before it uses any. With `withStats`, the per-warp key +// statistics are stored for readStats after the next block barrier. +__device__ int32_t loadSource( + SmemLayout const& s, float const* rowLogits, int32_t localLen, int32_t src, bool withStats) +{ + constexpr int32_t kLoadBatch = 8; + int32_t const n = sourceLength(localLen, src); + KeyStats st; + for (int32_t j0 = threadIdx.x; j0 < n; j0 += kLoadBatch * kThreads) + { + float v[kLoadBatch]; +#pragma unroll + for (int32_t i = 0; i < kLoadBatch; ++i) + { + int32_t const j = j0 + i * kThreads; + int32_t const page = j / kBlockTokens; + int32_t const col = page * kPageTokens + ((src + page) % kSourcesPerRank) * kBlockTokens + j % kBlockTokens; + v[i] = j < n ? rowLogits[col] : 0.F; + } +#pragma unroll + for (int32_t i = 0; i < kLoadBatch; ++i) + { + int32_t const j = j0 + i * kThreads; + if (j < n) + { + uint32_t const k = orderedKey(v[i]); + s.chunk[j] = k; + st.add(k); + } + } + } + if (withStats) + { + storeWarpStats(st, s.stats); + } + return n; +} + +// Whether source q withheld keys that may lie above the merged k-th key. +__device__ __forceinline__ bool sourceOverflows(SmemLayout const& s, int32_t q, uint32_t kth) +{ + uint32_t const cnt = s.scal[kSlotSrcCount + q]; + uint64_t const bound = s.scal[kSlotSrcBound + q]; + return cnt < s.scal[kSlotSrcLen + q] && bound > uint64_t{kth} + 1u; +} + +// Emit the top-kTopK of the merge candidates. Keys above `kth` are all taken; ties at `kth` are taken in candidate +// order when only some of them fit (deterministic on every rank), otherwise all of them. With `takeAll`, every +// candidate is taken and the row is padded with -1. The row is assembled in shared memory (the first histogram) +// and written with coalesced stores. +template +__device__ bool emitSelection(SmemLayout const& s, KeyAt const& keys, int32_t total, uint32_t kth, bool takeAll, + int32_t* outRow, int32_t overflowSources, int64_t* dbg = nullptr) +{ + uint32_t const want = static_cast(kTopK); + // Positions from one block scan over contiguous candidate ranges (no contended counter); the order is the + // candidate order, identical on every rank. Shared-memory candidates are read four at a time. + int32_t per = (total + kThreads - 1) / kThreads; + if constexpr (KeyAt::kVector) + { + per = (per + 3) / 4 * 4; + } + int32_t const j0 = min(total, static_cast(threadIdx.x) * per); + int32_t const j1 = min(total, j0 + per); + uint32_t cGt = 0; + uint32_t cEq = 0; + if constexpr (KeyAt::kVector) + { + for (int32_t j = j0; j < j1; j += 4) + { + uint4 const q = keys.keys4(j); + uint32_t const kk[4] = {q.x, q.y, q.z, q.w}; +#pragma unroll + for (int32_t i = 0; i < 4; ++i) + { + bool const in = j + i < j1; + cGt += (in && (takeAll || kk[i] > kth)) ? 1u : 0u; + cEq += (in && !takeAll && kk[i] == kth) ? 1u : 0u; + } + } + } + else + { + for (int32_t j = j0; j < j1; ++j) + { + uint32_t const k = keys(j); + cGt += (takeAll || k > kth) ? 1u : 0u; + cEq += (!takeAll && k == kth) ? 1u : 0u; + } + } + // The first `overflowSources` threads flag a source that withheld keys above the threshold (the flag slot is + // zero on entry); with a flag set, nothing is written and the caller runs the second round. + if (static_cast(threadIdx.x) < overflowSources && sourceOverflows(s, threadIdx.x, kth)) + { + atomicOr(&s.scal[kSlotCounter2], 1u); + } + uint32_t allGt; + uint32_t allEq; + uint32_t posGt; + uint32_t posEq; + if (total < 0x10000) + { + // Both counts fit 16 bits: one scan of the packed pair. + uint32_t packedTotal; + uint32_t const packed = cGt | (cEq << 16); + uint32_t const before = blockExclusiveScan(packed, s.warpSums, packedTotal); + allGt = packedTotal & 0xFFFFu; + allEq = packedTotal >> 16; + posGt = before & 0xFFFFu; + posEq = allGt + (before >> 16); + } + else + { + posGt = blockInclusiveScan(cGt, s.warpSums, allGt) - cGt; + __syncthreads(); + posEq = blockInclusiveScan(cEq, s.warpSums, allEq) - cEq; + posEq += allGt; + } + stamp(dbg, 59); + if (overflowSources > 0 && s.scal[kSlotCounter2] != 0u) + { + return true; + } + uint32_t* row = s.hist; + auto place = [&](uint32_t k, uint32_t idx) + { + if (takeAll || k > kth) + { + if (posGt < want) + { + row[posGt] = idx; + } + ++posGt; + } + else if (k == kth) + { + if (posEq < want) + { + row[posEq] = idx; + } + ++posEq; + } + }; + if constexpr (KeyAt::kVector) + { + for (int32_t j = j0; j < j1; j += 4) + { + uint4 const q = keys.keys4(j); + uint4 const x = keys.index4(j); + uint32_t const kk[4] = {q.x, q.y, q.z, q.w}; + uint32_t const xx[4] = {x.x, x.y, x.z, x.w}; +#pragma unroll + for (int32_t i = 0; i < 4; ++i) + { + if (j + i < j1) + { + place(kk[i], xx[i]); + } + } + } + } + else + { + for (int32_t j = j0; j < j1; ++j) + { + place(keys(j), keys.index(j)); + } + } + __syncthreads(); + stamp(dbg, 60); + uint32_t const selected = min(allGt + allEq, want); + for (uint32_t j = threadIdx.x; j < want; j += kThreads) + { + outRow[j] = j < selected ? static_cast(row[j]) : -1; + } + stamp(dbg, 61); + return false; +} + +// Threshold of an exact (cap = 0) merge select: keys above `kth` are taken, ties at `kth` while they fit. +__device__ __forceinline__ uint32_t mergeThreshold(SelectOut const& m, bool& takeAll) +{ + takeAll = m.all || (m.full && m.edge == 0u); + if (takeAll) + { + return 0u; + } + // A wholly taken bin: the keys at or above its edge are exactly kTopK. + return m.full ? m.edge - 1u : m.edge; +} + +// Reset the slots of Lamport buffer `buf` on this rank's copy, using one warp per slot. A slot was written iff its +// header was; the header's count bounds the candidates to reset. +__device__ void clearBuffer(CpTopkParams const& p, int32_t buf, int32_t clearer, int32_t numClearers) +{ + int32_t const warp = threadIdx.x / 32; + int32_t const lane = threadIdx.x % 32; + int32_t const round1Slots = kMaxSources * kMaxRows; + uint4 const empty = make_uint4(kEmptyWord, kEmptyWord, kEmptyWord, kEmptyWord); + for (int32_t slot = clearer * kWarps + warp; slot < 2 * round1Slots; slot += numClearers * kWarps) + { + bool const second = slot >= round1Slots; + int32_t const qr = second ? slot - round1Slots : slot; + uint32_t* base = p.ucWorkspace + + (second ? round2Offset(buf, qr / kMaxRows, qr % kMaxRows) + : round1Offset(buf, qr / kMaxRows, qr % kMaxRows)); + uint2 h = make_uint2(kEmptyWord, kEmptyWord); + if (lane == 0) + { + h = loadVolatile2(base); + } + h.x = __shfl_sync(0xFFFFFFFFu, h.x, 0); + h.y = __shfl_sync(0xFFFFFFFFu, h.y, 0); + if (h.x == kEmptyWord && h.y == kEmptyWord) + { + continue; + } + uint32_t const capacity = second ? kRound2Capacity : kSourceCapacity; + uint32_t const count = h.x == kEmptyWord ? capacity : min(h.x, capacity); + int32_t const vectors = static_cast(count + 1u) / 2; + for (int32_t v = lane; v < vectors; v += 32) + { + store4(base + kSlotHeaderWords + 4 * static_cast(v), empty); + } + if (lane == 0) + { + store2(base, kEmptyWord, kEmptyWord); + } + } +} + +// Merge of one row on this rank (one CTA): gathers every source's round-1 candidates, selects the exact +// top-kTopK, runs the second round when a source withheld keys above the threshold, and writes the row. +__device__ void mergeRow(CpTopkParams const& p, SmemLayout const& s, uint32_t const** segSlots, uint32_t* segBegin, + int32_t buf, int32_t row, int32_t rowLen, int32_t localLen, float const* rowLogits, int32_t* outRow, + WaitCtx const& wait) +{ + int32_t const world = p.worldSize; + int32_t const rank = p.rank; + int32_t const sources = numSources(world); + uint32_t const cap = static_cast(roundCap(world)); + __syncthreads(); + if (threadIdx.x == 0) + { + s.scal[kSlotCounter2] = 0u; + } + if (static_cast(threadIdx.x) < sources) + { + // Keys each source holds, computed while its header is in flight. + int32_t const q = threadIdx.x; + s.scal[kSlotSrcLen + q] = static_cast( + sourceLength(min(localLength(rowLen, q / kSourcesPerRank, world), p.localCapacity), q % kSourcesPerRank)); + uint2 const h = pollPair(p.ucWorkspace + round1Offset(buf, threadIdx.x, row), wait, kWaitRound1Header); + s.scal[kSlotSrcCount + threadIdx.x] = min(h.x, cap); + s.scal[kSlotSrcBound + threadIdx.x] = h.y; + } + __syncthreads(); + if (threadIdx.x < 32) + { + // Exclusive prefix of the source counts (at most 64 sources: two per lane). + uint32_t const a = static_cast(threadIdx.x) < sources ? s.scal[kSlotSrcCount + threadIdx.x] : 0u; + uint32_t const b + = static_cast(threadIdx.x) + 32 < sources ? s.scal[kSlotSrcCount + threadIdx.x + 32] : 0u; + uint32_t x = a; + uint32_t z = b; +#pragma unroll + for (int32_t o = 1; o < 32; o <<= 1) + { + uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, x, o); + uint32_t const w = __shfl_up_sync(0xFFFFFFFFu, z, o); + if (static_cast(threadIdx.x) >= o) + { + x += y; + z += w; + } + } + uint32_t const totalA = __shfl_sync(0xFFFFFFFFu, x, 31); + s.scal[kSlotSrcOffset + threadIdx.x] = x - a; + s.scal[kSlotSrcOffset + 32 + threadIdx.x] = totalA + z - b; + if (threadIdx.x == 31) + { + s.scal[kSlotSrcOffset + kMaxSources] = totalA + z; + } + } + __syncthreads(); + int32_t const total = static_cast(s.scal[kSlotSrcOffset + kMaxSources]); + markPhase(p, 4); + + // Candidates: padded (source, vector) positions; issue a batch of independent 16-byte loads, then re-poll the + // ones not yet written. The key statistics are tracked for the merge select. + int32_t const vectorsPerSource = static_cast(cap) / 2; + int32_t const padded = sources * vectorsPerSource; + KeyStats st; +#pragma unroll 1 + for (int32_t v0 = 0; v0 < padded; v0 += kThreads * kPollBatch) + { + uint4 v[kPollBatch]; +#pragma unroll + for (int32_t i = 0; i < kPollBatch; ++i) + { + int32_t const jv = v0 + i * kThreads + static_cast(threadIdx.x); + int32_t const q = jv / vectorsPerSource; + int32_t const kv = jv % vectorsPerSource; + v[i] = make_uint4(0u, 0u, 0u, 0u); + if (jv < padded && 2 * kv < static_cast(s.scal[kSlotSrcCount + q])) + { + v[i] = loadVolatile4(p.ucWorkspace + round1Offset(buf, q, row) + kSlotHeaderWords + 4 * kv); + } + } +#pragma unroll + for (int32_t i = 0; i < kPollBatch; ++i) + { + int32_t const jv = v0 + i * kThreads + static_cast(threadIdx.x); + int32_t const q = jv / vectorsPerSource; + int32_t const kv = jv % vectorsPerSource; + if (jv < padded && 2 * kv < static_cast(s.scal[kSlotSrcCount + q])) + { + bool const second = 2 * kv + 1 < static_cast(s.scal[kSlotSrcCount + q]); + if (!vectorReady(v[i], second)) + { + v[i] = pollVector(p.ucWorkspace + round1Offset(buf, q, row) + kSlotHeaderWords + 4 * kv, second, + wait, kWaitRound1Candidates); + } + int32_t const j = static_cast(s.scal[kSlotSrcOffset + q]) + 2 * kv; + s.mkey[j] = v[i].x; + s.midx[j] = v[i].y; + st.add(v[i].x); + if (second) + { + s.mkey[j + 1] = v[i].z; + s.midx[j + 1] = v[i].w; + st.add(v[i].z); + } + } + } + } + stamp(debugSlots(p, 0), 45); + storeWarpStats(st, s.stats); + __syncthreads(); + markPhase(p, 5); + + // Digits of the merge select span the candidates of the sign the cut lies in (see kPositiveKey). + SignRange const r = total > 0 ? signRange(readStats(s.stats), static_cast(kTopK)) + : SignRange{0u, uint64_t{kEmptyWord}, 32, 0u, false}; + SmemKeys const merged{s.mkey, s.midx}; + SelectOut const m = blockSelect( + s, merged, total, r.lo, r.hi, static_cast(kTopK), 0u, r.lo, r.bits, debugSlots(p, 48), r.aboveHi); + putDebug(debugSlots(p, 0), 62, total); + putDebug(debugSlots(p, 0), 63, m.passes + (r.negative ? 10 : 0)); + bool takeAll = false; + uint32_t const kth = mergeThreshold(m, takeAll); + markPhase(p, 6); + stamp(debugSlots(p, 0), 58); + if (!emitSelection(s, merged, total, kth, takeAll, outRow, sources, debugSlots(p, 0))) + { + return; + } + + // ---- Round 2 (rare; every rank takes it): withheld keys in [kth, bound) of overflowing sources ---------- + __syncthreads(); + // 1. This rank's overflowing sources send their top-kRound2Capacity withheld keys. + for (int32_t c = 0; c < kSourcesPerRank; ++c) + { + int32_t const q = rank * kSourcesPerRank + c; + if (!sourceOverflows(s, q, kth)) + { + continue; + } + uint64_t const b = s.scal[kSlotSrcBound + q]; + int32_t const nq = loadSource(s, rowLogits, localLen, c, false); + __syncthreads(); + SelectOut const sel2 = blockSelect( + s, SmemKeys{s.chunk, nullptr}, nq, kth, b, static_cast(kRound2Capacity), kRound2Capacity, 0u, 32); + (void) sendSource( + s, sel2, nq, kth, b, kRound2Capacity, p.mcWorkspace + round2Offset(buf, q, row), c, rank, world, nullptr); + __syncthreads(); + } + // 2. Gather the round-2 segments of every overflowing source (all ranks). + if (threadIdx.x == 0) + { + int32_t segs = 0; + uint32_t off = 0; + for (int32_t q = 0; q < sources; ++q) + { + if (!sourceOverflows(s, q, kth)) + { + continue; + } + uint32_t const* slot = p.ucWorkspace + round2Offset(buf, q, row); + uint2 const h = pollPair(slot, wait, kWaitRound2Header); + segSlots[segs] = slot; + segBegin[segs] = off; + off += min(h.x, static_cast(kRound2Capacity)); + ++segs; + } + segBegin[segs] = off; + s.scal[kSlotTotal] = static_cast(segs); + } + __syncthreads(); + int32_t const segs = static_cast(s.scal[kSlotTotal]); + int32_t const total2 = static_cast(segBegin[segs]); + MergeKeys const keys{s.mkey, s.midx, total, segSlots, segBegin, segs}; + for (int32_t k = threadIdx.x; k < total2; k += kThreads) + { + (void) pollPair(keys.pair(k), wait, kWaitRound2Candidates); + } + __syncthreads(); + SelectOut const m2 + = blockSelect(s, keys, total + total2, 0u, uint64_t{kEmptyWord}, static_cast(kTopK), 0u, 0u, 32); + bool takeAll2 = false; + uint32_t const kth2 = mergeThreshold(m2, takeAll2); + (void) emitSelection(s, keys, total + total2, kth2, takeAll2, outRow, 0); + if (threadIdx.x == 0) + { + atomicAdd(&p.state[2], 1); + } +} + +__global__ void __launch_bounds__(kThreads, 1) cpTopkKernel(CpTopkParams p) +{ + markPhase(p, 0); + extern __shared__ __align__(16) uint32_t smemRaw[]; + __shared__ uint32_t const* segSlots[kMaxSources]; + __shared__ uint32_t segBegin[kMaxSources + 1]; + int32_t const world = p.worldSize; + int32_t const rank = p.rank; + int32_t const src = blockIdx.x; + int32_t const row = blockIdx.y; + SmemLayout const s = carveSmem(smemRaw, p.localCapacity, world); + int32_t* outRow = p.out + static_cast(row) * kTopK; + // Every global read below (lengths, epoch, logits) comes after the wait, so a programmatic early launch is safe + // whatever the preceding kernel writes. + asm volatile("griddepcontrol.wait;" ::: "memory"); + int32_t const kvLen = p.kvLens[0]; + int32_t const epoch = *reinterpret_cast(&p.state[0]); + int32_t const buf = epoch % kNumLamportBuffers; + int32_t const rowLen = max(0, kvLen - p.nextN + row + 1); + // Identity launch (every row has at most kTopK tokens; the same decision on every rank): no exchange, and + // neither the workspace nor the epoch is touched, so identity launches never shorten the reuse chain of the + // Lamport buffers (see crossRankTopk.h). + bool const exchange = kvLen - p.nextN + p.numRows > kTopK; + __syncthreads(); + + if (rowLen <= kTopK) + { + // Every rank agrees on this branch: identity selection, no exchange for this row. + if (src == 0) + { + for (int32_t j = threadIdx.x; j < kTopK; j += kThreads) + { + outRow[j] = j < rowLen ? j : -1; + } + } + } + else + { + int32_t const localLen = min(localLength(rowLen, rank, world), p.localCapacity); + float const* rowLogits = p.logits + static_cast(row) * p.logitsStride; + + // ---- Round 1: this source's local top keys --------------------------------------------------------- + int32_t const n = loadSource(s, rowLogits, localLen, src, true); + __syncthreads(); + markPhase(p, 1); + uint32_t const cap = static_cast(roundCap(world)); + uint32_t const need = static_cast(roundNeed(world)); + SelectOut sel{0u, 0u, 0u, static_cast(n), true, false, 0}; + uint32_t selLo = 0u; + int32_t passInfo = 0; + if (n > static_cast(cap)) + { + // The cut lies among the keys of one sign (see kPositiveKey); the digits span that sign's key range. + SignRange const r = signRange(readStats(s.stats), need); + selLo = r.lo; + SmemKeys const keys{s.chunk, nullptr}; + sel = blockSelect(s, keys, n, r.lo, r.hi, need, cap, r.lo, r.bits, debugSlots(p, 32), r.aboveHi); + passInfo = sel.passes + (r.negative ? 10 : 0); + } + markPhase(p, 2); + uint32_t const sent = sendSource(s, sel, n, selLo, uint64_t{kEmptyWord}, cap, + p.mcWorkspace + round1Offset(buf, rank * kSourcesPerRank + src, row), src, rank, world, + reinterpret_cast(s.mkey), debugSlots(p, 0)); + putDebug(debugSlots(p, 0), 46, sent); + putDebug(debugSlots(p, 0), 47, passInfo); + markPhase(p, 3); + + // ---- Merge: one CTA per row and rank --------------------------------------------------------------- + if (src == 0) + { + WaitCtx const wait{p.state, p.timeoutNs, rank}; + mergeRow(p, s, segSlots, segBegin, buf, row, rowLen, localLen, rowLogits, outRow, wait); + markPhase(p, 7); + } + } + if (!exchange) + { + return; + } + + // The other CTAs reset the buffer of the previous exchange launch, whose candidates every rank has consumed: a + // peer writes it again only two exchange launches later, after this launch completed on this rank. + if (src != 0) + { + clearBuffer(p, (epoch + kNumLamportBuffers - 1) % kNumLamportBuffers, row * (kSourcesPerRank - 1) + src - 1, + static_cast(gridDim.y) * (kSourcesPerRank - 1)); + } + + // Advance the epoch once every CTA of this launch has read it. + __syncthreads(); + markPhase(p, 8); + if (threadIdx.x == 0) + { + __threadfence(); + int32_t const numCtas = static_cast(gridDim.x * gridDim.y); + int32_t const arrived = atomicAdd(&p.state[1], 1); + if (arrived == numCtas - 1) + { + p.state[1] = 0; + __threadfence(); + atomicAdd(&p.state[0], 1); + } + } +} + +#endif + +} // namespace + +size_t cpTopkSmemBytes(int32_t localCapacity, int32_t worldSize) +{ + return smemBytes(localCapacity, worldSize); +} + +int32_t cpTopkMaxLocalCapacity(int32_t worldSize) +{ + static int32_t cached[kMaxRanks + 1] = {}; + if (worldSize < 1 || worldSize > kMaxRanks) + { + return 0; + } + if (cached[worldSize] == 0) + { + int device = 0; + cudaGetDevice(&device); + int maxSmemOptin = 0; + cudaDeviceGetAttribute(&maxSmemOptin, cudaDevAttrMaxSharedMemoryPerBlockOptin, device); + int32_t cap = 0; + for (int32_t c = 1 << 22; c >= kPageTokens * kSourcesPerRank; c -= kPageTokens * kSourcesPerRank) + { + if (smemBytes(c, worldSize) + kStaticSmemBytes <= static_cast(maxSmemOptin)) + { + cap = c; + break; + } + } + cached[worldSize] = cap; + } + return cached[worldSize]; +} + +int64_t cpTopkDefaultTimeoutMs() +{ + static int64_t const timeoutMs = []() + { + char const* env = std::getenv("TRTLLM_STAIRCASE_GLM_LL_CROSS_RANK_TOPK_TIMEOUT_MS"); + if (env == nullptr || env[0] == '\0') + { + return kDefaultTimeoutMs; + } + char* end = nullptr; + long long const v = std::strtoll(env, &end, 10); + if (end == env || *end != '\0' || v <= 0) + { + throw std::invalid_argument( + std::string("glm_moe_dsa_cross_rank_topk: bad TRTLLM_STAIRCASE_GLM_LL_CROSS_RANK_TOPK_TIMEOUT_MS=") + + env); + } + return static_cast(v); + }(); + return timeoutMs; +} + +void invokeCpTopk(CpTopkParams const& inParams, cudaStream_t stream) +{ + if (inParams.numRows <= 0) + { + return; + } + CpTopkParams params = inParams; + if (params.timeoutNs == 0) + { + params.timeoutNs = static_cast(cpTopkDefaultTimeoutMs()) * 1000000ull; + } + if (params.numRows > kMaxRows || params.worldSize <= 0 || params.worldSize > kMaxRanks || params.rank < 0 + || params.rank >= params.worldSize) + { + throw std::invalid_argument("glm_moe_dsa_cross_rank_topk: unsupported rows/world/rank"); + } + size_t const smem = smemBytes(params.localCapacity, params.worldSize); + // The attribute is per function and device; raise it only when a larger footprint is requested. + static thread_local size_t configuredSmem[64] = {}; + int device = 0; + cudaGetDevice(&device); + cudaError_t err = cudaSuccess; + if (device >= 64 || configuredSmem[device] < smem) + { + err = cudaFuncSetAttribute(cpTopkKernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(smem)); + if (err != cudaSuccess) + { + throw std::runtime_error( + std::string("glm_moe_dsa_cross_rank_topk: smem attribute: ") + cudaGetErrorString(err)); + } + if (device < 64) + { + configuredSmem[device] = smem; + } + } + // Programmatic dependent launch as for the other TensorRT-LLM kernels (on unless TRTLLM_ENABLE_PDL=0). + static bool const enablePdl = []() + { + char const* env = std::getenv("TRTLLM_ENABLE_PDL"); + return !(env != nullptr && env[0] == '0' && env[1] == '\0'); + }(); + cudaLaunchConfig_t config{}; + config.gridDim = dim3(kSourcesPerRank, params.numRows); + config.blockDim = dim3(kThreads); + config.dynamicSmemBytes = smem; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = enablePdl ? 1 : 0; + config.attrs = attrs; + config.numAttrs = 1; + err = cudaLaunchKernelEx(&config, cpTopkKernel, params); + if (err == cudaSuccess) + { + err = cudaGetLastError(); + } + if (err != cudaSuccess) + { + throw std::runtime_error(std::string("glm_moe_dsa_cross_rank_topk: launch: ") + cudaGetErrorString(err)); + } +} + +} // namespace kernels::glm_moe_dsa::cross_rank_topk + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopk.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopk.h new file mode 100644 index 000000000000..2ff17f85e48e --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopk.h @@ -0,0 +1,162 @@ +/* + * 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 +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa::cross_rank_topk +{ + +// Context-partitioned DSA indexer top-k. +// +// The indexer context of one generation request is partitioned over the ranks of a tensor-parallel group by whole +// 64-token pages: global page p belongs to rank p % worldSize, and a rank's local column c maps to global token +// ((c / 64) * worldSize + rank) * 64 + c % 64. Every rank scores only its own pages (the caller supplies the local +// logits). Within a rank, each verify row is handled by kSourcesPerRank CTAs; CTA s owns one 8-token block of every +// local page (block (s + page) % kSourcesPerRank), so the (rank, CTA) "sources" interleave the context finely and a +// fully selected page is spread over all sources of its rank. Each source +// 1. selects its local top keys (kSourceBudget at 8 ranks, see below) with a block radix select, +// 2. broadcasts them as (order-preserving key, global token) pairs, plus a header (count, exclusive bound of the +// keys it withheld), through a multicast Lamport workspace. +// One CTA per row and rank then merges the candidates of all sources into the exact top-kTopK of the row, +// identically on every rank. A source that withheld a key above the merged threshold makes every rank run a +// second round in which that source's top keys between the threshold and its bound are sent (at most +// kRound2Capacity, by the merge CTA of the source's rank, from the local logits), so the result is the tie-aware +// exact top-kTopK for any input. In the second round, a source whose cut falls inside a run of equal keys sends the +// keys above the cut and the first equal ones in source order up to its capacity (round 1 sends only the keys above +// such a cut, so the two rounds never send the same key). +// Rows of at most kTopK valid tokens produce the identity selection followed by -1, without an exchange. +// A launch whose rows all have at most kTopK valid tokens (the same on every rank: the lengths come from the shared +// KV length) is an identity launch: it touches neither the workspace nor the epoch. Every other launch is an +// exchange launch. Exchange launches rotate kNumLamportBuffers buffers by a device epoch that counts exchange +// launches; in exchange launch x the CTAs that do not merge reset, on their own rank, the buffer of exchange launch +// x - 1. Every rank's merges of x - 1 have read their copy of it, and a peer writes it again only in exchange +// launch x + 2, which starts after the peer's merge of x + 1 received this rank's sources of x + 1, i.e. after this +// rank completed launch x (identity launches in between neither write nor wait, so they do not shorten this chain). +// +// A merge that waits longer than the timeout (CpTopkParams::timeoutNs, %globaltimer) for a peer's words records the +// wait in the state's error word and aborts the kernel with a trap: the exchange cannot complete correctly any more, +// and a late peer write could land in a buffer this rank reuses, so the run fails instead of continuing. + +//! Selected indices per verify row (DSA index_topk). +constexpr int32_t kTopK = 2048; +//! Ranks of the context-parallel group. +constexpr int32_t kMaxRanks = 8; +//! Verify rows per call. +constexpr int32_t kMaxRows = 8; +//! Tokens per indexer K page; the context is partitioned by whole pages. +constexpr int32_t kPageTokens = 64; +//! CTAs (exchange sources) per rank and verify row. +constexpr int32_t kSourcesPerRank = 8; +//! Exchange sources per row. +constexpr int32_t kMaxSources = kMaxRanks * kSourcesPerRank; +//! Round-1 slots per (source, row). +constexpr int32_t kSourceCapacity = 256; +//! Keys each saturated source aims to send in round 1 at 8 ranks (64 sources): three times the even share of kTopK +//! (the largest share of the exact top-2048 that one source of this layout held on real GLM-5.2 indexer logits was +//! 69). A source stops refining as soon as its keys at or above a bin edge number at most four times the even +//! share, so it sends between 96 and 128 keys at 64 sources. Other group sizes scale both with ceil(kTopK / sources). +constexpr int32_t kSourceBudget = 96; +//! Round-2 slots per (source, row). +constexpr int32_t kRound2Capacity = 2048; +//! Threads per CTA. +constexpr int32_t kThreads = 1024; +//! Lamport buffers rotated by the device epoch. +constexpr int32_t kNumLamportBuffers = 3; +//! 32-bit word that marks an empty Lamport slot; never a valid key, token index, count or bound. +constexpr uint32_t kEmptyWord = 0xFFFFFFFFu; + +//! 32-bit words of one (source, row) slot: a 16-byte header (count, bound, two unused words) then two words per +//! candidate, so candidates can be written and read with 16-byte vector accesses. +constexpr int64_t kRound1SlotWords = 4 + 2 * int64_t{kSourceCapacity}; +constexpr int64_t kRound2SlotWords = 4 + 2 * int64_t{kRound2Capacity}; +//! 32-bit words of one Lamport buffer (round-1 region, then round-2 region). +constexpr int64_t kBufferWords = int64_t{kMaxSources} * kMaxRows * (kRound1SlotWords + kRound2SlotWords); +//! 32-bit words of the whole exchange workspace; every word must start as kEmptyWord. +constexpr int64_t kWorkspaceWords = kNumLamportBuffers * kBufferWords; +//! int32 words of the per-rank device state: epoch (exchange launches), CTA arrivals, second rounds taken, error +//! word (nonzero = the wait that timed out, written right before the kernel aborts). +constexpr int32_t kStateWords = 4; +//! Default wait for a peer's Lamport words before the kernel aborts +//! (TRTLLM_STAIRCASE_GLM_LL_CROSS_RANK_TOPK_TIMEOUT_MS overrides it). +constexpr int64_t kDefaultTimeoutMs = 120000; + +//! Values of the error word. +enum CpTopkWait : int32_t +{ + kWaitRound1Header = 1, + kWaitRound1Candidates = 2, + kWaitRound2Header = 3, + kWaitRound2Candidates = 4 +}; + +struct CpTopkParams +{ + //! [numRows, logitsStride] fp32 logits of this rank's local columns. + float const* logits; + int64_t logitsStride; + //! Local columns the logits view holds (each row's local length is clamped to it). + int32_t localCapacity; + //! [1] int32 global indexer KV length of the request; row r has kvLen - nextN + r + 1 valid tokens. + int32_t const* kvLens; + //! [numRows, kTopK] int32 output: global token indices, -1 for unused slots. + int32_t* out; + int32_t numRows; + int32_t nextN; + int32_t rank; + int32_t worldSize; + //! Multicast address of the exchange workspace (kWorkspaceWords words, identical layout on every rank). + uint32_t* mcWorkspace; + //! This rank's unicast address of the same workspace. + uint32_t* ucWorkspace; + //! [kStateWords] int32 per-rank device state, zero-initialized. + int32_t* state; + //! Optional [numRows, kSourcesPerRank, kDebugSlots] int64 phase timestamps (diagnosis only; nullptr = off). + int64_t* debugTimes = nullptr; + //! Longest wait for a peer's words before the kernel aborts (ns); 0 = the default (kDefaultTimeoutMs, or + //! TRTLLM_STAIRCASE_GLM_LL_CROSS_RANK_TOPK_TIMEOUT_MS). + uint64_t timeoutNs = 0; +}; + +//! Timestamps per CTA in CpTopkParams::debugTimes: [0, 16) %globaltimer (ns), [16, 32) clock64 per phase, +//! [32, 42) local select (keys, then clock64 at the start of / after the histogram of / after the bin search of +//! each pass), [42, 45) send (counter ready, keys taken, stores issued), 45 candidates polled before the key-range +//! reduction, 46 round-1 keys sent by this source, 47 local select passes (+10 when the cut was searched among the +//! negative keys), [48, 58) merge select (as the local select), [58, 62) emission (threshold decided, positions +//! scanned, row assembled, row written), 62 merge candidates, 63 merge select passes (+10 likewise); clock64 and +//! values of thread 0. +constexpr int32_t kDebugSlots = 64; + +//! Effective wait timeout of a launch with params.timeoutNs == 0 (ms). +int64_t cpTopkDefaultTimeoutMs(); + +//! Dynamic shared memory the kernel needs (bytes). +size_t cpTopkSmemBytes(int32_t localCapacity, int32_t worldSize); + +//! Largest local capacity the kernel accepts for a world size. +int32_t cpTopkMaxLocalCapacity(int32_t worldSize); + +void invokeCpTopk(CpTopkParams const& params, cudaStream_t stream); + +} // namespace kernels::glm_moe_dsa::cross_rank_topk + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopkClusterMerge.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopkClusterMerge.cu new file mode 100644 index 000000000000..a0b9671537c5 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopkClusterMerge.cu @@ -0,0 +1,2420 @@ +/* + * 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/glmMoeDsa/crossRankTopkClusterMerge.h" + +#include + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa::cross_rank_topk_cluster_merge +{ +namespace +{ + +constexpr int32_t kBins = 2048; +static_assert(kBins >= kTopK, "the output row is staged in the histogram"); +constexpr int32_t kWarps = kThreads / 32; +//! Per-warp key statistics kept in shared memory (see KeyStats). +constexpr int32_t kNumKeyStats = 5; +//! Tokens of each page one source holds. +constexpr int32_t kBlockTokens = kPageTokens / kSourcesPerRank; +//! CTAs of one row's cluster. +constexpr int32_t kClusterCtas = kSourcesPerRank; +//! Histogram bins per group of the cluster bin search (one warp lane per bin). +constexpr int32_t kGroupBins = 32; +constexpr int32_t kMaxGroups = kBins / kGroupBins; +//! Static shared memory of the kernel (round-2 segment tables and the cluster exchange records), kept out of the +//! dynamic budget. +constexpr size_t kStaticSmemBytes = 8192; + +// Records the CTAs of a row's cluster send each other with st.async (static, so every CTA has them at the same +// address). Exchange x of a row completes on barrier bar[x % 2], phase x / 2; the record sets alternate by pass +// parity and are written again only after every CTA has completed the exchange in between. +struct ClusterRecords +{ + //! Transaction barriers of the exchanges (arrival count 1: this CTA's expect-tx). + uint64_t bar[2]; + //! [pass parity][source CTA][group] histogram totals per group of kGroupBins bins. + uint32_t groups[2][kClusterCtas][kMaxGroups]; + //! [pass parity][source CTA][bin] counts of the selected group's bins, ascending. + uint32_t bins[2][kClusterCtas][kGroupBins]; +}; + +// Scalar slots in shared memory. +enum ScalarSlot : int32_t +{ + kSlotBin = 0, + kSlotAbove, + kSlotBinCount, + kSlotTotal, + kSlotCounter2, + kSlotSrcCount, // [kMaxSources] + kSlotSrcBound = kSlotSrcCount + kMaxSources, // [kMaxSources] + kSlotSrcLen = kSlotSrcBound + kMaxSources, // [kMaxSources] + kSlotSrcOffset = kSlotSrcLen + kMaxSources, // [kMaxSources + 1] + kSlotCountBefore = kSlotSrcOffset + kMaxSources + 1, // candidates of the cluster CTAs before this one + kSlotCountMine, // candidates of this cluster CTA + kSlotCountTotal, // candidates of the row + kSlotPositivesBefore, // candidates >= +0.0 of the cluster CTAs before this one + kNumScalarSlots +}; + +__host__ __device__ inline int32_t numSources(int32_t world) +{ + return world * kSourcesPerRank; +} + +// Even share of kTopK per source. +__host__ __device__ inline int32_t evenShare(int32_t world) +{ + return (kTopK + numSources(world) - 1) / numSources(world); +} + +// Most keys a source sends in round 1 (an even number). +__host__ __device__ inline int32_t roundCap(int32_t world) +{ + return min(kSourceCapacity, 4 * evenShare(world)); +} + +// Keys one source can hold: one block of every local page. +__host__ __device__ inline int32_t chunkCapacity(int32_t localCapacity) +{ + int32_t const pages = (localCapacity + kPageTokens - 1) / kPageTokens; + return pages * kBlockTokens; +} + +__host__ __device__ inline int32_t mergeCapacity(int32_t world) +{ + return numSources(world) * roundCap(world); +} + +__host__ __device__ inline size_t smemBytes(int32_t localCapacity, int32_t world) +{ + size_t words = 2 * size_t{kBins} + 2 * size_t{kWarps} + size_t{kNumKeyStats} * kWarps + size_t{kNumScalarSlots}; + words = (words + 3) / 4 * 4; + words += static_cast(chunkCapacity(localCapacity)) + 2 * static_cast(mergeCapacity(world)); + return words * sizeof(uint32_t); +} + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 900) +// The kernel targets SM90+ (multicast Lamport exchange, thread-block clusters). +__global__ void __launch_bounds__(kThreads, 1) cpTopkZ9Kernel(CpTopkParams p) +{ + __trap(); +} +#else + +//! Bits of one radix-select digit (kBins bins). +constexpr int32_t kDigitBits = 11; +//! Order-preserving key of +0.0: the keys of non-negative values are at least this, those of negative values below. +//! Real indexer logits are mostly negative, and the cut of a row (and of many sources) often lies among small +//! negative values while the largest keys are positive: one key range across zero spans ~2^31 key units, where a +//! radix digit has no resolution. The selects therefore work on one sign at a time: with fewer than `need` keys at +//! or above kPositiveKey these are all taken and the cut is searched among the negative keys only, else among the +//! non-negative keys only; each range is narrow enough for one or two passes. +constexpr uint32_t kPositiveKey = 0x80000000u; +//! Independent 16-byte candidate loads one merge thread issues before it waits on any of them. +constexpr int32_t kPollBatch = 4; +//! Polls of a still-empty Lamport word between two reads of the global timer. +constexpr int32_t kPollsPerTimerRead = 256; +//! Warps of the cluster merge: warps [0, kPollWarps) poll the round-1 candidates of this CTA's sources (one 16-byte +//! vector per thread: at most 520 vectors per CTA, at world sizes 5 to 7), the next kHeaderWarps warps poll the +//! round-1 headers of every source (one per thread), and the remaining kResetWarps warps reset the previous exchange +//! launch's buffer meanwhile. +constexpr int32_t kPollWarps = 17; +constexpr int32_t kHeaderWarps = 2; +constexpr int32_t kFirstResetWarp = kPollWarps + kHeaderWarps; +constexpr int32_t kResetWarps = kWarps - kFirstResetWarp; +static_assert(kHeaderWarps * 32 >= kMaxSources, "one header per header-warp thread"); + +// Keys a saturated source aims to send in round 1 (kSourceBudget at 8 ranks). +__device__ __forceinline__ int32_t roundNeed(int32_t world) +{ + return min(kSourceCapacity, 3 * evenShare(world)); +} + +// Tokens of `rank` among the global tokens [0, n). +__host__ __device__ __forceinline__ int32_t localLength(int32_t n, int32_t rank, int32_t world) +{ + if (n <= 0) + { + return 0; + } + int32_t const fullPages = n / kPageTokens; + int32_t const rem = n % kPageTokens; + int32_t const ownedFull = (fullPages - rank + world - 1) / world; + int32_t len = ownedFull * kPageTokens; + if (rem > 0 && fullPages % world == rank) + { + len += rem; + } + return len; +} + +// Source `src` of a rank holds one block of kBlockTokens consecutive tokens of every local page: block +// (src + page) % kSourcesPerRank. A fully selected page is thus spread over all sources of its rank. +// Keys of source `src` among `localLen` local columns (only the last local page can be partial). +__host__ __device__ __forceinline__ int32_t sourceLength(int32_t localLen, int32_t src) +{ + int32_t const pages = (localLen + kPageTokens - 1) / kPageTokens; + if (pages == 0) + { + return 0; + } + int32_t const last = pages - 1; + int32_t const tail = localLen - last * kPageTokens - ((src + last) % kSourcesPerRank) * kBlockTokens; + return last * kBlockTokens + min(max(tail, 0), kBlockTokens); +} + +// Global token of a local column under the page-interleaved partition. +__device__ __forceinline__ uint32_t localToGlobal(int32_t col, int32_t rank, int32_t world) +{ + int32_t const page = col / kPageTokens; + return static_cast((page * world + rank) * kPageTokens + (col % kPageTokens)); +} + +struct SmemLayout +{ + uint32_t* hist; // [2 * kBins] radix histograms of alternate passes; the first one stages the output row + uint32_t* warpSums; // [2 * kWarps] + uint32_t* stats; // [kNumKeyStats * kWarps] per-warp key statistics + uint32_t* scal; // [kNumScalarSlots] + uint32_t* chunk; // [chunkCapacity] order-preserving keys of this source + uint32_t* mkey; // [mergeCapacity] staged round-1 pairs, then the merge candidates (merge CTA) + uint32_t* midx; // [mergeCapacity] +}; + +__device__ inline SmemLayout carveSmem(uint32_t* base, int32_t localCapacity, int32_t world) +{ + SmemLayout s; + uint32_t* p = base; + s.hist = p; + p += 2 * kBins; + s.warpSums = p; + p += 2 * kWarps; + s.stats = p; + p += kNumKeyStats * kWarps; + s.scal = p; + p += kNumScalarSlots; + p = base + ((p - base) + 3) / 4 * 4; + s.chunk = p; + p += chunkCapacity(localCapacity); + s.mkey = p; + p += mergeCapacity(world); + s.midx = p; + return s; +} + +__device__ __forceinline__ uint32_t orderedKey(float v) +{ + uint32_t const b = __float_as_uint(v); + uint32_t const k = (b & 0x80000000u) ? ~b : (b | 0x80000000u); + // NaN keys above +inf are clamped so that no key, and no exclusive bound key + 1, equals the empty marker. + return k < kEmptyWord - 1u ? k : kEmptyWord - 2u; +} + +// Bits of x (0 for 0): the digit range of keys in [base, base + x]. +__device__ __forceinline__ int32_t bitLength(uint32_t x) +{ + return x == 0u ? 0 : 32 - __clz(x); +} + +// Statistics of a set of keys: extremes overall and per sign, and the count of non-negative keys. +struct KeyStats +{ + uint32_t maxKey = 0u; // largest key (0 when none) + uint32_t minKey = kEmptyWord; // smallest key (kEmptyWord when none) + uint32_t positives = 0u; // keys >= kPositiveKey + uint32_t maxNeg = 0u; // largest key below kPositiveKey (0 when none) + uint32_t minPos = kEmptyWord; // smallest key >= kPositiveKey (kEmptyWord when none) + + __device__ __forceinline__ void add(uint32_t k) + { + bool const pos = k >= kPositiveKey; + maxKey = max(maxKey, k); + minKey = min(minKey, k); + positives += pos ? 1u : 0u; + maxNeg = pos ? maxNeg : max(maxNeg, k); + minPos = pos ? min(minPos, k) : minPos; + } +}; + +// Reduce a thread's statistics over its warp; lane 0 stores them in the per-warp array (read with readStats after +// the next block barrier). +__device__ __forceinline__ void storeWarpStats(KeyStats const& t, uint32_t* stats) +{ + uint32_t const v0 = __reduce_max_sync(0xFFFFFFFFu, t.maxKey); + uint32_t const v1 = __reduce_min_sync(0xFFFFFFFFu, t.minKey); + uint32_t const v2 = __reduce_add_sync(0xFFFFFFFFu, t.positives); + uint32_t const v3 = __reduce_max_sync(0xFFFFFFFFu, t.maxNeg); + uint32_t const v4 = __reduce_min_sync(0xFFFFFFFFu, t.minPos); + if (threadIdx.x % 32 == 0) + { + int32_t const w = threadIdx.x / 32; + stats[w] = v0; + stats[kWarps + w] = v1; + stats[2 * kWarps + w] = v2; + stats[3 * kWarps + w] = v3; + stats[4 * kWarps + w] = v4; + } +} + +// Block statistics from the per-warp array (every warp reduces it; no barrier). +__device__ __forceinline__ KeyStats readStats(uint32_t const* stats) +{ + int32_t const lane = threadIdx.x % 32; + KeyStats b; + b.maxKey = __reduce_max_sync(0xFFFFFFFFu, stats[lane]); + b.minKey = __reduce_min_sync(0xFFFFFFFFu, stats[kWarps + lane]); + b.positives = __reduce_add_sync(0xFFFFFFFFu, stats[2 * kWarps + lane]); + b.maxNeg = __reduce_max_sync(0xFFFFFFFFu, stats[3 * kWarps + lane]); + b.minPos = __reduce_min_sync(0xFFFFFFFFu, stats[4 * kWarps + lane]); + return b; +} + +// Digit setup of a select of the need-th largest key: over the non-negative keys when there are at least `need` of +// them, else over the negative keys with every non-negative key counted as taken (`aboveHi`). +struct SignRange +{ + uint32_t lo; // smallest key of the range (also the digit base) + uint64_t hi; // exclusive end of the histogrammed range + int32_t bits; // digit range of the range + uint32_t aboveHi; // keys at or above hi, all taken + bool negative; // the cut is among the negative keys +}; + +__device__ __forceinline__ SignRange signRange(KeyStats const& st, uint32_t need) +{ + if (st.positives >= need || st.positives == 0u) + { + uint32_t const lo = st.positives >= need ? st.minPos : st.minKey; + return SignRange{lo, uint64_t{kEmptyWord}, bitLength(st.maxKey - lo), 0u, st.positives == 0u}; + } + return SignRange{st.minKey, uint64_t{kPositiveKey}, bitLength(st.maxNeg - st.minKey), st.positives, true}; +} + +__device__ __forceinline__ int64_t round1Offset(int32_t buf, int32_t src, int32_t row) +{ + return int64_t{buf} * kBufferWords + (int64_t{src} * kMaxRows + row) * kRound1SlotWords; +} + +__device__ __forceinline__ int64_t round2Offset(int32_t buf, int32_t src, int32_t row) +{ + return int64_t{buf} * kBufferWords + int64_t{kMaxSources} * kMaxRows * kRound1SlotWords + + (int64_t{src} * kMaxRows + row) * kRound2SlotWords; +} + +__device__ __forceinline__ uint2 loadVolatile2(uint32_t const* ptr) +{ + uint2 v; + asm volatile("ld.volatile.global.v2.u32 {%0, %1}, [%2];" : "=r"(v.x), "=r"(v.y) : "l"(ptr) : "memory"); + return v; +} + +__device__ __forceinline__ uint4 loadVolatile4(uint32_t const* ptr) +{ + uint4 v; + asm volatile("ld.volatile.global.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(ptr) + : "memory"); + return v; +} + +__device__ __forceinline__ void store2(uint32_t* ptr, uint32_t a, uint32_t b) +{ + asm volatile("st.global.v2.u32 [%0], {%1, %2};" ::"l"(ptr), "r"(a), "r"(b) : "memory"); +} + +__device__ __forceinline__ void store4(uint32_t* ptr, uint4 v) +{ + asm volatile("st.global.v4.u32 [%0], {%1, %2, %3, %4};" ::"l"(ptr), "r"(v.x), "r"(v.y), "r"(v.z), "r"(v.w) + : "memory"); +} + +__device__ __forceinline__ bool isEmpty(uint2 v) +{ + return v.x == kEmptyWord || v.y == kEmptyWord; +} + +// Whether the first pair of a 16-byte candidate vector, and the second one when `second`, were written. +__device__ __forceinline__ bool vectorReady(uint4 v, bool second) +{ + return v.x != kEmptyWord && v.y != kEmptyWord && (!second || (v.z != kEmptyWord && v.w != kEmptyWord)); +} + +__device__ __forceinline__ uint64_t globalTimerNs() +{ + uint64_t t; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t)); + return t; +} + +// What a polling thread needs to give up: the error word, the timeout and this rank (for the report). +struct WaitCtx +{ + int32_t* state; + uint64_t timeoutNs; + int32_t rank; +}; + +// A peer did not deliver within the timeout. The selection cannot be completed correctly any more, and the peer's +// late writes could land in a buffer this rank resets or reuses, so record which wait failed and abort the kernel: +// the resulting CUDA error fails the run instead of letting any rank continue with a partial selection. +__device__ __noinline__ void abortOnTimeout(WaitCtx const& w, int32_t what) +{ + bool const first = atomicExch(&w.state[3], what) == 0; + __threadfence_system(); + if (first) + { + printf( + "[glm_moe_dsa] crossRankTopkClusterMerge: rank %d, row %d: wait %d for a peer exceeded %llu ms; aborting " + "the kernel\n", + w.rank, static_cast(blockIdx.y), what, static_cast(w.timeoutNs / 1000000ull)); + } + __trap(); +} + +// Spin until both words of a Lamport pair were written; aborts the kernel after the timeout. +__device__ uint2 pollPair(uint32_t const* ptr, WaitCtx const& w, int32_t what) +{ + uint2 v = loadVolatile2(ptr); + uint64_t start = 0; + for (int32_t polls = 1; isEmpty(v); ++polls) + { + if (polls % kPollsPerTimerRead == 0) + { + uint64_t const now = globalTimerNs(); + if (start == 0) + { + start = now; + } + else if (now - start > w.timeoutNs) + { + abortOnTimeout(w, what); + } + } + v = loadVolatile2(ptr); + } + return v; +} + +// Spin until a 16-byte candidate vector was written (see vectorReady); aborts the kernel after the timeout. +__device__ uint4 pollVector(uint32_t const* ptr, bool second, WaitCtx const& w, int32_t what) +{ + uint4 v = loadVolatile4(ptr); + uint64_t start = 0; + for (int32_t polls = 1; !vectorReady(v, second); ++polls) + { + if (polls % kPollsPerTimerRead == 0) + { + uint64_t const now = globalTimerNs(); + if (start == 0) + { + start = now; + } + else if (now - start > w.timeoutNs) + { + abortOnTimeout(w, what); + } + } + v = loadVolatile4(ptr); + } + return v; +} + +// Block-wide inclusive scan of one value per thread. Returns the inclusive prefix; `total` gets the block sum. +__device__ uint32_t blockInclusiveScan(uint32_t v, uint32_t* warpSums, uint32_t& total) +{ + int32_t const lane = threadIdx.x % 32; + int32_t const warp = threadIdx.x / 32; + uint32_t x = v; +#pragma unroll + for (int32_t o = 1; o < 32; o <<= 1) + { + uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, x, o); + if (lane >= o) + { + x += y; + } + } + if (lane == 31) + { + warpSums[warp] = x; + } + __syncthreads(); + if (warp == 0) + { + uint32_t w = warpSums[lane]; +#pragma unroll + for (int32_t o = 1; o < 32; o <<= 1) + { + uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, w, o); + if (lane >= o) + { + w += y; + } + } + warpSums[kWarps + lane] = w; + } + __syncthreads(); + total = warpSums[2 * kWarps - 1]; + uint32_t const before = warp > 0 ? warpSums[kWarps + warp - 1] : 0u; + return x + before; +} + +// Exclusive block-wide scan with a single barrier: every warp scans the per-warp sums itself. Returns this +// thread's exclusive prefix; `total` gets the block sum. warpSums must not be rewritten before the next barrier. +__device__ uint32_t blockExclusiveScan(uint32_t v, uint32_t* warpSums, uint32_t& total) +{ + int32_t const lane = threadIdx.x % 32; + int32_t const warp = threadIdx.x / 32; + uint32_t x = v; +#pragma unroll + for (int32_t o = 1; o < 32; o <<= 1) + { + uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, x, o); + if (lane >= o) + { + x += y; + } + } + if (lane == 31) + { + warpSums[warp] = x; + } + __syncthreads(); + uint32_t const w = warpSums[lane]; + uint32_t y = w; +#pragma unroll + for (int32_t o = 1; o < 32; o <<= 1) + { + uint32_t const z = __shfl_up_sync(0xFFFFFFFFu, y, o); + if (lane >= o) + { + y += z; + } + } + total = __shfl_sync(0xFFFFFFFFu, y, 31); + return __shfl_sync(0xFFFFFFFFu, y - w, warp) + x - v; +} + +// Find, counting from the top bin, the bin of `hist` that holds the target-th largest key, with two barriers: +// per-warp sums, then each warp scans them and only the warp holding the target scans its own bins. Returns the +// total count (every thread); when target <= total, the bin, the count above it and its own count go to the +// scalar slots. +__device__ uint32_t findBin(SmemLayout const& s, uint32_t const* hist, int32_t nbins, uint32_t target) +{ + int32_t const lane = threadIdx.x % 32; + int32_t const warp = threadIdx.x / 32; + int32_t const perThread = (nbins + kThreads - 1) / kThreads; + int32_t const top = nbins - 1 - perThread * static_cast(threadIdx.x); + uint32_t mine = 0; + for (int32_t i = 0; i < perThread && top - i >= 0; ++i) + { + mine += hist[top - i]; + } + uint32_t const warpTotal = __reduce_add_sync(0xFFFFFFFFu, mine); + if (lane == 0) + { + s.warpSums[warp] = warpTotal; + } + __syncthreads(); + uint32_t const w = s.warpSums[lane]; + uint32_t x = w; +#pragma unroll + for (int32_t o = 1; o < 32; o <<= 1) + { + uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, x, o); + if (lane >= o) + { + x += y; + } + } + uint32_t const total = __shfl_sync(0xFFFFFFFFu, x, 31); + uint32_t const warpBefore = __shfl_sync(0xFFFFFFFFu, x - w, warp); + if (warpBefore < target && target <= warpBefore + warpTotal) + { + uint32_t t = mine; +#pragma unroll + for (int32_t o = 1; o < 32; o <<= 1) + { + uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, t, o); + if (lane >= o) + { + t += y; + } + } + uint32_t acc = warpBefore + t - mine; + if (acc < target && target <= acc + mine) + { + for (int32_t i = 0; i < perThread && top - i >= 0; ++i) + { + uint32_t const c = hist[top - i]; + if (acc + c >= target) + { + s.scal[kSlotBin] = static_cast(top - i); + s.scal[kSlotAbove] = acc; + s.scal[kSlotBinCount] = c; + break; + } + acc += c; + } + } + } + __syncthreads(); + return total; +} + +struct SelectOut +{ + uint32_t edge; // lower edge of the selected bin (the exact key after the last digit) + uint32_t above; // predicate keys strictly above the selected bin + uint32_t binCount; // predicate keys in the selected bin + uint32_t total; // predicate keys + bool all; // take every predicate key (fewer than needed, or all fit the capacity) + bool full; // exact mode: the keys at or above `edge` are exactly the `need` largest + int32_t passes; // radix passes run +}; + +// Candidate keys and token indices held in shared memory (16-byte aligned, readable in groups of four up to the +// next multiple of four). +struct SmemKeys +{ + static constexpr bool kVector = true; + uint32_t const* keys; + uint32_t const* idx; + + __device__ uint4 keys4(int32_t j) const + { + return *reinterpret_cast(keys + j); + } + + __device__ uint4 index4(int32_t j) const + { + return *reinterpret_cast(idx + j); + } +}; + +// Merge candidates in shared memory followed by round-2 candidates read from their global Lamport slots. +struct MergeKeys +{ + static constexpr bool kVector = false; + uint32_t const* smemKeys; + uint32_t const* smemIdx; + int32_t smemCount; + uint32_t const* const* slots; // round-2 slot of each segment + uint32_t const* segBegin; // [numSegments + 1] candidate offsets + int32_t numSegments; + + __device__ uint32_t const* pair(int32_t k) const + { + // Last segment whose begin is <= k. + int32_t lo = 0; + int32_t hi = numSegments - 1; + while (lo < hi) + { + int32_t const mid = (lo + hi + 1) / 2; + if (static_cast(segBegin[mid]) <= k) + { + lo = mid; + } + else + { + hi = mid - 1; + } + } + return slots[lo] + kSlotHeaderWords + 2 * (k - static_cast(segBegin[lo])); + } + + __device__ uint32_t operator()(int32_t j) const + { + return j < smemCount ? smemKeys[j] : pair(j - smemCount)[0]; + } + + __device__ uint32_t index(int32_t j) const + { + return j < smemCount ? smemIdx[j] : pair(j - smemCount)[1]; + } +}; + +// Debug slots of this CTA starting at `offset` (nullptr without a debug buffer). +__device__ __forceinline__ int64_t* debugSlots(CpTopkParams const& p, int32_t offset) +{ + if (p.debugTimes == nullptr) + { + return nullptr; + } + return p.debugTimes + (static_cast(blockIdx.y) * kSourcesPerRank + blockIdx.x) * kDebugSlots + offset; +} + +// clock64 of thread 0 into a debug slot (no-op without a debug buffer). +__device__ __forceinline__ void stamp(int64_t* dbg, int32_t slot) +{ + if (dbg != nullptr && threadIdx.x == 0) + { + dbg[slot] = clock64(); + } +} + +// A value into a debug slot (thread 0; no-op without a debug buffer). +__device__ __forceinline__ void putDebug(int64_t* dbg, int32_t slot, int64_t value) +{ + if (dbg != nullptr && threadIdx.x == 0) + { + dbg[slot] = value; + } +} + +// Phase timestamp for diagnosis (thread 0 of the CTA; no-op without a debug buffer). +__device__ __forceinline__ void markPhase(CpTopkParams const& p, int32_t phase) +{ + if (p.debugTimes != nullptr && threadIdx.x == 0) + { + uint64_t t; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t)); + int64_t* d = p.debugTimes + (static_cast(blockIdx.y) * kSourcesPerRank + blockIdx.x) * kDebugSlots; + d[phase] = static_cast(t); + d[16 + phase] = static_cast(clock64()); + } +} + +// Histogram update of one key of a radix pass. +__device__ __forceinline__ void histogramKey(uint32_t* hist, bool inRange, uint32_t k, uint32_t lo, uint64_t hi, + uint32_t base, uint32_t prefMask, uint32_t prefVal, int32_t shift, uint32_t binMask) +{ + uint32_t const d = k - base; + if (inRange && k >= lo && static_cast(k) < hi && (d & prefMask) == prefVal) + { + atomicAdd(&hist[(d >> shift) & binMask], 1u); + } +} + +// Block radix select of the need-th largest of the n keys at or above lo, of which `aboveHi` (< need) lie at or +// above hi and are counted as taken without being histogrammed. Every key in [lo, hi) is at least `base`, and +// key - base < 2^topBits; digits of up to kDigitBits bits of key - base are refined from the top. With cap > 0 the +// select stops at the first bin edge whose keys at or above it number at most `cap`. With cap = 0 it runs to the +// exact key, or stops early when the selected bin is wholly taken (`full`). Alternate passes use alternate +// histograms, so the next one is cleared while the current one is searched. Counts in the result include aboveHi. +template +__device__ SelectOut blockSelect(SmemLayout const& s, KeyAt const& keyAt, int32_t n, uint32_t lo, uint64_t hi, + uint32_t need, uint32_t cap, uint32_t base, int32_t topBits, int64_t* dbg = nullptr, uint32_t aboveHi = 0u) +{ + if (dbg != nullptr && threadIdx.x == 0) + { + dbg[0] = n; + } + SelectOut out{lo, 0u, 0u, 0u, false, false, 0}; + uint32_t prefMask = 0; + uint32_t prefVal = 0; + uint32_t above = aboveHi; + int32_t remaining = max(topBits, 1); + uint32_t* hist = s.hist; + uint32_t* histNext = s.hist + kBins; + for (int32_t b = threadIdx.x; b < (1 << min(kDigitBits, remaining)); b += kThreads) + { + hist[b] = 0; + } + __syncthreads(); +#pragma unroll 1 + for (int32_t pass = 0; remaining > 0; ++pass) + { + int32_t const width = min(kDigitBits, remaining); + int32_t const shift = remaining - width; + int32_t const nbins = 1 << width; + uint32_t const binMask = static_cast(nbins - 1); + out.passes = pass + 1; + if (dbg != nullptr && threadIdx.x == 0 && pass < 3) + { + dbg[1 + 3 * pass] = clock64(); + } + if constexpr (KeyAt::kVector) + { + for (int32_t j0 = 4 * static_cast(threadIdx.x); j0 < n; j0 += 4 * kThreads) + { + uint4 const q = keyAt.keys4(j0); + histogramKey(hist, true, q.x, lo, hi, base, prefMask, prefVal, shift, binMask); + histogramKey(hist, j0 + 1 < n, q.y, lo, hi, base, prefMask, prefVal, shift, binMask); + histogramKey(hist, j0 + 2 < n, q.z, lo, hi, base, prefMask, prefVal, shift, binMask); + histogramKey(hist, j0 + 3 < n, q.w, lo, hi, base, prefMask, prefVal, shift, binMask); + } + } + else + { + for (int32_t j = threadIdx.x; j < n; j += kThreads) + { + histogramKey(hist, true, keyAt(j), lo, hi, base, prefMask, prefVal, shift, binMask); + } + } + __syncthreads(); + if (dbg != nullptr && threadIdx.x == 0 && pass < 3) + { + dbg[2 + 3 * pass] = clock64(); + } + if (shift > 0) + { + for (int32_t b = threadIdx.x; b < (1 << min(kDigitBits, shift)); b += kThreads) + { + histNext[b] = 0; + } + } + // Find the bin holding the (need - above)-th largest key, counting from the top bin. + uint32_t const target = need - above; + uint32_t const total = findBin(s, hist, nbins, target); + if (dbg != nullptr && threadIdx.x == 0 && pass < 3) + { + dbg[3 + 3 * pass] = clock64(); + } + if (pass == 0) + { + out.total = total + aboveHi; + if (total + aboveHi <= cap || target > total) + { + out.all = true; + out.above = total + aboveHi; + return out; + } + } + // The scalars are next written after two more barriers (the next histogram and scan). + uint32_t const bin = s.scal[kSlotBin]; + above += s.scal[kSlotAbove]; + out.binCount = s.scal[kSlotBinCount]; + prefVal |= bin << shift; + prefMask |= binMask << shift; + out.edge = base + prefVal; + out.above = above; + remaining = shift; + uint32_t* const t = hist; + hist = histNext; + histNext = t; + if (cap > 0 && above + out.binCount <= cap) + { + return out; + } + if (cap == 0 && remaining > 0 && above + out.binCount == need) + { + out.full = true; + return out; + } + } + return out; +} + +// Send one source's selected keys (keys in [lo, hi) of the chunk) to its slot; returns the count sent. +// Takes every predicate key at or above the edge when they fit `cap`. Otherwise (more than `cap` keys tie at the +// exact edge) it takes every key above the edge, plus, with kPartialTies, the first keys equal to it in source +// order up to `cap` (any subset of equal keys is a valid tie-aware choice). Round 1 sends no partial ties, so the +// keys a second round requests from a source (at most the edge) were never sent before. +// The header carries the count and an exclusive upper bound of the keys left unsent (0 when none); in round 1 +// (!kPartialTies) also the key statistics of the keys sent (see kSlotHeaderWords). +// With `stage` (round 1, at most kSourceCapacity keys), the pairs are gathered in shared memory first and written +// with coalesced 16-byte stores; otherwise each pair is stored where it is found. +template +__device__ uint32_t sendSource(SmemLayout const& s, SelectOut const& sel, int32_t n, uint32_t lo, uint64_t hi, + uint32_t cap, uint32_t* mcSlot, int32_t src, int32_t rank, int32_t world, uint2* stage, int64_t* dbg = nullptr) +{ + bool const strict = !sel.all && sel.above + sel.binCount > cap; + uint32_t const threshold = sel.all ? lo : sel.edge; + // Keys equal to the edge that fit after the keys above it (sel.above < need <= cap when strict). + uint32_t const tieBudget = kPartialTies && strict ? cap - sel.above : 0u; + uint32_t const sent = sel.all ? sel.total : (strict ? sel.above + tieBudget : sel.above + sel.binCount); + uint32_t const bound = sel.all ? lo : (strict ? sel.edge + 1u : sel.edge); + stamp(dbg, 42); + // One page block (kBlockTokens keys) per thread; positions from a block scan, so the slot holds the keys in + // source order: first those taken outright, then (strict only) the admitted keys equal to the edge. + int32_t const groups = (n + kBlockTokens - 1) / kBlockTokens; + uint32_t takenMain = 0; + uint32_t takenTies = 0; + uint32_t const tieBase = strict ? sel.above : 0u; + for (int32_t g0 = 0; g0 < groups; g0 += kThreads) + { + int32_t const g = g0 + static_cast(threadIdx.x); + uint32_t k[kBlockTokens]; + uint32_t mask = 0; + uint32_t tieMask = 0; + if (g < groups) + { + uint4 const a = *reinterpret_cast(s.chunk + kBlockTokens * g); + uint4 const b = *reinterpret_cast(s.chunk + kBlockTokens * g + 4); + k[0] = a.x; + k[1] = a.y; + k[2] = a.z; + k[3] = a.w; + k[4] = b.x; + k[5] = b.y; + k[6] = b.z; + k[7] = b.w; +#pragma unroll + for (int32_t i = 0; i < kBlockTokens; ++i) + { + bool const in = kBlockTokens * g + i < n && k[i] >= lo && static_cast(k[i]) < hi; + bool const take = in && (strict ? k[i] > threshold : k[i] >= threshold); + bool const tie = kPartialTies && in && strict && k[i] == threshold; + mask |= (take ? 1u : 0u) << i; + tieMask |= (tie ? 1u : 0u) << i; + } + } + // At most kThreads * kBlockTokens keys per pass: both counts fit 16 bits. + uint32_t const c = static_cast(__popc(mask)) | (static_cast(__popc(tieMask)) << 16); + uint32_t chunkTaken; + uint32_t const before = blockExclusiveScan(c, s.warpSums, chunkTaken); + uint32_t pos = takenMain + (before & 0xFFFFu); + uint32_t tiePos = takenTies + (before >> 16); + takenMain += chunkTaken & 0xFFFFu; + takenTies += chunkTaken >> 16; + if ((mask | tieMask) != 0u) + { + int32_t const colBase = g * kPageTokens + ((src + g) % kSourcesPerRank) * kBlockTokens; +#pragma unroll + for (int32_t i = 0; i < kBlockTokens; ++i) + { + uint32_t slot = kEmptyWord; + if ((mask >> i) & 1u) + { + slot = pos++; + } + else if constexpr (kPartialTies) + { + if ((tieMask >> i) & 1u) + { + slot = tiePos < tieBudget ? tieBase + tiePos : kEmptyWord; + ++tiePos; + } + } + if (slot != kEmptyWord) + { + uint32_t const token = localToGlobal(colBase + i, rank, world); + if (stage != nullptr) + { + stage[slot] = make_uint2(k[i], token); + } + else + { + store2(mcSlot + kSlotHeaderWords + 2 * static_cast(slot), k[i], token); + } + } + } + } + // The staged pairs are complete, and the scan's warp sums may be rewritten. + __syncthreads(); + } + stamp(dbg, 43); + if (stage != nullptr) + { + // Two pairs per 16-byte store; an odd count leaves the last half of the last vector empty. + int32_t const vectors = static_cast(sent + 1u) / 2; + for (int32_t v = threadIdx.x; v < vectors; v += kThreads) + { + uint2 const a = stage[2 * v]; + uint2 const b + = 2 * v + 1 < static_cast(sent) ? stage[2 * v + 1] : make_uint2(kEmptyWord, kEmptyWord); + store4(mcSlot + kSlotHeaderWords + 4 * static_cast(v), make_uint4(a.x, a.y, b.x, b.y)); + } + } + // Receivers check every candidate word, so the header needs no ordering against the candidates. + if constexpr (kPartialTies) + { + if (threadIdx.x == kThreads - 1) + { + store2(mcSlot, sent, bound); + } + } + else + { + // The statistics of the keys sent, from the staged pairs (round 1 always stages), by the last warp, which + // stores no vector (at most kSourceCapacity / 2), while the first warps store the vectors. + static_assert(kSourceCapacity / 2 <= kThreads - 32, "the last warp stores no candidate vector"); + if (threadIdx.x >= kThreads - 32) + { + int32_t const lane = threadIdx.x % 32; + KeyStats t; +#pragma unroll + for (int32_t j = 0; j < kSourceCapacity / 32; ++j) + { + int32_t const i = lane + 32 * j; + if (i < static_cast(sent)) + { + t.add(stage[i].x); + } + } + uint32_t const maxKey = __reduce_max_sync(0xFFFFFFFFu, t.maxKey); + uint32_t const minKey = __reduce_min_sync(0xFFFFFFFFu, t.minKey); + uint32_t const positives = __reduce_add_sync(0xFFFFFFFFu, t.positives); + uint32_t const maxNeg = __reduce_max_sync(0xFFFFFFFFu, t.maxNeg); + uint32_t const minPos = __reduce_min_sync(0xFFFFFFFFu, t.minPos); + if (lane == 31) + { + store4(mcSlot, make_uint4(sent, bound, maxKey, minKey + 1u)); + store4(mcSlot + 4, make_uint4(positives, maxNeg, minPos + 1u, 0u)); + } + } + } + stamp(dbg, 44); + return sent; +} + +// Load one source's keys (one block of kBlockTokens tokens of every local page) from the local logits into the +// chunk; returns the number of keys. Consecutive threads read consecutive keys (whole 32-byte blocks per eight +// threads), and each thread issues kLoadBatch loads before it uses any. With `withStats`, the per-warp key +// statistics are stored for readStats after the next block barrier. +__device__ int32_t loadSource( + SmemLayout const& s, float const* rowLogits, int32_t localLen, int32_t src, bool withStats) +{ + constexpr int32_t kLoadBatch = 8; + int32_t const n = sourceLength(localLen, src); + KeyStats st; + for (int32_t j0 = threadIdx.x; j0 < n; j0 += kLoadBatch * kThreads) + { + float v[kLoadBatch]; +#pragma unroll + for (int32_t i = 0; i < kLoadBatch; ++i) + { + int32_t const j = j0 + i * kThreads; + int32_t const page = j / kBlockTokens; + int32_t const col = page * kPageTokens + ((src + page) % kSourcesPerRank) * kBlockTokens + j % kBlockTokens; + v[i] = j < n ? rowLogits[col] : 0.F; + } +#pragma unroll + for (int32_t i = 0; i < kLoadBatch; ++i) + { + int32_t const j = j0 + i * kThreads; + if (j < n) + { + uint32_t const k = orderedKey(v[i]); + s.chunk[j] = k; + st.add(k); + } + } + } + if (withStats) + { + storeWarpStats(st, s.stats); + } + return n; +} + +// Whether source q withheld keys that may lie above the merged k-th key. +__device__ __forceinline__ bool sourceOverflows(SmemLayout const& s, int32_t q, uint32_t kth) +{ + uint32_t const cnt = s.scal[kSlotSrcCount + q]; + uint64_t const bound = s.scal[kSlotSrcBound + q]; + return cnt < s.scal[kSlotSrcLen + q] && bound > uint64_t{kth} + 1u; +} + +// Emit the top-kTopK of the merge candidates. Keys above `kth` are all taken; ties at `kth` are taken in candidate +// order when only some of them fit (deterministic on every rank), otherwise all of them. With `takeAll`, every +// candidate is taken and the row is padded with -1. The row is assembled in shared memory (the first histogram) +// and written with coalesced stores. +template +__device__ bool emitSelection(SmemLayout const& s, KeyAt const& keys, int32_t total, uint32_t kth, bool takeAll, + int32_t* outRow, int32_t overflowSources, int64_t* dbg = nullptr) +{ + uint32_t const want = static_cast(kTopK); + // Positions from one block scan over contiguous candidate ranges (no contended counter); the order is the + // candidate order, identical on every rank. Shared-memory candidates are read four at a time. + int32_t per = (total + kThreads - 1) / kThreads; + if constexpr (KeyAt::kVector) + { + per = (per + 3) / 4 * 4; + } + int32_t const j0 = min(total, static_cast(threadIdx.x) * per); + int32_t const j1 = min(total, j0 + per); + uint32_t cGt = 0; + uint32_t cEq = 0; + if constexpr (KeyAt::kVector) + { + for (int32_t j = j0; j < j1; j += 4) + { + uint4 const q = keys.keys4(j); + uint32_t const kk[4] = {q.x, q.y, q.z, q.w}; +#pragma unroll + for (int32_t i = 0; i < 4; ++i) + { + bool const in = j + i < j1; + cGt += (in && (takeAll || kk[i] > kth)) ? 1u : 0u; + cEq += (in && !takeAll && kk[i] == kth) ? 1u : 0u; + } + } + } + else + { + for (int32_t j = j0; j < j1; ++j) + { + uint32_t const k = keys(j); + cGt += (takeAll || k > kth) ? 1u : 0u; + cEq += (!takeAll && k == kth) ? 1u : 0u; + } + } + // The first `overflowSources` threads flag a source that withheld keys above the threshold (the flag slot is + // zero on entry); with a flag set, nothing is written and the caller runs the second round. + if (static_cast(threadIdx.x) < overflowSources && sourceOverflows(s, threadIdx.x, kth)) + { + atomicOr(&s.scal[kSlotCounter2], 1u); + } + uint32_t allGt; + uint32_t allEq; + uint32_t posGt; + uint32_t posEq; + if (total < 0x10000) + { + // Both counts fit 16 bits: one scan of the packed pair. + uint32_t packedTotal; + uint32_t const packed = cGt | (cEq << 16); + uint32_t const before = blockExclusiveScan(packed, s.warpSums, packedTotal); + allGt = packedTotal & 0xFFFFu; + allEq = packedTotal >> 16; + posGt = before & 0xFFFFu; + posEq = allGt + (before >> 16); + } + else + { + posGt = blockInclusiveScan(cGt, s.warpSums, allGt) - cGt; + __syncthreads(); + posEq = blockInclusiveScan(cEq, s.warpSums, allEq) - cEq; + posEq += allGt; + } + stamp(dbg, 59); + if (overflowSources > 0 && s.scal[kSlotCounter2] != 0u) + { + return true; + } + uint32_t* row = s.hist; + auto place = [&](uint32_t k, uint32_t idx) + { + if (takeAll || k > kth) + { + if (posGt < want) + { + row[posGt] = idx; + } + ++posGt; + } + else if (k == kth) + { + if (posEq < want) + { + row[posEq] = idx; + } + ++posEq; + } + }; + if constexpr (KeyAt::kVector) + { + for (int32_t j = j0; j < j1; j += 4) + { + uint4 const q = keys.keys4(j); + uint4 const x = keys.index4(j); + uint32_t const kk[4] = {q.x, q.y, q.z, q.w}; + uint32_t const xx[4] = {x.x, x.y, x.z, x.w}; +#pragma unroll + for (int32_t i = 0; i < 4; ++i) + { + if (j + i < j1) + { + place(kk[i], xx[i]); + } + } + } + } + else + { + for (int32_t j = j0; j < j1; ++j) + { + place(keys(j), keys.index(j)); + } + } + __syncthreads(); + stamp(dbg, 60); + uint32_t const selected = min(allGt + allEq, want); + for (uint32_t j = threadIdx.x; j < want; j += kThreads) + { + outRow[j] = j < selected ? static_cast(row[j]) : -1; + } + stamp(dbg, 61); + return false; +} + +// Threshold of an exact (cap = 0) merge select: keys above `kth` are taken, ties at `kth` while they fit. +__device__ __forceinline__ uint32_t mergeThreshold(SelectOut const& m, bool& takeAll) +{ + takeAll = m.all || (m.full && m.edge == 0u); + if (takeAll) + { + return 0u; + } + // A wholly taken bin: the keys at or above its edge are exactly kTopK. + return m.full ? m.edge - 1u : m.edge; +} + +// Reset of the previous exchange launch's buffer on this rank's copy. Its 2 * kMaxSources * kMaxRows slots (round 1, +// then round 2) are spread over the kResetWarps reset warps (the warps from kFirstResetWarp on) of every CTA of the +// launch, at most ten consecutive slots per warp, one per lane: each lane fetches its slot's header once the epoch is +// known (fetchResetHeaders, after the grid dependency wait and the first block barrier; an asynchronous copy into +// shared memory that completes during the local select and the send) and the warp resets the written slots while the +// poll warps wait for the candidates (resetSlots). A slot was written iff its header was; the header's count bounds +// the candidates to reset. +constexpr int32_t kResetSlots = 2 * kMaxSources * kMaxRows; + +// This CTA's index among the launch's clearers, and their number (from the grid; recomputed where used). +__device__ __forceinline__ int32_t resetClearer() +{ + return static_cast(blockIdx.y) * kSourcesPerRank + static_cast(blockIdx.x); +} + +__device__ __forceinline__ int32_t resetClearers() +{ + return static_cast(gridDim.y) * kSourcesPerRank; +} + +__device__ __forceinline__ int32_t resetSlotsPerWarp(int32_t numClearers) +{ + int32_t const warps = numClearers * kResetWarps; + return (kResetSlots + warps - 1) / warps; +} + +// First slot of this (reset) warp. +__device__ __forceinline__ int32_t resetFirstSlot(int32_t clearer, int32_t numClearers) +{ + int32_t const warp = static_cast(threadIdx.x) / 32 - kFirstResetWarp; + return (clearer * kResetWarps + warp) * resetSlotsPerWarp(numClearers); +} + +__device__ __forceinline__ int64_t resetSlotOffset(int32_t buf, int32_t slot) +{ + int32_t const round1Slots = kMaxSources * kMaxRows; + bool const second = slot >= round1Slots; + int32_t const qr = second ? slot - round1Slots : slot; + return second ? round2Offset(buf, qr / kMaxRows, qr % kMaxRows) : round1Offset(buf, qr / kMaxRows, qr % kMaxRows); +} + +__device__ __forceinline__ uint32_t* resetSlotBase(CpTopkParams const& p, int32_t buf, int32_t slot) +{ + return p.ucWorkspace + resetSlotOffset(buf, slot); +} + +//! Most slots one reset warp resets (one row: the fewest reset warps). +constexpr int32_t kMaxResetSlotsPerWarp + = (kResetSlots + kSourcesPerRank * kResetWarps - 1) / (kSourcesPerRank * kResetWarps); +static_assert(kMaxResetSlotsPerWarp <= 32, "one slot header per lane"); + +// Shared-memory home of the reset warps' slot headers (16 bytes per slot: count, bound and two further header words), +// in the merge candidate area beyond the round-1 staging of the send; it is free from the first block barrier until +// the reset. +constexpr int32_t kResetHeaderOffset = 2 * kSourceCapacity; +static_assert(kResetHeaderOffset + 4 * kMaxResetSlotsPerWarp * kResetWarps <= kSourcesPerRank * kSourceCapacity, + "the reset headers fit the smallest merge candidate area"); + +__device__ __forceinline__ uint4* resetHeaderHome(SmemLayout const& s) +{ + int32_t const warp = static_cast(threadIdx.x) / 32 - kFirstResetWarp; + return reinterpret_cast(s.mkey + kResetHeaderOffset) + warp * kMaxResetSlotsPerWarp; +} + +// Word offset, within one Lamport buffer, of the slot whose header this lane fetches for the reset; -1 for the lanes +// and warps without one. Arithmetic on the grid only (no memory access). +__device__ __forceinline__ int32_t resetHeaderOffset() +{ + int32_t const numClearers = resetClearers(); + int32_t const lane = threadIdx.x % 32; + if (static_cast(threadIdx.x) / 32 < kFirstResetWarp || lane >= resetSlotsPerWarp(numClearers)) + { + return -1; + } + int32_t const slot = resetFirstSlot(resetClearer(), numClearers) + lane; + if (slot >= kResetSlots) + { + return -1; + } + static_assert(kBufferWords <= int64_t{1} << 31, "slot offsets within a buffer fit int32"); + return static_cast(resetSlotOffset(0, slot)); +} + +// Starts the copy of this lane's reset slot header of buffer `buf` (offset from resetHeaderOffset) into the warp's +// home with an asynchronous copy (cp.async, generic proxy, completed in resetSlots): a pending copy holds no register +// and no block barrier waits for it. Lanes without a slot do nothing. +__device__ __forceinline__ void fetchResetHeader( + CpTopkParams const& p, SmemLayout const& s, int32_t buf, int32_t offset) +{ + if (offset >= 0) + { + uint32_t const dst = static_cast( + __cvta_generic_to_shared(resetHeaderHome(s) + static_cast(threadIdx.x % 32))); + uint32_t const* const src = p.ucWorkspace + int64_t{buf} * kBufferWords + offset; + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;" ::"r"(dst), "l"(src) : "memory"); + } +} + +// Reset this reset warp's slots of buffer `buf` (headers fetched by fetchResetHeaders); other warps return. +__device__ void resetSlots(CpTopkParams const& p, SmemLayout const& s, int32_t buf) +{ + if (static_cast(threadIdx.x) / 32 < kFirstResetWarp) + { + return; + } + int32_t const clearer = resetClearer(); + int32_t const numClearers = resetClearers(); + // Each lane completes its own header copy; the warp barrier then makes every lane's header visible to the warp. + asm volatile("cp.async.wait_all;" ::: "memory"); + __syncwarp(); + int32_t const per = resetSlotsPerWarp(numClearers); + int32_t const lane = threadIdx.x % 32; + int32_t const first = resetFirstSlot(clearer, numClearers); + uint4 const* const headers = resetHeaderHome(s); + uint4 const empty = make_uint4(kEmptyWord, kEmptyWord, kEmptyWord, kEmptyWord); + for (int32_t i = 0; i < per && first + i < kResetSlots; ++i) + { + uint4 const h = headers[i]; + if (h.x == kEmptyWord && h.y == kEmptyWord) + { + continue; + } + int32_t const slot = first + i; + uint32_t* base = resetSlotBase(p, buf, slot); + uint32_t const capacity = slot >= kMaxSources * kMaxRows ? kRound2Capacity : kSourceCapacity; + uint32_t const count = h.x == kEmptyWord ? capacity : min(h.x, capacity); + int32_t const vectors = static_cast(count + 1u) / 2; + for (int32_t v = lane; v < vectors; v += 32) + { + store4(base + kSlotHeaderWords + 4 * static_cast(v), empty); + } + if (lane < 2) + { + store4(base + 4 * lane, empty); + } + } +} + +// Merge of one row on this rank (one CTA): gathers every source's round-1 candidates, selects the exact +// top-kTopK, runs the second round when a source withheld keys above the threshold, and writes the row. +__device__ void mergeRow(CpTopkParams const& p, SmemLayout const& s, uint32_t const** segSlots, uint32_t* segBegin, + int32_t buf, int32_t row, int32_t rowLen, int32_t localLen, float const* rowLogits, int32_t* outRow, + WaitCtx const& wait) +{ + int32_t const world = p.worldSize; + int32_t const rank = p.rank; + int32_t const sources = numSources(world); + uint32_t const cap = static_cast(roundCap(world)); + __syncthreads(); + if (threadIdx.x == 0) + { + s.scal[kSlotCounter2] = 0u; + } + if (static_cast(threadIdx.x) < sources) + { + // Keys each source holds, computed while its header is in flight. + int32_t const q = threadIdx.x; + s.scal[kSlotSrcLen + q] = static_cast( + sourceLength(min(localLength(rowLen, q / kSourcesPerRank, world), p.localCapacity), q % kSourcesPerRank)); + uint2 const h = pollPair(p.ucWorkspace + round1Offset(buf, threadIdx.x, row), wait, kWaitRound1Header); + s.scal[kSlotSrcCount + threadIdx.x] = min(h.x, cap); + s.scal[kSlotSrcBound + threadIdx.x] = h.y; + } + __syncthreads(); + if (threadIdx.x < 32) + { + // Exclusive prefix of the source counts (at most 64 sources: two per lane). + uint32_t const a = static_cast(threadIdx.x) < sources ? s.scal[kSlotSrcCount + threadIdx.x] : 0u; + uint32_t const b + = static_cast(threadIdx.x) + 32 < sources ? s.scal[kSlotSrcCount + threadIdx.x + 32] : 0u; + uint32_t x = a; + uint32_t z = b; +#pragma unroll + for (int32_t o = 1; o < 32; o <<= 1) + { + uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, x, o); + uint32_t const w = __shfl_up_sync(0xFFFFFFFFu, z, o); + if (static_cast(threadIdx.x) >= o) + { + x += y; + z += w; + } + } + uint32_t const totalA = __shfl_sync(0xFFFFFFFFu, x, 31); + s.scal[kSlotSrcOffset + threadIdx.x] = x - a; + s.scal[kSlotSrcOffset + 32 + threadIdx.x] = totalA + z - b; + if (threadIdx.x == 31) + { + s.scal[kSlotSrcOffset + kMaxSources] = totalA + z; + } + } + __syncthreads(); + int32_t const total = static_cast(s.scal[kSlotSrcOffset + kMaxSources]); + markPhase(p, 4); + + // Candidates: padded (source, vector) positions; issue a batch of independent 16-byte loads, then re-poll the + // ones not yet written. The key statistics are tracked for the merge select. + int32_t const vectorsPerSource = static_cast(cap) / 2; + int32_t const padded = sources * vectorsPerSource; + KeyStats st; +#pragma unroll 1 + for (int32_t v0 = 0; v0 < padded; v0 += kThreads * kPollBatch) + { + uint4 v[kPollBatch]; +#pragma unroll + for (int32_t i = 0; i < kPollBatch; ++i) + { + int32_t const jv = v0 + i * kThreads + static_cast(threadIdx.x); + int32_t const q = jv / vectorsPerSource; + int32_t const kv = jv % vectorsPerSource; + v[i] = make_uint4(0u, 0u, 0u, 0u); + if (jv < padded && 2 * kv < static_cast(s.scal[kSlotSrcCount + q])) + { + v[i] = loadVolatile4(p.ucWorkspace + round1Offset(buf, q, row) + kSlotHeaderWords + 4 * kv); + } + } +#pragma unroll + for (int32_t i = 0; i < kPollBatch; ++i) + { + int32_t const jv = v0 + i * kThreads + static_cast(threadIdx.x); + int32_t const q = jv / vectorsPerSource; + int32_t const kv = jv % vectorsPerSource; + if (jv < padded && 2 * kv < static_cast(s.scal[kSlotSrcCount + q])) + { + bool const second = 2 * kv + 1 < static_cast(s.scal[kSlotSrcCount + q]); + if (!vectorReady(v[i], second)) + { + v[i] = pollVector(p.ucWorkspace + round1Offset(buf, q, row) + kSlotHeaderWords + 4 * kv, second, + wait, kWaitRound1Candidates); + } + int32_t const j = static_cast(s.scal[kSlotSrcOffset + q]) + 2 * kv; + s.mkey[j] = v[i].x; + s.midx[j] = v[i].y; + st.add(v[i].x); + if (second) + { + s.mkey[j + 1] = v[i].z; + s.midx[j + 1] = v[i].w; + st.add(v[i].z); + } + } + } + } + stamp(debugSlots(p, 0), 45); + storeWarpStats(st, s.stats); + __syncthreads(); + markPhase(p, 5); + + // Digits of the merge select span the candidates of the sign the cut lies in (see kPositiveKey). + SignRange const r = total > 0 ? signRange(readStats(s.stats), static_cast(kTopK)) + : SignRange{0u, uint64_t{kEmptyWord}, 32, 0u, false}; + SmemKeys const merged{s.mkey, s.midx}; + SelectOut const m = blockSelect( + s, merged, total, r.lo, r.hi, static_cast(kTopK), 0u, r.lo, r.bits, debugSlots(p, 48), r.aboveHi); + putDebug(debugSlots(p, 0), 62, total); + putDebug(debugSlots(p, 0), 63, m.passes + (r.negative ? 10 : 0)); + bool takeAll = false; + uint32_t const kth = mergeThreshold(m, takeAll); + markPhase(p, 6); + stamp(debugSlots(p, 0), 58); + if (!emitSelection(s, merged, total, kth, takeAll, outRow, sources, debugSlots(p, 0))) + { + return; + } + + // ---- Round 2 (rare; every rank takes it): withheld keys in [kth, bound) of overflowing sources ---------- + __syncthreads(); + // 1. This rank's overflowing sources send their top-kRound2Capacity withheld keys. + for (int32_t c = 0; c < kSourcesPerRank; ++c) + { + int32_t const q = rank * kSourcesPerRank + c; + if (!sourceOverflows(s, q, kth)) + { + continue; + } + uint64_t const b = s.scal[kSlotSrcBound + q]; + int32_t const nq = loadSource(s, rowLogits, localLen, c, false); + __syncthreads(); + SelectOut const sel2 = blockSelect( + s, SmemKeys{s.chunk, nullptr}, nq, kth, b, static_cast(kRound2Capacity), kRound2Capacity, 0u, 32); + (void) sendSource( + s, sel2, nq, kth, b, kRound2Capacity, p.mcWorkspace + round2Offset(buf, q, row), c, rank, world, nullptr); + __syncthreads(); + } + // 2. Gather the round-2 segments of every overflowing source (all ranks). + if (threadIdx.x == 0) + { + int32_t segs = 0; + uint32_t off = 0; + for (int32_t q = 0; q < sources; ++q) + { + if (!sourceOverflows(s, q, kth)) + { + continue; + } + uint32_t const* slot = p.ucWorkspace + round2Offset(buf, q, row); + uint2 const h = pollPair(slot, wait, kWaitRound2Header); + segSlots[segs] = slot; + segBegin[segs] = off; + off += min(h.x, static_cast(kRound2Capacity)); + ++segs; + } + segBegin[segs] = off; + s.scal[kSlotTotal] = static_cast(segs); + } + __syncthreads(); + int32_t const segs = static_cast(s.scal[kSlotTotal]); + int32_t const total2 = static_cast(segBegin[segs]); + MergeKeys const keys{s.mkey, s.midx, total, segSlots, segBegin, segs}; + for (int32_t k = threadIdx.x; k < total2; k += kThreads) + { + (void) pollPair(keys.pair(k), wait, kWaitRound2Candidates); + } + __syncthreads(); + SelectOut const m2 + = blockSelect(s, keys, total + total2, 0u, uint64_t{kEmptyWord}, static_cast(kTopK), 0u, 0u, 32); + bool takeAll2 = false; + uint32_t const kth2 = mergeThreshold(m2, takeAll2); + (void) emitSelection(s, keys, total + total2, kth2, takeAll2, outRow, 0); + if (threadIdx.x == 0) + { + atomicAdd(&p.state[2], 1); + } +} + +// ---- Thread-block cluster primitives (one cluster of kClusterCtas CTAs per row; cluster rank = blockIdx.x) ---------- + +__device__ __forceinline__ uint32_t smemAddr(void const* ptr) +{ + return static_cast(__cvta_generic_to_shared(ptr)); +} + +// Shared-memory window address of `ptr` (this CTA's shared memory) in CTA `rank` of the cluster. +__device__ __forceinline__ uint32_t clusterAddr(void const* ptr, int32_t rank) +{ + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(remote) : "r"(smemAddr(ptr)), "r"(rank)); + return remote; +} + +// Initialize the exchange barriers and announce it to the cluster (relaxed arrive; clusterInitWait completes it +// before the first exchange). Touches shared memory only. +__device__ __forceinline__ void clusterInit(ClusterRecords& rec) +{ + if (threadIdx.x == 0) + { + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(smemAddr(&rec.bar[0])) : "memory"); + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(smemAddr(&rec.bar[1])) : "memory"); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + asm volatile("barrier.cluster.arrive.relaxed.aligned;" ::: "memory"); +} + +__device__ __forceinline__ void clusterInitWait() +{ + asm volatile("barrier.cluster.wait.aligned;" ::: "memory"); +} + +// This CTA's single arrival on exchange barrier `bar`, expecting `bytes` from the cluster's st.async stores. +__device__ __forceinline__ void expectBytes(uint64_t* bar, uint32_t bytes) +{ + asm volatile( + "{\n\t.reg .b64 st;\n\tmbarrier.arrive.expect_tx.shared::cta.b64 st, [%0], %1;\n\t}" ::"r"(smemAddr(bar)), + "r"(bytes) + : "memory"); +} + +// Wait until the phase of exchange barrier `bar` with parity `parity` completed (every expected byte arrived). +__device__ __forceinline__ void waitExchange(uint64_t* bar, uint32_t parity) +{ + uint32_t done; + do + { + asm volatile( + "{\n\t.reg .pred p;\n\tmbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 p, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, p;\n\t}" + : "=r"(done) + : "r"(smemAddr(bar)), "r"(parity) + : "memory"); + } while (done == 0u); +} + +// Store 16 / 8 bytes at `ptr` (a record of this CTA's layout) in CTA `rank`, completing them on its barrier `bar`. +__device__ __forceinline__ void sendTo(int32_t rank, void const* ptr, uint64_t const* bar, uint4 v) +{ + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [%0], {%1, %2, %3, %4}, [%5];" ::"r"( + clusterAddr(ptr, rank)), + "r"(v.x), "r"(v.y), "r"(v.z), "r"(v.w), "r"(clusterAddr(bar, rank)) + : "memory"); +} + +__device__ __forceinline__ void sendTo(int32_t rank, void const* ptr, uint64_t const* bar, uint2 v) +{ + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.v2.b32 [%0], {%1, %2}, [%3];" ::"r"( + clusterAddr(ptr, rank)), + "r"(v.x), "r"(v.y), "r"(clusterAddr(bar, rank)) + : "memory"); +} + +// Exchange sequence of a row: per pass x = 2 * pass group totals and x = 2 * pass + 1 the selected group's bins. +// Every CTA sends the records of exchange x + 1 only after its own exchange x completed, so a barrier phase (and an +// alternating record set) is reused only after every CTA consumed the previous one. +__device__ __forceinline__ uint64_t* exchangeBar(ClusterRecords& rec, uint32_t x) +{ + return &rec.bar[x & 1u]; +} + +__device__ __forceinline__ uint32_t exchangeParity(uint32_t x) +{ + return (x >> 1) & 1u; +} + +// Inclusive scan over the lanes of a warp. +__device__ __forceinline__ uint32_t warpInclusiveScan(uint32_t v) +{ + int32_t const lane = threadIdx.x % 32; +#pragma unroll + for (int32_t o = 1; o < 32; o <<= 1) + { + uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, v, o); + if (lane >= o) + { + v += y; + } + } + return v; +} + +// Tallies of the cluster select over the candidates of the cluster CTAs before this one (in candidate order), kept +// by warp 0 and published in shared memory: those strictly above the selected bins of all passes so far, and those +// in the selected bin of the last pass. +struct CtaTally +{ + uint32_t aboveBefore; + uint32_t binBefore; +}; + +// Cluster radix select of the need-th largest key among the candidates of every CTA of the cluster (this thread +// holds k0 and k1; kEmptyWord is no candidate, as it lies above every range). It computes exactly what blockSelect +// with cap = 0 computes over the union of the candidates: the same digits (from `base` and topBits), the same bins +// and counts from the cluster-summed histograms, and the same stop rules. Per pass: every CTA histograms its own +// candidates and sends its totals per group of kGroupBins bins to every CTA (st.async, two groups per warp); warp 0 +// of every CTA finds, from the top, the group holding the target, sends its own counts of that group's bins to every +// CTA, and finds the bin from the eight CTAs' counts. From the same records, warp 0 also tallies the candidates of +// the CTAs before this one above the selected bins and in the last selected bin (`tally`), which give this CTA's +// emission offsets without another exchange. +__device__ SelectOut clusterSelect(SmemLayout const& s, ClusterRecords& rec, CtaTally& tally, uint32_t k0, uint32_t k1, + uint32_t lo, uint64_t hi, uint32_t need, uint32_t base, int32_t topBits, uint32_t aboveHi, int64_t* dbg) +{ + int32_t const me = blockIdx.x; + int32_t const lane = threadIdx.x % 32; + int32_t const warp = threadIdx.x / 32; + SelectOut out{lo, 0u, 0u, 0u, false, false, 0}; + uint32_t prefMask = 0; + uint32_t prefVal = 0; + uint32_t above = aboveHi; + int32_t remaining = max(topBits, 1); + uint32_t* hist = s.hist; + uint32_t* histNext = s.hist + kBins; + // Warp 0: the running tallies of the CTAs before this one. + uint32_t aboveBefore = 0u; + uint32_t binBefore = 0u; + for (int32_t b = threadIdx.x; b < (1 << min(kDigitBits, remaining)); b += kThreads) + { + hist[b] = 0; + } + __syncthreads(); +#pragma unroll 1 + for (int32_t pass = 0; remaining > 0; ++pass) + { + int32_t const width = min(kDigitBits, remaining); + int32_t const shift = remaining - width; + int32_t const nbins = 1 << width; + uint32_t const binMask = static_cast(nbins - 1); + int32_t const groups = (nbins + kGroupBins - 1) / kGroupBins; + int32_t const pairs = (groups + 1) / 2; + uint32_t const xg = 2u * pass; + uint32_t const xb = xg + 1u; + out.passes = pass + 1; + if (dbg != nullptr && threadIdx.x == 0 && pass < 3) + { + dbg[1 + 3 * pass] = clock64(); + } + if (threadIdx.x == 0) + { + expectBytes(exchangeBar(rec, xg), kClusterCtas * pairs * 8); + } + // The next pass's histogram: its last reads (the previous pass's bins sent by warp 0) are complete. + if (shift > 0) + { + for (int32_t b = threadIdx.x; b < (1 << min(kDigitBits, shift)); b += kThreads) + { + histNext[b] = 0; + } + } + histogramKey(hist, true, k0, lo, hi, base, prefMask, prefVal, shift, binMask); + histogramKey(hist, true, k1, lo, hi, base, prefMask, prefVal, shift, binMask); + __syncthreads(); + uint32_t(*const groupRec)[kMaxGroups] = rec.groups[pass & 1]; + uint32_t(*const binRec)[kGroupBins] = rec.bins[pass & 1]; + if (warp < pairs) + { + int32_t const b0 = 2 * warp * kGroupBins + lane; + int32_t const b1 = b0 + kGroupBins; + uint32_t const t0 = __reduce_add_sync(0xFFFFFFFFu, b0 < nbins ? hist[b0] : 0u); + uint32_t const t1 = __reduce_add_sync(0xFFFFFFFFu, b1 < nbins ? hist[b1] : 0u); + if (lane < kClusterCtas) + { + sendTo(lane, &groupRec[me][2 * warp], exchangeBar(rec, xg), make_uint2(t0, t1)); + } + } + uint32_t const target = need - above; + if (warp == 0) + { + waitExchange(exchangeBar(rec, xg), exchangeParity(xg)); + if (dbg != nullptr && lane == 0 && pass < 3) + { + dbg[2 + 3 * pass] = clock64(); + } + // Lane l sums the groups at positions 2l and 2l + 1 from the top over the cluster. + int32_t const ga = groups - 1 - 2 * lane; + int32_t const gb = ga - 1; + uint32_t a = 0; + uint32_t b = 0; + uint32_t beforeA = 0; + uint32_t beforeB = 0; +#pragma unroll + for (int32_t c = 0; c < kClusterCtas; ++c) + { + uint32_t const x = ga >= 0 ? groupRec[c][ga] : 0u; + uint32_t const y = gb >= 0 ? groupRec[c][gb] : 0u; + a += x; + b += y; + beforeA += c < me ? x : 0u; + beforeB += c < me ? y : 0u; + } + uint32_t const incl = warpInclusiveScan(a + b); + uint32_t const before = incl - a - b; + uint32_t const total = __shfl_sync(0xFFFFFFFFu, incl, 31); + bool const holds = before < target && target <= incl; + uint32_t const holder = __ballot_sync(0xFFFFFFFFu, holds); + if (lane == 0) + { + s.scal[kSlotTotal] = total; + } + if (holder != 0u) + { + int32_t const hl = __ffs(holder) - 1; + bool const first = target <= before + a; + int32_t const gSel = __shfl_sync(0xFFFFFFFFu, first ? ga : gb, hl); + uint32_t const gAbove = __shfl_sync(0xFFFFFFFFu, first ? before : before + a, hl); + if (lane == 0) + { + expectBytes(exchangeBar(rec, xb), kClusterCtas * kGroupBins * 4); + } + // This CTA's counts of the selected group's bins to every CTA: lane l sends vector l % 8 to CTAs + // l / 8 and l / 8 + 4. + int32_t const v = lane % 8; + int32_t const b0 = gSel * kGroupBins + 4 * v; + uint4 const q = make_uint4(b0 < nbins ? hist[b0] : 0u, b0 + 1 < nbins ? hist[b0 + 1] : 0u, + b0 + 2 < nbins ? hist[b0 + 2] : 0u, b0 + 3 < nbins ? hist[b0 + 3] : 0u); + sendTo(lane / 8, &binRec[me][4 * v], exchangeBar(rec, xb), q); + sendTo(lane / 8 + 4, &binRec[me][4 * v], exchangeBar(rec, xb), q); + // While the bins travel: the earlier CTAs' candidates in the groups above the selected one (both + // groups of the lanes before the holder, and the holder's first group when it selected the second). + bool const selFirst = __shfl_sync(0xFFFFFFFFu, first, hl); + aboveBefore += __reduce_add_sync( + 0xFFFFFFFFu, lane < hl ? beforeA + beforeB : (lane == hl && !selFirst ? beforeA : 0u)); + waitExchange(exchangeBar(rec, xb), exchangeParity(xb)); + // Bins of the selected group from the top: lane l holds bin 32 * gSel + 31 - l. + int32_t const bin = gSel * kGroupBins + kGroupBins - 1 - lane; + uint32_t c = 0; + uint32_t cBefore = 0; +#pragma unroll + for (int32_t d = 0; d < kClusterCtas; ++d) + { + uint32_t const x = binRec[d][kGroupBins - 1 - lane]; + c += x; + cBefore += d < me ? x : 0u; + } + uint32_t const y = warpInclusiveScan(c); + uint32_t const acc = gAbove + y - c; + bool const found = acc < target && target <= acc + c; + if (found) + { + s.scal[kSlotBin] = static_cast(bin); + s.scal[kSlotAbove] = acc; + s.scal[kSlotBinCount] = c; + } + // The earlier CTAs' candidates in the bins above the selected one (lanes before the finder) and in + // the selected bin (the finder's). + int32_t const finder = __ffs(__ballot_sync(0xFFFFFFFFu, found)) - 1; + aboveBefore += __reduce_add_sync(0xFFFFFFFFu, lane < finder ? cBefore : 0u); + binBefore = __shfl_sync(0xFFFFFFFFu, cBefore, finder); + } + else if (lane == 0) + { + // Every key is taken (pass 0 only, the same decision in every CTA): the bin exchange still + // completes, with no bytes, so that the exchange sequence stays aligned. + expectBytes(exchangeBar(rec, xb), 0u); + waitExchange(exchangeBar(rec, xb), exchangeParity(xb)); + } + if (lane == 0) + { + tally.aboveBefore = aboveBefore; + tally.binBefore = binBefore; + } + } + __syncthreads(); + if (dbg != nullptr && threadIdx.x == 0 && pass < 3) + { + dbg[3 + 3 * pass] = clock64(); + } + uint32_t const total = s.scal[kSlotTotal]; + if (pass == 0) + { + out.total = total + aboveHi; + if (total + aboveHi == 0u || target > total) + { + out.all = true; + out.above = total + aboveHi; + return out; + } + } + // The scalars are next written by warp 0 after the next histogram's block barrier. + uint32_t const bin = s.scal[kSlotBin]; + above += s.scal[kSlotAbove]; + out.binCount = s.scal[kSlotBinCount]; + prefVal |= bin << shift; + prefMask |= binMask << shift; + out.edge = base + prefVal; + out.above = above; + remaining = shift; + uint32_t* const t = hist; + hist = histNext; + histNext = t; + if (remaining > 0 && above + out.binCount == need) + { + out.full = true; + return out; + } + } + return out; +} + +// Spin until the header of a round-1 slot was written and, when the slot's count covers candidate vector kv, that +// vector too; returns the header's count (clamped to `cap`) and bound. Both are polled together, so neither waits +// for the other; aborts the kernel after the timeout. +__device__ uint2 pollSlotVector(uint32_t const* slot, int32_t kv, uint32_t cap, uint4& v, WaitCtx const& w) +{ + uint32_t const* vec = slot + kSlotHeaderWords + 4 * kv; + uint2 h = loadVolatile2(slot); + v = loadVolatile4(vec); + uint64_t start = 0; + for (int32_t polls = 1;; ++polls) + { + bool const headerReady = !isEmpty(h); + uint32_t const cnt = headerReady ? min(h.x, cap) : 0u; + if (headerReady && (2u * kv >= cnt || vectorReady(v, 2u * kv + 1u < cnt))) + { + return make_uint2(cnt, h.y); + } + if (polls % kPollsPerTimerRead == 0) + { + uint64_t const now = globalTimerNs(); + if (start == 0) + { + start = now; + } + else if (now - start > w.timeoutNs) + { + abortOnTimeout(w, headerReady ? kWaitRound1Candidates : kWaitRound1Header); + } + } + if (!headerReady) + { + h = loadVolatile2(slot); + } + v = loadVolatile4(vec); + } +} + +// Spin until the whole 32-byte round-1 header of a slot was written; aborts the kernel after the timeout. +__device__ void pollHeader(uint32_t const* slot, uint4& a, uint4& b, WaitCtx const& w) +{ + a = loadVolatile4(slot); + b = loadVolatile4(slot + 4); + uint64_t start = 0; + for (int32_t polls = 1; !(vectorReady(a, true) && vectorReady(b, true)); ++polls) + { + if (polls % kPollsPerTimerRead == 0) + { + uint64_t const now = globalTimerNs(); + if (start == 0) + { + start = now; + } + else if (now - start > w.timeoutNs) + { + abortOnTimeout(w, kWaitRound1Header); + } + } + a = loadVolatile4(slot); + b = loadVolatile4(slot + 4); + } +} + +// Merge of one row by its cluster. CTA c gathers the round-1 candidates of sources [c * world, (c + 1) * world) (one +// 16-byte vector, two candidate slots, per thread of the poll warps, kept in registers: thread order is crossRankTopk's +// candidate order) while its header warps read every source's header (counts, bounds and the key statistics of the keys +// sent, so every CTA has the row's statistics and can decide every source's overflow without an exchange). The cluster +// runs crossRankTopk's merge select on the summed histograms (clusterSelect), whose records also give every CTA's +// number of keys above the threshold and of ties, and every CTA writes its selected candidates at crossRankTopk's +// positions: keys above the threshold in candidate order, then ties in candidate order while they fit, then -1. Returns +// true, writing nothing, when a source withheld keys that may lie above the threshold (the same decision on every CTA +// and rank); the row is then merged by mergeRow. +__device__ bool clusterMergeRow(CpTopkParams const& p, SmemLayout const& s, ClusterRecords& rec, CtaTally& tally, + int32_t buf, int32_t row, int32_t rowLen, int32_t* outRow, WaitCtx const& wait) +{ + int32_t const world = p.worldSize; + int32_t const sources = numSources(world); + int32_t const me = blockIdx.x; + int32_t const tid = threadIdx.x; + int32_t const lane = tid % 32; + int32_t const warp = tid / 32; + uint32_t const cap = static_cast(roundCap(world)); + int32_t const vectorsPerSource = static_cast(cap) / 2; + int32_t const q0 = me * world; + int64_t* const dbg = debugSlots(p, 0); + // This thread's candidates (kEmptyWord: none; never a key). + uint32_t k0 = kEmptyWord; + uint32_t k1 = kEmptyWord; + if (tid == 0) + { + s.scal[kSlotCounter2] = 0u; + } + if (tid < world * vectorsPerSource) + { + int32_t const q = q0 + tid / vectorsPerSource; + int32_t const kv = tid % vectorsPerSource; + uint4 v; + uint2 const h = pollSlotVector(p.ucWorkspace + round1Offset(buf, q, row), kv, cap, v, wait); + if (2u * kv < h.x) + { + k0 = v.x; + if (2u * kv + 1u < h.x) + { + k1 = v.z; + } + } + // The tokens wait in shared memory until the emission (s.midx is free in the cluster merge). + *reinterpret_cast(s.midx + 2 * tid) = make_uint2(v.y, v.w); + } + else if (warp >= kPollWarps && warp < kFirstResetWarp) + { + // Header warps: thread h reads source h's header; per source its count (clamped to the round-1 capacity), + // bound and length, and per header warp the reduced key statistics. + int32_t const q = tid - kPollWarps * 32; + KeyStats st; + if (q < sources) + { + uint4 a; + uint4 b; + pollHeader(p.ucWorkspace + round1Offset(buf, q, row), a, b, wait); + s.scal[kSlotSrcCount + q] = min(a.x, cap); + s.scal[kSlotSrcBound + q] = a.y; + s.scal[kSlotSrcLen + q] = static_cast(sourceLength( + min(localLength(rowLen, q / kSourcesPerRank, world), p.localCapacity), q % kSourcesPerRank)); + st.maxKey = a.z; + st.minKey = a.w - 1u; + st.positives = b.x; + st.maxNeg = b.y; + st.minPos = b.z - 1u; + } + uint32_t const cnt = q < sources ? s.scal[kSlotSrcCount + q] : 0u; + // Candidates and keys >= +0.0 of the cluster CTAs before this one (CTA c merges sources + // [c * world, (c + 1) * world)), of this one and of the row. + int32_t const c = q < sources ? q / world : kClusterCtas; + uint32_t const nBefore = __reduce_add_sync(0xFFFFFFFFu, c < me ? cnt : 0u); + uint32_t const nMine = __reduce_add_sync(0xFFFFFFFFu, c == me ? cnt : 0u); + uint32_t const nAll = __reduce_add_sync(0xFFFFFFFFu, cnt); + uint32_t const pBefore = __reduce_add_sync(0xFFFFFFFFu, c < me ? st.positives : 0u); + if (lane == 0) + { + atomicAdd(&s.scal[kSlotCountBefore], nBefore); + atomicAdd(&s.scal[kSlotCountMine], nMine); + atomicAdd(&s.scal[kSlotCountTotal], nAll); + atomicAdd(&s.scal[kSlotPositivesBefore], pBefore); + } + uint32_t const v0 = __reduce_max_sync(0xFFFFFFFFu, st.maxKey); + uint32_t const v1 = __reduce_min_sync(0xFFFFFFFFu, st.minKey); + uint32_t const v2 = __reduce_add_sync(0xFFFFFFFFu, st.positives); + uint32_t const v3 = __reduce_max_sync(0xFFFFFFFFu, st.maxNeg); + uint32_t const v4 = __reduce_min_sync(0xFFFFFFFFu, st.minPos); + if (lane == 0) + { + // In the scan scratch (s.stats may still be read by the send's last warp; s.warpSums is next written + // by the emission scan, after the reads below). + uint32_t* const out = s.warpSums + kNumKeyStats * (warp - kPollWarps); + out[0] = v0; + out[1] = v1; + out[2] = v2; + out[3] = v3; + out[4] = v4; + } + } + stamp(dbg, 45); + __syncthreads(); + markPhase(p, 4); + + // Statistics of the row's candidates from every source's header (identical in every CTA and rank). + KeyStats g; +#pragma unroll + for (int32_t w = 0; w < kHeaderWarps; ++w) + { + uint32_t const* const in = s.warpSums + kNumKeyStats * w; + g.maxKey = max(g.maxKey, in[0]); + g.minKey = min(g.minKey, in[1]); + g.positives += in[2]; + g.maxNeg = max(g.maxNeg, in[3]); + g.minPos = min(g.minPos, in[4]); + } + uint32_t const total = s.scal[kSlotCountTotal]; + stamp(dbg, 61); + markPhase(p, 5); + + SignRange const r + = total > 0u ? signRange(g, static_cast(kTopK)) : SignRange{0u, uint64_t{kEmptyWord}, 32, 0u, false}; + clusterInitWait(); + SelectOut const m = clusterSelect( + s, rec, tally, k0, k1, r.lo, r.hi, static_cast(kTopK), r.lo, r.bits, r.aboveHi, debugSlots(p, 48)); + putDebug(debugSlots(p, 48), 0, s.scal[kSlotCountMine]); + putDebug(dbg, 62, total); + putDebug(dbg, 63, m.passes + (r.negative ? 10 : 0)); + bool takeAll = false; + uint32_t const kth = mergeThreshold(m, takeAll); + markPhase(p, 6); + + // Keys above the threshold and ties of the cluster CTAs before this one and of the row (identical in every CTA), + // from the select's tallies; with the threshold below a wholly taken bin the ties fall beyond kTopK and are taken + // as keys above it. + uint32_t gtBefore; + uint32_t eqBefore; + uint32_t allGt; + uint32_t allEq; + if (takeAll) + { + gtBefore = s.scal[kSlotCountBefore]; + eqBefore = 0u; + allGt = total; + allEq = 0u; + } + else + { + uint32_t const aboveBefore = tally.aboveBefore + (r.aboveHi > 0u ? s.scal[kSlotPositivesBefore] : 0u); + gtBefore = m.full ? aboveBefore + tally.binBefore : aboveBefore; + eqBefore = m.full ? 0u : tally.binBefore; + allGt = m.full ? m.above + m.binCount : m.above; + allEq = m.full ? 0u : m.binCount; + } + // This CTA's positions, scanned in thread order; a source of the row that withheld keys above the threshold + // sends the row to the second round. + bool const gt0 = k0 != kEmptyWord && (takeAll || k0 > kth); + bool const eq0 = k0 != kEmptyWord && !takeAll && k0 == kth; + bool const gt1 = k1 != kEmptyWord && (takeAll || k1 > kth); + bool const eq1 = k1 != kEmptyWord && !takeAll && k1 == kth; + if (tid < sources && sourceOverflows(s, tid, kth)) + { + atomicOr(&s.scal[kSlotCounter2], 1u); + } + // Exclusive prefix in thread order (k0 before k1): within the warp from ballots, across warps from the warp + // totals (one block barrier). + uint32_t const below = (1u << lane) - 1u; + uint32_t const bGt0 = __ballot_sync(0xFFFFFFFFu, gt0); + uint32_t const bGt1 = __ballot_sync(0xFFFFFFFFu, gt1); + uint32_t const bEq0 = __ballot_sync(0xFFFFFFFFu, eq0); + uint32_t const bEq1 = __ballot_sync(0xFFFFFFFFu, eq1); + if (lane == 0) + { + s.warpSums[warp] = static_cast(__popc(bGt0) + __popc(bGt1)) + + (static_cast(__popc(bEq0) + __popc(bEq1)) << 16); + } + __syncthreads(); + uint32_t const warpsBefore = __reduce_add_sync(0xFFFFFFFFu, lane < warp ? s.warpSums[lane] : 0u); + uint32_t const before = warpsBefore + static_cast(__popc(bGt0 & below) + __popc(bGt1 & below)) + + (static_cast(__popc(bEq0 & below) + __popc(bEq1 & below)) << 16); + stamp(dbg, 58); + if (s.scal[kSlotCounter2] != 0u) + { + return true; + } + uint32_t const want = static_cast(kTopK); + uint32_t posGt = gtBefore + (before & 0xFFFFu); + uint32_t posEq = allGt + eqBefore + (before >> 16); + uint32_t x0 = 0u; + uint32_t x1 = 0u; + if (k0 != kEmptyWord) + { + uint2 const x = *reinterpret_cast(s.midx + 2 * tid); + x0 = x.x; + x1 = x.y; + } + if (gt0 && posGt < want) + { + outRow[posGt] = static_cast(x0); + } + posGt += gt0 ? 1u : 0u; + if (eq0 && posEq < want) + { + outRow[posEq] = static_cast(x0); + } + posEq += eq0 ? 1u : 0u; + if (gt1 && posGt < want) + { + outRow[posGt] = static_cast(x1); + } + if (eq1 && posEq < want) + { + outRow[posEq] = static_cast(x1); + } + uint32_t const selected = min(allGt + allEq, want); + for (uint32_t j = selected + static_cast(me * kThreads + tid); j < want; j += kClusterCtas * kThreads) + { + outRow[j] = -1; + } + stamp(dbg, 60); + return false; +} + +__global__ void __launch_bounds__(kThreads, 1) cpTopkZ9Kernel(CpTopkParams p) +{ + markPhase(p, 0); + extern __shared__ __align__(16) uint32_t smemRaw[]; + __shared__ uint32_t const* segSlots[kMaxSources]; + __shared__ uint32_t segBegin[kMaxSources + 1]; + __shared__ __align__(16) ClusterRecords rec; + __shared__ CtaTally tally; + __shared__ int32_t epochShared; + int32_t const world = p.worldSize; + int32_t const rank = p.rank; + int32_t const src = blockIdx.x; + int32_t const row = blockIdx.y; + SmemLayout const s = carveSmem(smemRaw, p.localCapacity, world); + int32_t* outRow = p.out + static_cast(row) * kTopK; + clusterInit(rec); + if (threadIdx.x <= kSlotPositivesBefore - kSlotCountBefore) + { + s.scal[kSlotCountBefore + threadIdx.x] = 0u; + } + int32_t const resetOffset = resetHeaderOffset(); + // Every global read (lengths, logits, the epoch word, the workspace) comes after the grid dependency wait, so a + // programmatic early launch is safe whatever the preceding kernels write. + asm volatile("griddepcontrol.wait;" ::: "memory"); + markPhase(p, 10); + // Thread 0 and the reset warps read this CTA's epoch word, written only by the previous exchange launch of this + // op, before the key loads; nothing waits for the value until the key loads are done. Thread 0 publishes it to the + // CTA at the block barrier that follows the key loads, and the reset warps start the copies of the previous + // buffer's headers right before that barrier (slot offsets computed before the wait, from the grid only). + int32_t* const epochWord = p.state + kEpochBase + row * kSourcesPerRank + src; + int32_t epochRead = 0; + if (threadIdx.x == 0 || static_cast(threadIdx.x) >= kFirstResetWarp * 32) + { + epochRead = *reinterpret_cast(epochWord); + } + int32_t const prevBufRead = (epochRead + kNumLamportBuffers - 1) % kNumLamportBuffers; + int32_t const kvLen = p.kvLens[0]; + int32_t const rowLen = max(0, kvLen - p.nextN + row + 1); + // Identity launch (every row has at most kTopK tokens; the same decision on every rank): no exchange, and + // neither the workspace nor the epoch words are touched, so identity launches never shorten the reuse chain of + // the Lamport buffers (see crossRankTopkClusterMerge.h). + bool const exchange = kvLen - p.nextN + p.numRows > kTopK; + // Thread 0 advances this CTA's epoch word (the CTAs of the last row also those of the rows this launch does not + // have, so every word counts every exchange launch), after a block barrier that follows every read of the word in + // this CTA. Nothing else of this launch reads the epoch words, and the next launch reads them only after its grid + // dependency wait. + auto advanceEpoch = [&]() + { + if (exchange && threadIdx.x == 0) + { + int32_t const rows = row == static_cast(gridDim.y) - 1 ? kMaxRows - row : 1; + for (int32_t r = 0; r < rows; ++r) + { + epochWord[r * kSourcesPerRank] = epochRead + 1; + } + } + }; + + if (rowLen <= kTopK) + { + // Every rank agrees on this branch: identity selection, no exchange for this row. + if (src == 0) + { + for (int32_t j = threadIdx.x; j < kTopK; j += kThreads) + { + outRow[j] = j < rowLen ? j : -1; + } + } + if (exchange) + { + fetchResetHeader(p, s, prevBufRead, resetOffset); + resetSlots(p, s, prevBufRead); + markPhase(p, 9); + } + __syncthreads(); + advanceEpoch(); + clusterInitWait(); + } + else + { + int32_t const localLen = min(localLength(rowLen, rank, world), p.localCapacity); + float const* rowLogits = p.logits + static_cast(row) * p.logitsStride; + + // ---- Round 1: this source's local top keys --------------------------------------------------------- + int32_t const n = loadSource(s, rowLogits, localLen, src, true); + if (threadIdx.x == 0) + { + epochShared = epochRead; + } + // The reset warps start fetching the headers of their share of the previous exchange launch's buffer (an + // exchange row implies an exchange launch); the values are first used by the reset after the send. + fetchResetHeader(p, s, prevBufRead, resetOffset); + __syncthreads(); + markPhase(p, 1); + advanceEpoch(); + uint32_t const cap = static_cast(roundCap(world)); + uint32_t const need = static_cast(roundNeed(world)); + SelectOut sel{0u, 0u, 0u, static_cast(n), true, false, 0}; + uint32_t selLo = 0u; + int32_t passInfo = 0; + if (n > static_cast(cap)) + { + // The cut lies among the keys of one sign (see kPositiveKey); the digits span that sign's key range. + SignRange const r = signRange(readStats(s.stats), need); + selLo = r.lo; + SmemKeys const keys{s.chunk, nullptr}; + sel = blockSelect(s, keys, n, r.lo, r.hi, need, cap, r.lo, r.bits, debugSlots(p, 32), r.aboveHi); + passInfo = sel.passes + (r.negative ? 10 : 0); + } + markPhase(p, 2); + int32_t const sendBuf = epochShared % kNumLamportBuffers; + uint32_t const sent = sendSource(s, sel, n, selLo, uint64_t{kEmptyWord}, cap, + p.mcWorkspace + round1Offset(sendBuf, rank * kSourcesPerRank + src, row), src, rank, world, + reinterpret_cast(s.mkey), debugSlots(p, 0)); + putDebug(debugSlots(p, 0), 46, sent); + putDebug(debugSlots(p, 0), 47, passInfo); + markPhase(p, 3); + + // The reset warps reset the previous exchange launch's buffer, whose candidates every rank has consumed (a + // peer writes it again only two exchange launches later, after this launch completed on this rank), while + // the poll and header warps wait for this launch's candidates. + resetSlots(p, s, (epochShared + kNumLamportBuffers - 1) % kNumLamportBuffers); + int32_t const buf = epochShared % kNumLamportBuffers; + if (threadIdx.x == kFirstResetWarp * 32) + { + int64_t* const d = debugSlots(p, 0); + if (d != nullptr) + { + d[9] = static_cast(globalTimerNs()); + } + } + + // ---- Merge: the row's cluster ----------------------------------------------------------------------- + WaitCtx const wait{p.state, p.timeoutNs, rank}; + if (clusterMergeRow(p, s, rec, tally, buf, row, rowLen, outRow, wait) && src == 0) + { + // A source withheld keys above the threshold: crossRankTopk's merge, with its second round, on CTA 0. + mergeRow(p, s, segSlots, segBegin, buf, row, rowLen, localLen, rowLogits, outRow, wait); + int64_t* const dbg = debugSlots(p, 0); + putDebug(dbg, 63, dbg != nullptr ? dbg[63] + 100 : 0); + } + markPhase(p, 7); + } + markPhase(p, 8); +} + +#endif + +} // namespace + +size_t cpTopkSmemBytes(int32_t localCapacity, int32_t worldSize) +{ + return smemBytes(localCapacity, worldSize); +} + +int32_t cpTopkMaxLocalCapacity(int32_t worldSize) +{ + static int32_t cached[kMaxRanks + 1] = {}; + if (worldSize < 1 || worldSize > kMaxRanks) + { + return 0; + } + if (cached[worldSize] == 0) + { + int device = 0; + cudaGetDevice(&device); + int maxSmemOptin = 0; + cudaDeviceGetAttribute(&maxSmemOptin, cudaDevAttrMaxSharedMemoryPerBlockOptin, device); + int32_t cap = 0; + for (int32_t c = 1 << 22; c >= kPageTokens * kSourcesPerRank; c -= kPageTokens * kSourcesPerRank) + { + if (smemBytes(c, worldSize) + kStaticSmemBytes <= static_cast(maxSmemOptin)) + { + cap = c; + break; + } + } + cached[worldSize] = cap; + } + return cached[worldSize]; +} + +int32_t cpTopkMaxActiveClusters(int32_t worldSize) +{ + // Queried once per world size and device (the value + 1; 0 = not queried yet). + constexpr int32_t kMaxDevices = 64; + static int32_t cached[kMaxDevices][kMaxRanks + 1] = {}; + if (worldSize < 1 || worldSize > kMaxRanks) + { + return 0; + } + int device = 0; + cudaGetDevice(&device); + if (device >= 0 && device < kMaxDevices && cached[device][worldSize] > 0) + { + return cached[device][worldSize] - 1; + } + int32_t const capacity = cpTopkMaxLocalCapacity(worldSize); + if (capacity == 0) + { + return 0; + } + size_t const smem = smemBytes(capacity, worldSize); + if (cudaFuncSetAttribute(cpTopkZ9Kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(smem)) + != cudaSuccess) + { + (void) cudaGetLastError(); + return 0; + } + cudaLaunchConfig_t config{}; + config.gridDim = dim3(kSourcesPerRank, kMaxRows); + config.blockDim = dim3(kThreads); + config.dynamicSmemBytes = smem; + cudaLaunchAttribute attr{}; + attr.id = cudaLaunchAttributeClusterDimension; + attr.val.clusterDim.x = kSourcesPerRank; + attr.val.clusterDim.y = 1; + attr.val.clusterDim.z = 1; + config.attrs = &attr; + config.numAttrs = 1; + int clusters = 0; + if (cudaOccupancyMaxActiveClusters(&clusters, cpTopkZ9Kernel, &config) != cudaSuccess) + { + (void) cudaGetLastError(); + return 0; + } + if (device >= 0 && device < kMaxDevices) + { + cached[device][worldSize] = clusters + 1; + } + return clusters; +} + +int64_t cpTopkDefaultTimeoutMs() +{ + static int64_t const timeoutMs = []() + { + char const* env = std::getenv("TRTLLM_STAIRCASE_GLM_LL_CROSS_RANK_TOPK_TIMEOUT_MS"); + if (env == nullptr || env[0] == '\0') + { + return kDefaultTimeoutMs; + } + char* end = nullptr; + long long const v = std::strtoll(env, &end, 10); + if (end == env || *end != '\0' || v <= 0) + { + throw std::invalid_argument(std::string("glm_moe_dsa_cross_rank_topk_cluster_merge: bad " + "TRTLLM_STAIRCASE_GLM_LL_CROSS_RANK_TOPK_TIMEOUT_MS=") + + env); + } + return static_cast(v); + }(); + return timeoutMs; +} + +void invokeCpTopk(CpTopkParams const& inParams, cudaStream_t stream) +{ + if (inParams.numRows <= 0) + { + return; + } + CpTopkParams params = inParams; + if (params.timeoutNs == 0) + { + params.timeoutNs = static_cast(cpTopkDefaultTimeoutMs()) * 1000000ull; + } + if (params.numRows > kMaxRows || params.worldSize <= 0 || params.worldSize > kMaxRanks || params.rank < 0 + || params.rank >= params.worldSize) + { + throw std::invalid_argument("glm_moe_dsa_cross_rank_topk_cluster_merge: unsupported rows/world/rank"); + } + size_t const smem = smemBytes(params.localCapacity, params.worldSize); + // The attribute is per function and device; raise it only when a larger footprint is requested. + static thread_local size_t configuredSmem[64] = {}; + int device = 0; + cudaGetDevice(&device); + cudaError_t err = cudaSuccess; + if (device >= 64 || configuredSmem[device] < smem) + { + err = cudaFuncSetAttribute(cpTopkZ9Kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(smem)); + if (err != cudaSuccess) + { + throw std::runtime_error( + std::string("glm_moe_dsa_cross_rank_topk_cluster_merge: smem attribute: ") + cudaGetErrorString(err)); + } + if (device < 64) + { + configuredSmem[device] = smem; + } + } + // Programmatic dependent launch as for the other TensorRT-LLM kernels (on unless TRTLLM_ENABLE_PDL=0). + static bool const enablePdl = []() + { + char const* env = std::getenv("TRTLLM_ENABLE_PDL"); + return !(env != nullptr && env[0] == '0' && env[1] == '\0'); + }(); + cudaLaunchConfig_t config{}; + config.gridDim = dim3(kSourcesPerRank, params.numRows); + config.blockDim = dim3(kThreads); + config.dynamicSmemBytes = smem; + config.stream = stream; + cudaLaunchAttribute attrs[2]; + attrs[0].id = cudaLaunchAttributeClusterDimension; + attrs[0].val.clusterDim.x = kSourcesPerRank; + attrs[0].val.clusterDim.y = 1; + attrs[0].val.clusterDim.z = 1; + attrs[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[1].val.programmaticStreamSerializationAllowed = enablePdl ? 1 : 0; + config.attrs = attrs; + config.numAttrs = 2; + err = cudaLaunchKernelEx(&config, cpTopkZ9Kernel, params); + if (err == cudaSuccess) + { + err = cudaGetLastError(); + } + if (err != cudaSuccess) + { + throw std::runtime_error( + std::string("glm_moe_dsa_cross_rank_topk_cluster_merge: launch: ") + cudaGetErrorString(err)); + } +} + +} // namespace kernels::glm_moe_dsa::cross_rank_topk_cluster_merge + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopkClusterMerge.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopkClusterMerge.h new file mode 100644 index 000000000000..0f6d53a52e92 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopkClusterMerge.h @@ -0,0 +1,191 @@ +/* + * 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 +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa::cross_rank_topk_cluster_merge +{ + +// Context-partitioned DSA indexer top-k with a cluster merge. +// +// It selects exactly what the kernel of glmMoeDsa/crossRankTopk.h selects, in the same output order, with the +// same context partition, exchange sources, round-1 / round-2 protocol and Lamport workspace layout: +// - global 64-token page p belongs to rank p % worldSize; within a rank, each verify row is handled by +// kSourcesPerRank CTAs, CTA s owning the 8-token block (s + page) % kSourcesPerRank of every local page +// (exchange source q = rank * kSourcesPerRank + s); +// - each source selects its local top keys with crossRankTopk's rule and broadcasts them as (order-preserving key, +// global token) pairs in source order, plus a header (count, exclusive bound of the withheld keys; this kernel adds +// the key statistics of the keys sent), through the multicast Lamport workspace. +// What differs is how the row is merged and how a launch ends: +// - the kSourcesPerRank CTAs of one row form a thread-block cluster. CTA s merges the candidates of sources +// [s * worldSize, (s + 1) * worldSize) (at 8 ranks: the 8 sources of rank s). The cluster runs crossRankTopk's +// merge select (the same digits and early-stop rules) on the cluster-summed histograms, and every CTA writes its +// selected candidates directly into the output row at crossRankTopk's positions (keys above the threshold in +// candidate order, then ties in candidate order, then -1). Each round-1 slot header carries its source's key +// statistics, so every CTA derives the row's statistics from the headers it polls; the CTAs exchange only the +// select's histogram totals, with st.async stores into each other's shared memory completed on transaction +// barriers, and derive every CTA's output counts from the per-CTA tallies of the select. The only cluster-wide +// barrier orders the barriers' initialization before the first exchange. A row in which a source withheld keys +// above the threshold (identical decision on every rank) is merged instead by CTA 0 with crossRankTopk's single-CTA +// merge, including its second round; +// - every CTA keeps its own epoch word (one per (row, source CTA) slot): its thread 0 reads the word after the grid +// dependency wait (only the previous exchange launch of this op writes it), the value is published to the CTA at +// the block barrier that follows the key loads, and the CTA advances the word with a plain store after that +// barrier (nothing else in the launch reads it; the next launch reads it after its own grid dependency wait); the +// CTAs of the last row also advance the words of the rows the launch does not have, so every word counts the +// exchange launches. There is no arrival counter and no end-of-kernel fence. +// +// Rows of at most kTopK valid tokens produce the identity selection followed by -1, without an exchange. A launch whose +// rows all have at most kTopK valid tokens (the same on every rank: the lengths come from the shared KV length) is an +// identity launch: it writes neither the workspace nor the epoch words. Every other launch is an exchange launch. +// Exchange launches rotate kNumLamportBuffers buffers by the epoch, which counts exchange launches; in exchange launch +// x every CTA resets its share of its own rank's copy of the buffer of exchange launch x - 1 between its send and its +// merge (the headers of its share are fetched once the epoch is known, after the grid dependency wait). Every rank's +// merges of x - 1 have read their copy of it, and a peer writes it again only in exchange launch x + 2, which starts +// after the peer's merges of x + 1 received this rank's sources of x + 1, i.e. after this rank completed launch x +// (identity launches in between neither write nor wait, so they do not shorten this chain). Every global memory +// access of the kernel follows its grid dependency wait, which orders it after the completion (and memory +// visibility) of the preceding kernels of the stream, including the previous launch of this op; the kernel never +// triggers its dependents early. +// +// A merge that waits longer than the timeout (CpTopkParams::timeoutNs, %globaltimer) for a peer's words records the +// wait in the state's error word and aborts the kernel with a trap: the exchange cannot complete correctly any more, +// and a late peer write could land in a buffer this rank reuses, so the run fails instead of continuing. + +//! Selected indices per verify row (DSA index_topk). +constexpr int32_t kTopK = 2048; +//! Ranks of the context-parallel group. +constexpr int32_t kMaxRanks = 8; +//! Verify rows per call. +constexpr int32_t kMaxRows = 8; +//! Tokens per indexer K page; the context is partitioned by whole pages. +constexpr int32_t kPageTokens = 64; +//! CTAs (exchange sources) per rank and verify row; also the cluster size (one cluster per row). +constexpr int32_t kSourcesPerRank = 8; +//! Exchange sources per row. +constexpr int32_t kMaxSources = kMaxRanks * kSourcesPerRank; +//! Round-1 slots per (source, row). +constexpr int32_t kSourceCapacity = 256; +//! Keys each saturated source aims to send in round 1 at 8 ranks (crossRankTopk's rule: three times the even share of +//! kTopK; refining stops at the first bin edge with at most four times the even share of keys at or above it). +constexpr int32_t kSourceBudget = 96; +//! Round-2 slots per (source, row). +constexpr int32_t kRound2Capacity = 2048; +//! Threads per CTA. +constexpr int32_t kThreads = 1024; +//! Lamport buffers rotated by the epoch. +constexpr int32_t kNumLamportBuffers = 3; +//! 32-bit word that marks an empty Lamport slot; never a valid key, token index, count or bound. +constexpr uint32_t kEmptyWord = 0xFFFFFFFFu; + +//! 32-bit words of one (source, row) slot: a 32-byte header then two words per candidate, so candidates can be written +//! and read with 16-byte vector accesses. Round-1 header: count, exclusive bound of the withheld keys, then the key +//! statistics of the keys sent (largest key, smallest key + 1, keys >= +0.0, largest key below +0.0, smallest key +//! >= +0.0 plus 1; the "+ 1" maps the "none" value 0xFFFFFFFF to 0), then 0. Round-2 header: count, bound. +constexpr int64_t kSlotHeaderWords = 8; +constexpr int64_t kRound1SlotWords = kSlotHeaderWords + 2 * int64_t{kSourceCapacity}; +constexpr int64_t kRound2SlotWords = kSlotHeaderWords + 2 * int64_t{kRound2Capacity}; +//! 32-bit words of one Lamport buffer (round-1 region, then round-2 region). +constexpr int64_t kBufferWords = int64_t{kMaxSources} * kMaxRows * (kRound1SlotWords + kRound2SlotWords); +//! 32-bit words of the whole exchange workspace; every word must start as kEmptyWord. +constexpr int64_t kWorkspaceWords = kNumLamportBuffers * kBufferWords; +//! int32 words of the per-rank device state, zero-initialized: [0, 2) unused, [2] second rounds taken, [3] error word +//! (nonzero = the wait that timed out, written right before the kernel aborts), then one epoch word per +//! (row, source CTA) slot: state[kEpochBase + row * kSourcesPerRank + s]. +constexpr int32_t kEpochBase = 4; +constexpr int32_t kStateWords = kEpochBase + kMaxRows * kSourcesPerRank; +//! Default wait for a peer's Lamport words before the kernel aborts +//! (TRTLLM_STAIRCASE_GLM_LL_CROSS_RANK_TOPK_TIMEOUT_MS overrides it). +constexpr int64_t kDefaultTimeoutMs = 120000; + +//! Values of the error word. +enum CpTopkWait : int32_t +{ + kWaitRound1Header = 1, + kWaitRound1Candidates = 2, + kWaitRound2Header = 3, + kWaitRound2Candidates = 4 +}; + +struct CpTopkParams +{ + //! [numRows, logitsStride] fp32 logits of this rank's local columns. + float const* logits; + int64_t logitsStride; + //! Local columns the logits view holds (each row's local length is clamped to it). + int32_t localCapacity; + //! [1] int32 global indexer KV length of the request; row r has kvLen - nextN + r + 1 valid tokens. + int32_t const* kvLens; + //! [numRows, kTopK] int32 output: global token indices, -1 for unused slots. + int32_t* out; + int32_t numRows; + int32_t nextN; + int32_t rank; + int32_t worldSize; + //! Multicast address of the exchange workspace (kWorkspaceWords words, identical layout on every rank). + uint32_t* mcWorkspace; + //! This rank's unicast address of the same workspace. + uint32_t* ucWorkspace; + //! [kStateWords] int32 per-rank device state, zero-initialized. + int32_t* state; + //! Optional [numRows, kSourcesPerRank, kDebugSlots] int64 phase timestamps (diagnosis only; nullptr = off). + int64_t* debugTimes = nullptr; + //! Longest wait for a peer's words before the kernel aborts (ns); 0 = the default (kDefaultTimeoutMs, or + //! TRTLLM_STAIRCASE_GLM_LL_CROSS_RANK_TOPK_TIMEOUT_MS). + uint64_t timeoutNs = 0; +}; + +//! Timestamps per CTA in CpTopkParams::debugTimes: [0, 16) %globaltimer (ns) and [16, 32) clock64 per phase (0 start, +//! 10 grid dependency wait returned, 1 keys loaded, 2 local select, 3 sent, 9 previous buffer reset (%globaltimer of +//! the first reset warp only), 4 candidates polled, 5 cluster statistics agreed, 6 merge select, 7 row written, 8 end), +//! [32, 42) local select (keys, then clock64 at the start of / after the histogram of / after the bin search of each +//! pass), [42, 45) send (start, keys taken, stores issued), 46 round-1 keys sent by this source, 47 local select passes +//! (+10 when the cut was searched among the negative keys), [48, 58) merge select (keys = this CTA's candidates; per +//! pass clock64 at its start / after the group totals arrived / after the bin was decided), clock64 of the merge: 45 +//! candidates polled, 61 statistics combined, 58 emission positions scanned, 60 row written; 62 candidates of the row +//! (all CTAs), 63 merge select passes (+10 likewise; +100 when the row took crossRankTopk's single-CTA merge, which +//! then reuses 45 and 58-61 as crossRankTopk does); clock64 and values of thread 0. +constexpr int32_t kDebugSlots = 64; + +//! Effective wait timeout of a launch with params.timeoutNs == 0 (ms). +int64_t cpTopkDefaultTimeoutMs(); + +//! Dynamic shared memory the kernel needs (bytes). +size_t cpTopkSmemBytes(int32_t localCapacity, int32_t worldSize); + +//! Largest local capacity the kernel accepts for a world size. +int32_t cpTopkMaxLocalCapacity(int32_t worldSize); + +//! Clusters of the kernel's launch configuration (largest local capacity) that can be resident at once on this +//! device (cudaOccupancyMaxActiveClusters, queried once per world size and device); 0 when the query fails. No row +//! waits for another row of its rank, but a row's cluster waits for the same row's sources on every peer rank, and +//! the ranks may make their rows resident in different orders; the caller therefore requires kMaxRows co-resident +//! clusters, so every cluster of a launch can be resident at once (other kernels sharing the SMs are finite and do +//! not wait for this one). +int32_t cpTopkMaxActiveClusters(int32_t worldSize); + +void invokeCpTopk(CpTopkParams const& params, cudaStream_t stream); + +} // namespace kernels::glm_moe_dsa::cross_rank_topk_cluster_merge + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopkGlobalEmit.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopkGlobalEmit.cu new file mode 100644 index 000000000000..685a2da67a97 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopkGlobalEmit.cu @@ -0,0 +1,2757 @@ +/* + * 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/glmMoeDsa/crossRankTopkGlobalEmit.h" + +#include + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa::cross_rank_topk_global_emit +{ +namespace +{ + +constexpr int32_t kBins = 2048; +static_assert(kBins >= kTopK, "the output row is staged in the histogram"); +constexpr int32_t kWarps = kThreads / 32; +//! Per-warp key statistics kept in shared memory (see KeyStats). +constexpr int32_t kNumKeyStats = 5; +//! Tokens of each page one source holds. +constexpr int32_t kBlockTokens = kPageTokens / kSourcesPerRank; +//! CTAs of one row's cluster. +constexpr int32_t kClusterCtas = kSourcesPerRank; +//! Histogram bins per group of the cluster bin search (one warp lane per bin). +constexpr int32_t kGroupBins = 32; +constexpr int32_t kMaxGroups = kBins / kGroupBins; +//! Static shared memory of the kernel (round-2 segment tables and the cluster exchange records), kept out of the +//! dynamic budget. +constexpr size_t kStaticSmemBytes = 8192; + +// Records the CTAs of a row's cluster send each other with st.async (static, so every CTA has them at the same +// address). Exchange x of a row completes on barrier bar[x % 2], phase x / 2; the record sets alternate by pass +// parity and are written again only after every CTA has completed the exchange in between. +struct ClusterRecords +{ + //! Transaction barriers of the exchanges (arrival count 1: this CTA's expect-tx). + uint64_t bar[2]; + //! [pass parity][source CTA][group] histogram totals per group of kGroupBins bins. + uint32_t groups[2][kClusterCtas][kMaxGroups]; + //! [pass parity][source CTA][bin] counts of the selected group's bins, ascending. + uint32_t bins[2][kClusterCtas][kGroupBins]; +}; + +// Scalar slots in shared memory. +enum ScalarSlot : int32_t +{ + kSlotBin = 0, + kSlotAbove, + kSlotBinCount, + kSlotTotal, + kSlotCounter2, + kSlotSrcCount, // [kMaxSources] + kSlotSrcBound = kSlotSrcCount + kMaxSources, // [kMaxSources] + kSlotSrcLen = kSlotSrcBound + kMaxSources, // [kMaxSources] + kSlotSrcOffset = kSlotSrcLen + kMaxSources, // [kMaxSources + 1] + kSlotCountBefore = kSlotSrcOffset + kMaxSources + 1, // candidates of the cluster CTAs before this one + kSlotCountMine, // candidates of this cluster CTA + kSlotCountTotal, // candidates of the row + kSlotPositivesBefore, // candidates >= +0.0 of the cluster CTAs before this one + kNumScalarSlots +}; + +__host__ __device__ inline int32_t numSources(int32_t world) +{ + return world * kSourcesPerRank; +} + +// Even share of kTopK per source. +__host__ __device__ inline int32_t evenShare(int32_t world) +{ + return (kTopK + numSources(world) - 1) / numSources(world); +} + +// Most keys a source sends in round 1 (an even number). +__host__ __device__ inline int32_t roundCap(int32_t world) +{ + return min(kSourceCapacity, 4 * evenShare(world)); +} + +// Keys one source can hold: one block of every local page. +__host__ __device__ inline int32_t chunkCapacity(int32_t localCapacity) +{ + int32_t const pages = (localCapacity + kPageTokens - 1) / kPageTokens; + return pages * kBlockTokens; +} + +__host__ __device__ inline int32_t mergeCapacity(int32_t world) +{ + return numSources(world) * roundCap(world); +} + +__host__ __device__ inline size_t smemBytes(int32_t localCapacity, int32_t world) +{ + size_t words = 2 * size_t{kBins} + 2 * size_t{kWarps} + size_t{kNumKeyStats} * kWarps + size_t{kNumScalarSlots}; + words = (words + 3) / 4 * 4; + words += static_cast(chunkCapacity(localCapacity)) + 2 * static_cast(mergeCapacity(world)); + return words * sizeof(uint32_t); +} + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 900) +// The kernel targets SM90+ (multicast Lamport exchange, thread-block clusters). +__global__ void __launch_bounds__(kThreads, 1) cpTopkZ11Kernel(CpTopkParams p) +{ + __trap(); +} +#else + +//! Bits of one radix-select digit (kBins bins). +constexpr int32_t kDigitBits = 11; +//! Order-preserving key of +0.0: the keys of non-negative values are at least this, those of negative values below. +//! Real indexer logits are mostly negative, and the cut of a row (and of many sources) often lies among small +//! negative values while the largest keys are positive: one key range across zero spans ~2^31 key units, where a +//! radix digit has no resolution. The selects therefore work on one sign at a time: with fewer than `need` keys at +//! or above kPositiveKey these are all taken and the cut is searched among the negative keys only, else among the +//! non-negative keys only; each range is narrow enough for one or two passes. +constexpr uint32_t kPositiveKey = 0x80000000u; +//! Independent 16-byte candidate loads one merge thread issues before it waits on any of them. +constexpr int32_t kPollBatch = 4; +//! Polls of a still-empty Lamport word between two reads of the global timer. +constexpr int32_t kPollsPerTimerRead = 256; +//! Warps of the cluster merge: warps [0, kPollWarps) poll the round-1 candidates of this CTA's sources (one 16-byte +//! vector per thread: at most 520 vectors per CTA, at world sizes 5 to 7), the next kHeaderWarps warps poll the +//! round-1 headers of every source (one per thread), and the remaining kResetWarps warps reset the previous exchange +//! launch's buffer meanwhile. +constexpr int32_t kPollWarps = 17; +constexpr int32_t kHeaderWarps = 2; +constexpr int32_t kFirstResetWarp = kPollWarps + kHeaderWarps; +constexpr int32_t kResetWarps = kWarps - kFirstResetWarp; +static_assert(kHeaderWarps * 32 >= kMaxSources, "one header per header-warp thread"); + +// Keys a saturated source aims to send in round 1 (kSourceBudget at 8 ranks). +__device__ __forceinline__ int32_t roundNeed(int32_t world) +{ + return min(kSourceCapacity, 3 * evenShare(world)); +} + +// Tokens of `rank` among the global tokens [0, n). +__host__ __device__ __forceinline__ int32_t localLength(int32_t n, int32_t rank, int32_t world) +{ + if (n <= 0) + { + return 0; + } + int32_t const fullPages = n / kPageTokens; + int32_t const rem = n % kPageTokens; + int32_t const ownedFull = (fullPages - rank + world - 1) / world; + int32_t len = ownedFull * kPageTokens; + if (rem > 0 && fullPages % world == rank) + { + len += rem; + } + return len; +} + +// Source `src` of a rank holds one block of kBlockTokens consecutive tokens of every local page: block +// (src + page) % kSourcesPerRank. A fully selected page is thus spread over all sources of its rank. +// Keys of source `src` among `localLen` local columns (only the last local page can be partial). +__host__ __device__ __forceinline__ int32_t sourceLength(int32_t localLen, int32_t src) +{ + int32_t const pages = (localLen + kPageTokens - 1) / kPageTokens; + if (pages == 0) + { + return 0; + } + int32_t const last = pages - 1; + int32_t const tail = localLen - last * kPageTokens - ((src + last) % kSourcesPerRank) * kBlockTokens; + return last * kBlockTokens + min(max(tail, 0), kBlockTokens); +} + +// Global token of a local column under the page-interleaved partition. +__device__ __forceinline__ uint32_t localToGlobal(int32_t col, int32_t rank, int32_t world) +{ + int32_t const page = col / kPageTokens; + return static_cast((page * world + rank) * kPageTokens + (col % kPageTokens)); +} + +struct SmemLayout +{ + uint32_t* hist; // [2 * kBins] radix histograms of alternate passes; the first one stages the output row + uint32_t* warpSums; // [2 * kWarps] + uint32_t* stats; // [kNumKeyStats * kWarps] per-warp key statistics + uint32_t* scal; // [kNumScalarSlots] + uint32_t* chunk; // [chunkCapacity] order-preserving keys of this source + uint32_t* mkey; // [mergeCapacity] staged round-1 pairs, then the merge candidates (merge CTA) + uint32_t* midx; // [mergeCapacity] +}; + +__device__ inline SmemLayout carveSmem(uint32_t* base, int32_t localCapacity, int32_t world) +{ + SmemLayout s; + uint32_t* p = base; + s.hist = p; + p += 2 * kBins; + s.warpSums = p; + p += 2 * kWarps; + s.stats = p; + p += kNumKeyStats * kWarps; + s.scal = p; + p += kNumScalarSlots; + p = base + ((p - base) + 3) / 4 * 4; + s.chunk = p; + p += chunkCapacity(localCapacity); + s.mkey = p; + p += mergeCapacity(world); + s.midx = p; + return s; +} + +__device__ __forceinline__ uint32_t orderedKey(float v) +{ + uint32_t const b = __float_as_uint(v); + uint32_t const k = (b & 0x80000000u) ? ~b : (b | 0x80000000u); + // NaN keys above +inf are clamped so that no key, and no exclusive bound key + 1, equals the empty marker. + return k < kEmptyWord - 1u ? k : kEmptyWord - 2u; +} + +// Bits of x (0 for 0): the digit range of keys in [base, base + x]. +__device__ __forceinline__ int32_t bitLength(uint32_t x) +{ + return x == 0u ? 0 : 32 - __clz(x); +} + +// Statistics of a set of keys: extremes overall and per sign, and the count of non-negative keys. +struct KeyStats +{ + uint32_t maxKey = 0u; // largest key (0 when none) + uint32_t minKey = kEmptyWord; // smallest key (kEmptyWord when none) + uint32_t positives = 0u; // keys >= kPositiveKey + uint32_t maxNeg = 0u; // largest key below kPositiveKey (0 when none) + uint32_t minPos = kEmptyWord; // smallest key >= kPositiveKey (kEmptyWord when none) + + __device__ __forceinline__ void add(uint32_t k) + { + bool const pos = k >= kPositiveKey; + maxKey = max(maxKey, k); + minKey = min(minKey, k); + positives += pos ? 1u : 0u; + maxNeg = pos ? maxNeg : max(maxNeg, k); + minPos = pos ? min(minPos, k) : minPos; + } +}; + +// Reduce a thread's statistics over its warp; lane 0 stores them in the per-warp array (read with readStats after +// the next block barrier). +__device__ __forceinline__ void storeWarpStats(KeyStats const& t, uint32_t* stats) +{ + uint32_t const v0 = __reduce_max_sync(0xFFFFFFFFu, t.maxKey); + uint32_t const v1 = __reduce_min_sync(0xFFFFFFFFu, t.minKey); + uint32_t const v2 = __reduce_add_sync(0xFFFFFFFFu, t.positives); + uint32_t const v3 = __reduce_max_sync(0xFFFFFFFFu, t.maxNeg); + uint32_t const v4 = __reduce_min_sync(0xFFFFFFFFu, t.minPos); + if (threadIdx.x % 32 == 0) + { + int32_t const w = threadIdx.x / 32; + stats[w] = v0; + stats[kWarps + w] = v1; + stats[2 * kWarps + w] = v2; + stats[3 * kWarps + w] = v3; + stats[4 * kWarps + w] = v4; + } +} + +// Block statistics from the per-warp array (every warp reduces it; no barrier). +__device__ __forceinline__ KeyStats readStats(uint32_t const* stats) +{ + int32_t const lane = threadIdx.x % 32; + KeyStats b; + b.maxKey = __reduce_max_sync(0xFFFFFFFFu, stats[lane]); + b.minKey = __reduce_min_sync(0xFFFFFFFFu, stats[kWarps + lane]); + b.positives = __reduce_add_sync(0xFFFFFFFFu, stats[2 * kWarps + lane]); + b.maxNeg = __reduce_max_sync(0xFFFFFFFFu, stats[3 * kWarps + lane]); + b.minPos = __reduce_min_sync(0xFFFFFFFFu, stats[4 * kWarps + lane]); + return b; +} + +// Digit setup of a select of the need-th largest key: over the non-negative keys when there are at least `need` of +// them, else over the negative keys with every non-negative key counted as taken (`aboveHi`). +struct SignRange +{ + uint32_t lo; // smallest key of the range (also the digit base) + uint64_t hi; // exclusive end of the histogrammed range + int32_t bits; // digit range of the range + uint32_t aboveHi; // keys at or above hi, all taken + bool negative; // the cut is among the negative keys +}; + +__device__ __forceinline__ SignRange signRange(KeyStats const& st, uint32_t need) +{ + if (st.positives >= need || st.positives == 0u) + { + uint32_t const lo = st.positives >= need ? st.minPos : st.minKey; + return SignRange{lo, uint64_t{kEmptyWord}, bitLength(st.maxKey - lo), 0u, st.positives == 0u}; + } + return SignRange{st.minKey, uint64_t{kPositiveKey}, bitLength(st.maxNeg - st.minKey), st.positives, true}; +} + +__device__ __forceinline__ int64_t round1Offset(int32_t buf, int32_t src, int32_t row) +{ + return int64_t{buf} * kBufferWords + (int64_t{src} * kMaxRows + row) * kRound1SlotWords; +} + +__device__ __forceinline__ int64_t round2Offset(int32_t buf, int32_t src, int32_t row) +{ + return int64_t{buf} * kBufferWords + int64_t{kMaxSources} * kMaxRows * kRound1SlotWords + + (int64_t{src} * kMaxRows + row) * kRound2SlotWords; +} + +__device__ __forceinline__ uint2 loadVolatile2(uint32_t const* ptr) +{ + uint2 v; + asm volatile("ld.volatile.global.v2.u32 {%0, %1}, [%2];" : "=r"(v.x), "=r"(v.y) : "l"(ptr) : "memory"); + return v; +} + +__device__ __forceinline__ uint4 loadVolatile4(uint32_t const* ptr) +{ + uint4 v; + asm volatile("ld.volatile.global.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(ptr) + : "memory"); + return v; +} + +__device__ __forceinline__ void store2(uint32_t* ptr, uint32_t a, uint32_t b) +{ + asm volatile("st.global.v2.u32 [%0], {%1, %2};" ::"l"(ptr), "r"(a), "r"(b) : "memory"); +} + +__device__ __forceinline__ void store4(uint32_t* ptr, uint4 v) +{ + asm volatile("st.global.v4.u32 [%0], {%1, %2, %3, %4};" ::"l"(ptr), "r"(v.x), "r"(v.y), "r"(v.z), "r"(v.w) + : "memory"); +} + +__device__ __forceinline__ bool isEmpty(uint2 v) +{ + return v.x == kEmptyWord || v.y == kEmptyWord; +} + +// Whether the first pair of a 16-byte candidate vector, and the second one when `second`, were written. +__device__ __forceinline__ bool vectorReady(uint4 v, bool second) +{ + return v.x != kEmptyWord && v.y != kEmptyWord && (!second || (v.z != kEmptyWord && v.w != kEmptyWord)); +} + +__device__ __forceinline__ uint64_t globalTimerNs() +{ + uint64_t t; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t)); + return t; +} + +// What a polling thread needs to give up: the error word, the timeout and this rank (for the report). +struct WaitCtx +{ + int32_t* state; + uint64_t timeoutNs; + int32_t rank; +}; + +// A peer did not deliver within the timeout. The selection cannot be completed correctly any more, and the peer's +// late writes could land in a buffer this rank resets or reuses, so record which wait failed and abort the kernel: +// the resulting CUDA error fails the run instead of letting any rank continue with a partial selection. +__device__ __noinline__ void abortOnTimeout(WaitCtx const& w, int32_t what) +{ + bool const first = atomicExch(&w.state[3], what) == 0; + __threadfence_system(); + if (first) + { + printf( + "[glm_moe_dsa] crossRankTopkGlobalEmit: rank %d, row %d: wait %d for a peer exceeded %llu ms; aborting the " + "kernel\n", + w.rank, static_cast(blockIdx.y), what, static_cast(w.timeoutNs / 1000000ull)); + } + __trap(); +} + +// Spin until both words of a Lamport pair were written; aborts the kernel after the timeout. +__device__ uint2 pollPair(uint32_t const* ptr, WaitCtx const& w, int32_t what) +{ + uint2 v = loadVolatile2(ptr); + uint64_t start = 0; + for (int32_t polls = 1; isEmpty(v); ++polls) + { + if (polls % kPollsPerTimerRead == 0) + { + uint64_t const now = globalTimerNs(); + if (start == 0) + { + start = now; + } + else if (now - start > w.timeoutNs) + { + abortOnTimeout(w, what); + } + } + v = loadVolatile2(ptr); + } + return v; +} + +// Spin until a 16-byte candidate vector was written (see vectorReady); aborts the kernel after the timeout. +__device__ uint4 pollVector(uint32_t const* ptr, bool second, WaitCtx const& w, int32_t what) +{ + uint4 v = loadVolatile4(ptr); + uint64_t start = 0; + for (int32_t polls = 1; !vectorReady(v, second); ++polls) + { + if (polls % kPollsPerTimerRead == 0) + { + uint64_t const now = globalTimerNs(); + if (start == 0) + { + start = now; + } + else if (now - start > w.timeoutNs) + { + abortOnTimeout(w, what); + } + } + v = loadVolatile4(ptr); + } + return v; +} + +// Block-wide inclusive scan of one value per thread. Returns the inclusive prefix; `total` gets the block sum. +__device__ uint32_t blockInclusiveScan(uint32_t v, uint32_t* warpSums, uint32_t& total) +{ + int32_t const lane = threadIdx.x % 32; + int32_t const warp = threadIdx.x / 32; + uint32_t x = v; +#pragma unroll + for (int32_t o = 1; o < 32; o <<= 1) + { + uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, x, o); + if (lane >= o) + { + x += y; + } + } + if (lane == 31) + { + warpSums[warp] = x; + } + __syncthreads(); + if (warp == 0) + { + uint32_t w = warpSums[lane]; +#pragma unroll + for (int32_t o = 1; o < 32; o <<= 1) + { + uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, w, o); + if (lane >= o) + { + w += y; + } + } + warpSums[kWarps + lane] = w; + } + __syncthreads(); + total = warpSums[2 * kWarps - 1]; + uint32_t const before = warp > 0 ? warpSums[kWarps + warp - 1] : 0u; + return x + before; +} + +// Exclusive block-wide scan with a single barrier: every warp scans the per-warp sums itself. Returns this +// thread's exclusive prefix; `total` gets the block sum. warpSums must not be rewritten before the next barrier. +__device__ uint32_t blockExclusiveScan(uint32_t v, uint32_t* warpSums, uint32_t& total) +{ + int32_t const lane = threadIdx.x % 32; + int32_t const warp = threadIdx.x / 32; + uint32_t x = v; +#pragma unroll + for (int32_t o = 1; o < 32; o <<= 1) + { + uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, x, o); + if (lane >= o) + { + x += y; + } + } + if (lane == 31) + { + warpSums[warp] = x; + } + __syncthreads(); + uint32_t const w = warpSums[lane]; + uint32_t y = w; +#pragma unroll + for (int32_t o = 1; o < 32; o <<= 1) + { + uint32_t const z = __shfl_up_sync(0xFFFFFFFFu, y, o); + if (lane >= o) + { + y += z; + } + } + total = __shfl_sync(0xFFFFFFFFu, y, 31); + return __shfl_sync(0xFFFFFFFFu, y - w, warp) + x - v; +} + +// Find, counting from the top bin, the bin of `hist` that holds the target-th largest key, with two barriers: +// per-warp sums, then each warp scans them and only the warp holding the target scans its own bins. Returns the +// total count (every thread); when target <= total, the bin, the count above it and its own count go to the +// scalar slots. +__device__ uint32_t findBin(SmemLayout const& s, uint32_t const* hist, int32_t nbins, uint32_t target) +{ + int32_t const lane = threadIdx.x % 32; + int32_t const warp = threadIdx.x / 32; + int32_t const perThread = (nbins + kThreads - 1) / kThreads; + int32_t const top = nbins - 1 - perThread * static_cast(threadIdx.x); + uint32_t mine = 0; + for (int32_t i = 0; i < perThread && top - i >= 0; ++i) + { + mine += hist[top - i]; + } + uint32_t const warpTotal = __reduce_add_sync(0xFFFFFFFFu, mine); + if (lane == 0) + { + s.warpSums[warp] = warpTotal; + } + __syncthreads(); + uint32_t const w = s.warpSums[lane]; + uint32_t x = w; +#pragma unroll + for (int32_t o = 1; o < 32; o <<= 1) + { + uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, x, o); + if (lane >= o) + { + x += y; + } + } + uint32_t const total = __shfl_sync(0xFFFFFFFFu, x, 31); + uint32_t const warpBefore = __shfl_sync(0xFFFFFFFFu, x - w, warp); + if (warpBefore < target && target <= warpBefore + warpTotal) + { + uint32_t t = mine; +#pragma unroll + for (int32_t o = 1; o < 32; o <<= 1) + { + uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, t, o); + if (lane >= o) + { + t += y; + } + } + uint32_t acc = warpBefore + t - mine; + if (acc < target && target <= acc + mine) + { + for (int32_t i = 0; i < perThread && top - i >= 0; ++i) + { + uint32_t const c = hist[top - i]; + if (acc + c >= target) + { + s.scal[kSlotBin] = static_cast(top - i); + s.scal[kSlotAbove] = acc; + s.scal[kSlotBinCount] = c; + break; + } + acc += c; + } + } + } + __syncthreads(); + return total; +} + +struct SelectOut +{ + uint32_t edge; // lower edge of the selected bin (the exact key after the last digit) + uint32_t above; // predicate keys strictly above the selected bin + uint32_t binCount; // predicate keys in the selected bin + uint32_t total; // predicate keys + bool all; // take every predicate key (fewer than needed, or all fit the capacity) + bool full; // exact mode: the keys at or above `edge` are exactly the `need` largest + int32_t passes; // radix passes run +}; + +// Candidate keys and token indices held in shared memory (16-byte aligned, readable in groups of four up to the +// next multiple of four). +struct SmemKeys +{ + static constexpr bool kVector = true; + uint32_t const* keys; + uint32_t const* idx; + + __device__ uint4 keys4(int32_t j) const + { + return *reinterpret_cast(keys + j); + } + + __device__ uint4 index4(int32_t j) const + { + return *reinterpret_cast(idx + j); + } +}; + +// Merge candidates in shared memory followed by round-2 candidates read from their global Lamport slots. +struct MergeKeys +{ + static constexpr bool kVector = false; + uint32_t const* smemKeys; + uint32_t const* smemIdx; + int32_t smemCount; + uint32_t const* const* slots; // round-2 slot of each segment + uint32_t const* segBegin; // [numSegments + 1] candidate offsets + int32_t numSegments; + + __device__ uint32_t const* pair(int32_t k) const + { + // Last segment whose begin is <= k. + int32_t lo = 0; + int32_t hi = numSegments - 1; + while (lo < hi) + { + int32_t const mid = (lo + hi + 1) / 2; + if (static_cast(segBegin[mid]) <= k) + { + lo = mid; + } + else + { + hi = mid - 1; + } + } + return slots[lo] + kSlotHeaderWords + 2 * (k - static_cast(segBegin[lo])); + } + + __device__ uint32_t operator()(int32_t j) const + { + return j < smemCount ? smemKeys[j] : pair(j - smemCount)[0]; + } + + __device__ uint32_t index(int32_t j) const + { + return j < smemCount ? smemIdx[j] : pair(j - smemCount)[1]; + } +}; + +// Debug slots of this CTA starting at `offset` (nullptr without a debug buffer). +__device__ __forceinline__ int64_t* debugSlots(CpTopkParams const& p, int32_t offset) +{ + if (p.debugTimes == nullptr) + { + return nullptr; + } + return p.debugTimes + (static_cast(blockIdx.y) * kSourcesPerRank + blockIdx.x) * kDebugSlots + offset; +} + +// clock64 of thread 0 into a debug slot (no-op without a debug buffer). +__device__ __forceinline__ void stamp(int64_t* dbg, int32_t slot) +{ + if (dbg != nullptr && threadIdx.x == 0) + { + dbg[slot] = clock64(); + } +} + +// A value into a debug slot (thread 0; no-op without a debug buffer). +__device__ __forceinline__ void putDebug(int64_t* dbg, int32_t slot, int64_t value) +{ + if (dbg != nullptr && threadIdx.x == 0) + { + dbg[slot] = value; + } +} + +// Phase timestamp for diagnosis (thread 0 of the CTA; no-op without a debug buffer). +__device__ __forceinline__ void markPhase(CpTopkParams const& p, int32_t phase) +{ + if (p.debugTimes != nullptr && threadIdx.x == 0) + { + uint64_t t; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t)); + int64_t* d = p.debugTimes + (static_cast(blockIdx.y) * kSourcesPerRank + blockIdx.x) * kDebugSlots; + d[phase] = static_cast(t); + d[16 + phase] = static_cast(clock64()); + } +} + +// ---- Bridge outputs (topkBridgeKernel's values, written at the emission) -------------------------------------------- + +// Row scalars of the bridge in an exchange row, in static shared memory, from warp kFirstResetWarp: +// - the row's request index before the bias (lane kReqLane's asynchronous copy, started right after the grid +// dependency wait and completed before the block barrier that follows the key loads); +// - the layer ids and the accepted count (asynchronous copies started right after the grid dependency wait and +// completed with the warp's cp.async.wait_all in resetSlots); +// - the row's pool row and copy row (stored by lane 0 after resetSlots); +// the last two groups are read after the block barrier that follows the candidate poll. +struct BridgeShared +{ + int64_t acc64; + int32_t* poolRow; + int32_t* copyRow; + int32_t acc32; + int32_t reqRaw; + int32_t layer[kMaxGroupSize]; +}; + +// Lane of warp kFirstResetWarp that copies the request index. +constexpr int32_t kReqLane = 30; + +// Request of an exchange row from its staged index. +__device__ __forceinline__ int32_t sharedRequest(CpTopkBridgeParams const& b, BridgeShared const& bs) +{ + return bs.reqRaw - b.reqIdxBias; +} + +// Pool base of token `tok` read from the block table (the bridge's rule: -1 for tok < 0 or a page beyond the table). +__device__ __forceinline__ int32_t directBase(CpTopkBridgeParams const& b, int32_t req, int32_t tok) +{ + if (tok < 0) + { + return -1; + } + int32_t const page = tok / kPageTokens; + if (page >= b.maxBlocksPerReq) + { + return -1; + } + return b.blockTable[req * b.btStride0 + page * b.btStride1]; +} + +// Copy row of verify row `row` (-1: the row is not copied). Copy mode 0: rows [0, copyRows) to copyRowStart + row. +// Copy mode 1: request i = row / copyNextN < copyRows copies its last accepted row i * copyNextN + clamp(acc - 1, 0, +// copyNextN - 1) (copyNextN - 1 without accepted counts) to copyRowStart + i; that row is the only one of request i +// that can match, so this is the bridge's selection seen from the source row. +__device__ __forceinline__ int32_t bridgeCopyRow(CpTopkBridgeParams const& b, int32_t row, int64_t acc) +{ + if (b.copyDst == nullptr) + { + return -1; + } + if (b.copyMode == 0) + { + return row < b.copyRows ? b.copyRowStart + row : -1; + } + int32_t const i = row / b.copyNextN; + if (i >= b.copyRows) + { + return -1; + } + int64_t off = b.copyNextN - 1; + if (b.numAccepted != nullptr) + { + off = acc - 1; + off = off < 0 ? 0 : (off > b.copyNextN - 1 ? b.copyNextN - 1 : off); + } + int32_t const srcRow = static_cast(static_cast(i) * b.copyNextN + off); + return srcRow == row ? b.copyRowStart + i : -1; +} + +// Whether verify row `row` needs an accepted count (copy mode 1 with counts, a request below copyRows). +__device__ __forceinline__ bool needsAccepted(CpTopkBridgeParams const& b, int32_t row) +{ + return b.copyDst != nullptr && b.copyMode == 1 && b.numAccepted != nullptr && row / b.copyNextN < b.copyRows; +} + +// Accepted count staged in shared memory (BridgeShared) in the dtype of the counts. +__device__ __forceinline__ int64_t sharedAccepted(CpTopkBridgeParams const& b, BridgeShared const& bs) +{ + return b.numAcceptedIs64 ? bs.acc64 : static_cast(bs.acc32); +} + +// Accepted count of the request of verify row `row`, read from global memory (0 when not needed). +__device__ __forceinline__ int64_t loadAccepted(CpTopkBridgeParams const& b, int32_t row) +{ + if (!needsAccepted(b, row)) + { + return 0; + } + int32_t const i = row / b.copyNextN; + return b.numAcceptedIs64 ? static_cast(b.numAccepted)[i] + : static_cast(static_cast(b.numAccepted)[i]); +} + +// Destinations of one verify row's bridge outputs: the pool row of group member 0 (member g at g * strideG further) +// and the copy row (nullptr: not copied). +struct BridgeRowOut +{ + int32_t* pool; + int64_t strideG; + int32_t* copy; +}; + +__device__ __forceinline__ BridgeRowOut bridgeRowOut(CpTopkBridgeParams const& b, int32_t row, int32_t copyRow) +{ + return BridgeRowOut{ + b.out + row * b.outStride0, b.outStrideG, copyRow >= 0 ? b.copyDst + copyRow * b.copyStride0 : nullptr}; +} + +// The bridge outputs of output position `pos`: for every group member g, the bridge's remapOne(tok, layer g) = +// base * strideFactor + (tok % kPageTokens + layer[g] * kPageTokens), or -1 for tok < 0 or base < 0, with its int32 +// wraparound arithmetic (evaluated in uint32, where the grouping of the sum does not change the bits); then the copy. +template +__device__ __forceinline__ void writeBridge( + CpTopkBridgeParams const& b, BridgeRowOut const& o, uint32_t pos, int32_t tok, int32_t base, LayerAt const& layerAt) +{ + bool const valid = tok >= 0 && base >= 0; + uint32_t const common = static_cast(base) * static_cast(b.strideFactor) + + (static_cast(tok) & static_cast(kPageTokens - 1)); + int32_t* dst = o.pool + pos; + for (int32_t g = 0; g < b.groupSize; ++g) + { + *dst = valid ? static_cast(common + static_cast(layerAt(g)) * kPageTokens) : -1; + dst += o.strideG; + } + if (o.copy != nullptr) + { + o.copy[pos] = tok; + } +} + +// Warp kFirstResetWarp of an exchange row: start the asynchronous copies of the row's request index, layer ids (or +// store the single layer id) and its request's accepted count into `bs` (cp.async, generic proxy; see BridgeShared). +__device__ __forceinline__ void fetchRowScalars(CpTopkParams const& p, BridgeShared& bs, int32_t row) +{ + CpTopkBridgeParams const& b = p.bridge; + int32_t const lane = static_cast(threadIdx.x) % 32; + if (lane < b.groupSize) + { + if (b.layerIds != nullptr) + { + uint32_t const dst = static_cast(__cvta_generic_to_shared(&bs.layer[lane])); + asm volatile("cp.async.ca.shared.global [%0], [%1], 4;" ::"r"(dst), "l"(b.layerIds + lane) : "memory"); + } + else + { + bs.layer[lane] = b.layerId; + } + } + static_assert(kMaxGroupSize < kReqLane && kReqLane < 31, "the request and accepted-count lanes hold no layer id"); + if (lane == kReqLane) + { + // A group of its own, completed (cp.async.wait_group 0) before the block barrier that follows the key loads. + uint32_t const dst = static_cast(__cvta_generic_to_shared(&bs.reqRaw)); + asm volatile("cp.async.ca.shared.global [%0], [%1], 4;" ::"r"(dst), "l"(b.reqIdx + row) : "memory"); + asm volatile("cp.async.commit_group;" ::: "memory"); + } + if (lane == 31 && needsAccepted(b, row)) + { + int32_t const i = row / b.copyNextN; + if (b.numAcceptedIs64) + { + uint32_t const dst = static_cast(__cvta_generic_to_shared(&bs.acc64)); + asm volatile("cp.async.ca.shared.global [%0], [%1], 8;" ::"r"(dst), + "l"(static_cast(b.numAccepted) + i) + : "memory"); + } + else + { + uint32_t const dst = static_cast(__cvta_generic_to_shared(&bs.acc32)); + asm volatile("cp.async.ca.shared.global [%0], [%1], 4;" ::"r"(dst), + "l"(static_cast(b.numAccepted) + i) + : "memory"); + } + } +} + +// Start the copy of candidate token `tok`'s pool base (its block-table entry, 4-byte cp.async) into `dst`, or store +// -1 for tok < 0 or a page beyond the table (the bridge's rule). Only the calling thread reads `dst`, after its +// cp.async.wait_all. +__device__ __forceinline__ void fetchBase(CpTopkBridgeParams const& b, int32_t req, int32_t tok, uint32_t* dst) +{ + int32_t const page = tok / kPageTokens; + if (tok < 0 || page >= b.maxBlocksPerReq) + { + *dst = 0xFFFFFFFFu; + return; + } + uint32_t const addr = static_cast(__cvta_generic_to_shared(dst)); + asm volatile("cp.async.ca.shared.global [%0], [%1], 4;" ::"r"(addr), + "l"(b.blockTable + req * b.btStride0 + page * b.btStride1) + : "memory"); +} + +// Cluster-merge emission of one thread: its (at most two) candidates at pos0 / pos1 and its padding position posPad +// (kTopK: none), ids and bridge outputs (same values as writeBridge): the id row outRow, the pool row of group member +// g at poolRow + g * strideG (kGroup members; 0: groupSize, at most kMaxGroupSize) and the copy row (nullptr: none), +// with the pool bases fetched at the poll and the layer ids in shared memory. +// +// It is a called function (not inlined) so that its instructions are at one address: warp kFirstResetWarp's lane 0 +// runs it once into shared scratch while it waits for the poll (warmEmission), so the emission does not fetch them +// from memory on the critical path. The emission runs once per launch, at its end; when the L2 has been streamed since +// the previous launch (as by a layer's weights in the served step), its instructions would otherwise come from DRAM. +template +__device__ __noinline__ void emitCandidates(int32_t* outRow, int32_t* poolRow, int32_t* copyRow, int64_t strideG, + uint32_t strideFactor, int32_t groupSize, int32_t const* layer, uint32_t pos0, uint32_t pos1, uint32_t posPad, + uint2 tokens, uint2 bases) +{ + uint32_t const want = static_cast(kTopK); + auto emit = [&](uint32_t pos, int32_t tok, int32_t base) + { + outRow[pos] = tok; + bool const valid = tok >= 0 && base >= 0; + uint32_t const common = static_cast(base) * strideFactor + + (static_cast(tok) & static_cast(kPageTokens - 1)); + int32_t* const dst = poolRow + pos; +#pragma unroll + for (int32_t g = 0; g < (kGroup > 0 ? kGroup : kMaxGroupSize); ++g) + { + if (kGroup > 0 || g < groupSize) + { + dst[g * strideG] + = valid ? static_cast(common + static_cast(layer[g]) * kPageTokens) : -1; + } + } + if (copyRow != nullptr) + { + copyRow[pos] = tok; + } + }; + if (pos0 < want) + { + emit(pos0, static_cast(tokens.x), static_cast(bases.x)); + } + if (pos1 < want) + { + emit(pos1, static_cast(tokens.y), static_cast(bases.y)); + } + if (posPad < want) + { + emit(posPad, -1, -1); + } +} + +// emitCandidates for a run-time group size; the served sizes (per-layer remaps and groups of four) are unrolled. +__device__ __forceinline__ void emitCandidatesFor(int32_t groupSize, int32_t* outRow, int32_t* poolRow, + int32_t* copyRow, int64_t strideG, uint32_t strideFactor, int32_t const* layer, uint32_t pos0, uint32_t pos1, + uint32_t posPad, uint2 tokens, uint2 bases) +{ + switch (groupSize) + { + case 1: + emitCandidates<1>( + outRow, poolRow, copyRow, strideG, strideFactor, groupSize, layer, pos0, pos1, posPad, tokens, bases); + break; + case 4: + emitCandidates<4>( + outRow, poolRow, copyRow, strideG, strideFactor, groupSize, layer, pos0, pos1, posPad, tokens, bases); + break; + default: + emitCandidates<0>( + outRow, poolRow, copyRow, strideG, strideFactor, groupSize, layer, pos0, pos1, posPad, tokens, bases); + break; + } +} + +//! Words of the shared scratch of warmEmission: ids [0, 4), pool rows [4, 36) (stride 4), copy row [36, 40). +constexpr int32_t kWarmScratchWords = 40; + +// One thread runs the emission of this launch's group size once into `scratch` (three positions, every store taken +// once, no global access), so that its instructions are cached on the SM before the cluster emission runs. +__device__ __forceinline__ void warmEmission(int32_t groupSize, int32_t* scratch, int32_t const* layer) +{ + emitCandidatesFor(groupSize, scratch, scratch + 4, scratch + 36, 4, 0u, layer, 0u, 1u, 2u, make_uint2(0u, 0u), + make_uint2(~0u, ~0u)); +} + +// Histogram update of one key of a radix pass. +__device__ __forceinline__ void histogramKey(uint32_t* hist, bool inRange, uint32_t k, uint32_t lo, uint64_t hi, + uint32_t base, uint32_t prefMask, uint32_t prefVal, int32_t shift, uint32_t binMask) +{ + uint32_t const d = k - base; + if (inRange && k >= lo && static_cast(k) < hi && (d & prefMask) == prefVal) + { + atomicAdd(&hist[(d >> shift) & binMask], 1u); + } +} + +// Block radix select of the need-th largest of the n keys at or above lo, of which `aboveHi` (< need) lie at or +// above hi and are counted as taken without being histogrammed. Every key in [lo, hi) is at least `base`, and +// key - base < 2^topBits; digits of up to kDigitBits bits of key - base are refined from the top. With cap > 0 the +// select stops at the first bin edge whose keys at or above it number at most `cap`. With cap = 0 it runs to the +// exact key, or stops early when the selected bin is wholly taken (`full`). Alternate passes use alternate +// histograms, so the next one is cleared while the current one is searched. Counts in the result include aboveHi. +template +__device__ SelectOut blockSelect(SmemLayout const& s, KeyAt const& keyAt, int32_t n, uint32_t lo, uint64_t hi, + uint32_t need, uint32_t cap, uint32_t base, int32_t topBits, int64_t* dbg = nullptr, uint32_t aboveHi = 0u) +{ + if (dbg != nullptr && threadIdx.x == 0) + { + dbg[0] = n; + } + SelectOut out{lo, 0u, 0u, 0u, false, false, 0}; + uint32_t prefMask = 0; + uint32_t prefVal = 0; + uint32_t above = aboveHi; + int32_t remaining = max(topBits, 1); + uint32_t* hist = s.hist; + uint32_t* histNext = s.hist + kBins; + for (int32_t b = threadIdx.x; b < (1 << min(kDigitBits, remaining)); b += kThreads) + { + hist[b] = 0; + } + __syncthreads(); +#pragma unroll 1 + for (int32_t pass = 0; remaining > 0; ++pass) + { + int32_t const width = min(kDigitBits, remaining); + int32_t const shift = remaining - width; + int32_t const nbins = 1 << width; + uint32_t const binMask = static_cast(nbins - 1); + out.passes = pass + 1; + if (dbg != nullptr && threadIdx.x == 0 && pass < 3) + { + dbg[1 + 3 * pass] = clock64(); + } + if constexpr (KeyAt::kVector) + { + for (int32_t j0 = 4 * static_cast(threadIdx.x); j0 < n; j0 += 4 * kThreads) + { + uint4 const q = keyAt.keys4(j0); + histogramKey(hist, true, q.x, lo, hi, base, prefMask, prefVal, shift, binMask); + histogramKey(hist, j0 + 1 < n, q.y, lo, hi, base, prefMask, prefVal, shift, binMask); + histogramKey(hist, j0 + 2 < n, q.z, lo, hi, base, prefMask, prefVal, shift, binMask); + histogramKey(hist, j0 + 3 < n, q.w, lo, hi, base, prefMask, prefVal, shift, binMask); + } + } + else + { + for (int32_t j = threadIdx.x; j < n; j += kThreads) + { + histogramKey(hist, true, keyAt(j), lo, hi, base, prefMask, prefVal, shift, binMask); + } + } + __syncthreads(); + if (dbg != nullptr && threadIdx.x == 0 && pass < 3) + { + dbg[2 + 3 * pass] = clock64(); + } + if (shift > 0) + { + for (int32_t b = threadIdx.x; b < (1 << min(kDigitBits, shift)); b += kThreads) + { + histNext[b] = 0; + } + } + // Find the bin holding the (need - above)-th largest key, counting from the top bin. + uint32_t const target = need - above; + uint32_t const total = findBin(s, hist, nbins, target); + if (dbg != nullptr && threadIdx.x == 0 && pass < 3) + { + dbg[3 + 3 * pass] = clock64(); + } + if (pass == 0) + { + out.total = total + aboveHi; + if (total + aboveHi <= cap || target > total) + { + out.all = true; + out.above = total + aboveHi; + return out; + } + } + // The scalars are next written after two more barriers (the next histogram and scan). + uint32_t const bin = s.scal[kSlotBin]; + above += s.scal[kSlotAbove]; + out.binCount = s.scal[kSlotBinCount]; + prefVal |= bin << shift; + prefMask |= binMask << shift; + out.edge = base + prefVal; + out.above = above; + remaining = shift; + uint32_t* const t = hist; + hist = histNext; + histNext = t; + if (cap > 0 && above + out.binCount <= cap) + { + return out; + } + if (cap == 0 && remaining > 0 && above + out.binCount == need) + { + out.full = true; + return out; + } + } + return out; +} + +// Send one source's selected keys (keys in [lo, hi) of the chunk) to its slot; returns the count sent. +// Takes every predicate key at or above the edge when they fit `cap`. Otherwise (more than `cap` keys tie at the +// exact edge) it takes every key above the edge, plus, with kPartialTies, the first keys equal to it in source +// order up to `cap` (any subset of equal keys is a valid tie-aware choice). Round 1 sends no partial ties, so the +// keys a second round requests from a source (at most the edge) were never sent before. +// The header carries the count and an exclusive upper bound of the keys left unsent (0 when none); in round 1 +// (!kPartialTies) also the key statistics of the keys sent (see kSlotHeaderWords). +// With `stage` (round 1, at most kSourceCapacity keys), the pairs are gathered in shared memory first and written +// with coalesced 16-byte stores; otherwise each pair is stored where it is found. +template +__device__ uint32_t sendSource(SmemLayout const& s, SelectOut const& sel, int32_t n, uint32_t lo, uint64_t hi, + uint32_t cap, uint32_t* mcSlot, int32_t src, int32_t rank, int32_t world, uint2* stage, int64_t* dbg = nullptr) +{ + bool const strict = !sel.all && sel.above + sel.binCount > cap; + uint32_t const threshold = sel.all ? lo : sel.edge; + // Keys equal to the edge that fit after the keys above it (sel.above < need <= cap when strict). + uint32_t const tieBudget = kPartialTies && strict ? cap - sel.above : 0u; + uint32_t const sent = sel.all ? sel.total : (strict ? sel.above + tieBudget : sel.above + sel.binCount); + uint32_t const bound = sel.all ? lo : (strict ? sel.edge + 1u : sel.edge); + stamp(dbg, 42); + // One page block (kBlockTokens keys) per thread; positions from a block scan, so the slot holds the keys in + // source order: first those taken outright, then (strict only) the admitted keys equal to the edge. + int32_t const groups = (n + kBlockTokens - 1) / kBlockTokens; + uint32_t takenMain = 0; + uint32_t takenTies = 0; + uint32_t const tieBase = strict ? sel.above : 0u; + for (int32_t g0 = 0; g0 < groups; g0 += kThreads) + { + int32_t const g = g0 + static_cast(threadIdx.x); + uint32_t k[kBlockTokens]; + uint32_t mask = 0; + uint32_t tieMask = 0; + if (g < groups) + { + uint4 const a = *reinterpret_cast(s.chunk + kBlockTokens * g); + uint4 const b = *reinterpret_cast(s.chunk + kBlockTokens * g + 4); + k[0] = a.x; + k[1] = a.y; + k[2] = a.z; + k[3] = a.w; + k[4] = b.x; + k[5] = b.y; + k[6] = b.z; + k[7] = b.w; +#pragma unroll + for (int32_t i = 0; i < kBlockTokens; ++i) + { + bool const in = kBlockTokens * g + i < n && k[i] >= lo && static_cast(k[i]) < hi; + bool const take = in && (strict ? k[i] > threshold : k[i] >= threshold); + bool const tie = kPartialTies && in && strict && k[i] == threshold; + mask |= (take ? 1u : 0u) << i; + tieMask |= (tie ? 1u : 0u) << i; + } + } + // At most kThreads * kBlockTokens keys per pass: both counts fit 16 bits. + uint32_t const c = static_cast(__popc(mask)) | (static_cast(__popc(tieMask)) << 16); + uint32_t chunkTaken; + uint32_t const before = blockExclusiveScan(c, s.warpSums, chunkTaken); + uint32_t pos = takenMain + (before & 0xFFFFu); + uint32_t tiePos = takenTies + (before >> 16); + takenMain += chunkTaken & 0xFFFFu; + takenTies += chunkTaken >> 16; + if ((mask | tieMask) != 0u) + { + int32_t const colBase = g * kPageTokens + ((src + g) % kSourcesPerRank) * kBlockTokens; +#pragma unroll + for (int32_t i = 0; i < kBlockTokens; ++i) + { + uint32_t slot = kEmptyWord; + if ((mask >> i) & 1u) + { + slot = pos++; + } + else if constexpr (kPartialTies) + { + if ((tieMask >> i) & 1u) + { + slot = tiePos < tieBudget ? tieBase + tiePos : kEmptyWord; + ++tiePos; + } + } + if (slot != kEmptyWord) + { + uint32_t const token = localToGlobal(colBase + i, rank, world); + if (stage != nullptr) + { + stage[slot] = make_uint2(k[i], token); + } + else + { + store2(mcSlot + kSlotHeaderWords + 2 * static_cast(slot), k[i], token); + } + } + } + } + // The staged pairs are complete, and the scan's warp sums may be rewritten. + __syncthreads(); + } + stamp(dbg, 43); + if (stage != nullptr) + { + // Two pairs per 16-byte store; an odd count leaves the last half of the last vector empty. + int32_t const vectors = static_cast(sent + 1u) / 2; + for (int32_t v = threadIdx.x; v < vectors; v += kThreads) + { + uint2 const a = stage[2 * v]; + uint2 const b + = 2 * v + 1 < static_cast(sent) ? stage[2 * v + 1] : make_uint2(kEmptyWord, kEmptyWord); + store4(mcSlot + kSlotHeaderWords + 4 * static_cast(v), make_uint4(a.x, a.y, b.x, b.y)); + } + } + // Receivers check every candidate word, so the header needs no ordering against the candidates. + if constexpr (kPartialTies) + { + if (threadIdx.x == kThreads - 1) + { + store2(mcSlot, sent, bound); + } + } + else + { + // The statistics of the keys sent, from the staged pairs (round 1 always stages), by the last warp, which + // stores no vector (at most kSourceCapacity / 2), while the first warps store the vectors. + static_assert(kSourceCapacity / 2 <= kThreads - 32, "the last warp stores no candidate vector"); + if (threadIdx.x >= kThreads - 32) + { + int32_t const lane = threadIdx.x % 32; + KeyStats t; +#pragma unroll + for (int32_t j = 0; j < kSourceCapacity / 32; ++j) + { + int32_t const i = lane + 32 * j; + if (i < static_cast(sent)) + { + t.add(stage[i].x); + } + } + uint32_t const maxKey = __reduce_max_sync(0xFFFFFFFFu, t.maxKey); + uint32_t const minKey = __reduce_min_sync(0xFFFFFFFFu, t.minKey); + uint32_t const positives = __reduce_add_sync(0xFFFFFFFFu, t.positives); + uint32_t const maxNeg = __reduce_max_sync(0xFFFFFFFFu, t.maxNeg); + uint32_t const minPos = __reduce_min_sync(0xFFFFFFFFu, t.minPos); + if (lane == 31) + { + store4(mcSlot, make_uint4(sent, bound, maxKey, minKey + 1u)); + store4(mcSlot + 4, make_uint4(positives, maxNeg, minPos + 1u, 0u)); + } + } + } + stamp(dbg, 44); + return sent; +} + +// Load one source's keys (one block of kBlockTokens tokens of every local page) from the local logits into the +// chunk; returns the number of keys. Consecutive threads read consecutive keys (whole 32-byte blocks per eight +// threads), and each thread issues kLoadBatch loads before it uses any. With `withStats`, the per-warp key +// statistics are stored for readStats after the next block barrier. +__device__ int32_t loadSource( + SmemLayout const& s, float const* rowLogits, int32_t localLen, int32_t src, bool withStats) +{ + constexpr int32_t kLoadBatch = 8; + int32_t const n = sourceLength(localLen, src); + KeyStats st; + for (int32_t j0 = threadIdx.x; j0 < n; j0 += kLoadBatch * kThreads) + { + float v[kLoadBatch]; +#pragma unroll + for (int32_t i = 0; i < kLoadBatch; ++i) + { + int32_t const j = j0 + i * kThreads; + int32_t const page = j / kBlockTokens; + int32_t const col = page * kPageTokens + ((src + page) % kSourcesPerRank) * kBlockTokens + j % kBlockTokens; + v[i] = j < n ? rowLogits[col] : 0.F; + } +#pragma unroll + for (int32_t i = 0; i < kLoadBatch; ++i) + { + int32_t const j = j0 + i * kThreads; + if (j < n) + { + uint32_t const k = orderedKey(v[i]); + s.chunk[j] = k; + st.add(k); + } + } + } + if (withStats) + { + storeWarpStats(st, s.stats); + } + return n; +} + +// Whether source q withheld keys that may lie above the merged k-th key. +__device__ __forceinline__ bool sourceOverflows(SmemLayout const& s, int32_t q, uint32_t kth) +{ + uint32_t const cnt = s.scal[kSlotSrcCount + q]; + uint64_t const bound = s.scal[kSlotSrcBound + q]; + return cnt < s.scal[kSlotSrcLen + q] && bound > uint64_t{kth} + 1u; +} + +// Emit the top-kTopK of the merge candidates. Keys above `kth` are all taken; ties at `kth` are taken in candidate +// order when only some of them fit (deterministic on every rank), otherwise all of them. With `takeAll`, every +// candidate is taken and the row is padded with -1. The row is assembled in shared memory (the first histogram) +// and written with coalesced stores, together with the bridge outputs of verify row `verifyRow` (pool bases read from +// the block table; the row scalars from `bs`). +template +__device__ bool emitSelection(SmemLayout const& s, KeyAt const& keys, int32_t total, uint32_t kth, bool takeAll, + int32_t* outRow, int32_t overflowSources, CpTopkParams const& p, BridgeShared const& bs, int32_t verifyRow, + int64_t* dbg = nullptr) +{ + uint32_t const want = static_cast(kTopK); + // Positions from one block scan over contiguous candidate ranges (no contended counter); the order is the + // candidate order, identical on every rank. Shared-memory candidates are read four at a time. + int32_t per = (total + kThreads - 1) / kThreads; + if constexpr (KeyAt::kVector) + { + per = (per + 3) / 4 * 4; + } + int32_t const j0 = min(total, static_cast(threadIdx.x) * per); + int32_t const j1 = min(total, j0 + per); + uint32_t cGt = 0; + uint32_t cEq = 0; + if constexpr (KeyAt::kVector) + { + for (int32_t j = j0; j < j1; j += 4) + { + uint4 const q = keys.keys4(j); + uint32_t const kk[4] = {q.x, q.y, q.z, q.w}; +#pragma unroll + for (int32_t i = 0; i < 4; ++i) + { + bool const in = j + i < j1; + cGt += (in && (takeAll || kk[i] > kth)) ? 1u : 0u; + cEq += (in && !takeAll && kk[i] == kth) ? 1u : 0u; + } + } + } + else + { + for (int32_t j = j0; j < j1; ++j) + { + uint32_t const k = keys(j); + cGt += (takeAll || k > kth) ? 1u : 0u; + cEq += (!takeAll && k == kth) ? 1u : 0u; + } + } + // The first `overflowSources` threads flag a source that withheld keys above the threshold (the flag slot is + // zero on entry); with a flag set, nothing is written and the caller runs the second round. + if (static_cast(threadIdx.x) < overflowSources && sourceOverflows(s, threadIdx.x, kth)) + { + atomicOr(&s.scal[kSlotCounter2], 1u); + } + uint32_t allGt; + uint32_t allEq; + uint32_t posGt; + uint32_t posEq; + if (total < 0x10000) + { + // Both counts fit 16 bits: one scan of the packed pair. + uint32_t packedTotal; + uint32_t const packed = cGt | (cEq << 16); + uint32_t const before = blockExclusiveScan(packed, s.warpSums, packedTotal); + allGt = packedTotal & 0xFFFFu; + allEq = packedTotal >> 16; + posGt = before & 0xFFFFu; + posEq = allGt + (before >> 16); + } + else + { + posGt = blockInclusiveScan(cGt, s.warpSums, allGt) - cGt; + __syncthreads(); + posEq = blockInclusiveScan(cEq, s.warpSums, allEq) - cEq; + posEq += allGt; + } + stamp(dbg, 59); + if (overflowSources > 0 && s.scal[kSlotCounter2] != 0u) + { + return true; + } + uint32_t* row = s.hist; + auto place = [&](uint32_t k, uint32_t idx) + { + if (takeAll || k > kth) + { + if (posGt < want) + { + row[posGt] = idx; + } + ++posGt; + } + else if (k == kth) + { + if (posEq < want) + { + row[posEq] = idx; + } + ++posEq; + } + }; + if constexpr (KeyAt::kVector) + { + for (int32_t j = j0; j < j1; j += 4) + { + uint4 const q = keys.keys4(j); + uint4 const x = keys.index4(j); + uint32_t const kk[4] = {q.x, q.y, q.z, q.w}; + uint32_t const xx[4] = {x.x, x.y, x.z, x.w}; +#pragma unroll + for (int32_t i = 0; i < 4; ++i) + { + if (j + i < j1) + { + place(kk[i], xx[i]); + } + } + } + } + else + { + for (int32_t j = j0; j < j1; ++j) + { + place(keys(j), keys.index(j)); + } + } + __syncthreads(); + stamp(dbg, 60); + uint32_t const selected = min(allGt + allEq, want); + BridgeRowOut const bridgeOut{bs.poolRow, p.bridge.outStrideG, bs.copyRow}; + int32_t const req = sharedRequest(p.bridge, bs); + for (uint32_t j = threadIdx.x; j < want; j += kThreads) + { + int32_t const tok = j < selected ? static_cast(row[j]) : -1; + outRow[j] = tok; + writeBridge( + p.bridge, bridgeOut, j, tok, directBase(p.bridge, req, tok), [&](int32_t g) { return bs.layer[g]; }); + } + stamp(dbg, 61); + return false; +} + +// Threshold of an exact (cap = 0) merge select: keys above `kth` are taken, ties at `kth` while they fit. +__device__ __forceinline__ uint32_t mergeThreshold(SelectOut const& m, bool& takeAll) +{ + takeAll = m.all || (m.full && m.edge == 0u); + if (takeAll) + { + return 0u; + } + // A wholly taken bin: the keys at or above its edge are exactly kTopK. + return m.full ? m.edge - 1u : m.edge; +} + +// Reset of the previous exchange launch's buffer on this rank's copy. Its 2 * kMaxSources * kMaxRows slots (round 1, +// then round 2) are spread over the kResetWarps reset warps (the warps from kFirstResetWarp on) of every CTA of the +// launch, at most ten consecutive slots per warp, one per lane: each lane fetches its slot's header once the epoch is +// known (fetchResetHeaders, after the grid dependency wait and the first block barrier; an asynchronous copy into +// shared memory that completes during the local select and the send) and the warp resets the written slots while the +// poll warps wait for the candidates (resetSlots). A slot was written iff its header was; the header's count bounds +// the candidates to reset. +constexpr int32_t kResetSlots = 2 * kMaxSources * kMaxRows; + +// This CTA's index among the launch's clearers, and their number (from the grid; recomputed where used). +__device__ __forceinline__ int32_t resetClearer() +{ + return static_cast(blockIdx.y) * kSourcesPerRank + static_cast(blockIdx.x); +} + +__device__ __forceinline__ int32_t resetClearers() +{ + return static_cast(gridDim.y) * kSourcesPerRank; +} + +__device__ __forceinline__ int32_t resetSlotsPerWarp(int32_t numClearers) +{ + int32_t const warps = numClearers * kResetWarps; + return (kResetSlots + warps - 1) / warps; +} + +// First slot of this (reset) warp. +__device__ __forceinline__ int32_t resetFirstSlot(int32_t clearer, int32_t numClearers) +{ + int32_t const warp = static_cast(threadIdx.x) / 32 - kFirstResetWarp; + return (clearer * kResetWarps + warp) * resetSlotsPerWarp(numClearers); +} + +__device__ __forceinline__ int64_t resetSlotOffset(int32_t buf, int32_t slot) +{ + int32_t const round1Slots = kMaxSources * kMaxRows; + bool const second = slot >= round1Slots; + int32_t const qr = second ? slot - round1Slots : slot; + return second ? round2Offset(buf, qr / kMaxRows, qr % kMaxRows) : round1Offset(buf, qr / kMaxRows, qr % kMaxRows); +} + +__device__ __forceinline__ uint32_t* resetSlotBase(CpTopkParams const& p, int32_t buf, int32_t slot) +{ + return p.ucWorkspace + resetSlotOffset(buf, slot); +} + +//! Most slots one reset warp resets (one row: the fewest reset warps). +constexpr int32_t kMaxResetSlotsPerWarp + = (kResetSlots + kSourcesPerRank * kResetWarps - 1) / (kSourcesPerRank * kResetWarps); +static_assert(kMaxResetSlotsPerWarp <= 32, "one slot header per lane"); + +// Shared-memory home of the reset warps' slot headers (16 bytes per slot: count, bound and two further header words), +// in the merge candidate area beyond the round-1 staging of the send; it is free from the first block barrier until +// the reset. +constexpr int32_t kResetHeaderOffset = 2 * kSourceCapacity; +static_assert(kResetHeaderOffset + 4 * kMaxResetSlotsPerWarp * kResetWarps <= kSourcesPerRank * kSourceCapacity, + "the reset headers fit the smallest merge candidate area"); + +__device__ __forceinline__ uint4* resetHeaderHome(SmemLayout const& s) +{ + int32_t const warp = static_cast(threadIdx.x) / 32 - kFirstResetWarp; + return reinterpret_cast(s.mkey + kResetHeaderOffset) + warp * kMaxResetSlotsPerWarp; +} + +// Word offset, within one Lamport buffer, of the slot whose header this lane fetches for the reset; -1 for the lanes +// and warps without one. Arithmetic on the grid only (no memory access). +__device__ __forceinline__ int32_t resetHeaderOffset() +{ + int32_t const numClearers = resetClearers(); + int32_t const lane = threadIdx.x % 32; + if (static_cast(threadIdx.x) / 32 < kFirstResetWarp || lane >= resetSlotsPerWarp(numClearers)) + { + return -1; + } + int32_t const slot = resetFirstSlot(resetClearer(), numClearers) + lane; + if (slot >= kResetSlots) + { + return -1; + } + static_assert(kBufferWords <= int64_t{1} << 31, "slot offsets within a buffer fit int32"); + return static_cast(resetSlotOffset(0, slot)); +} + +// Starts the copy of this lane's reset slot header of buffer `buf` (offset from resetHeaderOffset) into the warp's +// home with an asynchronous copy (cp.async, generic proxy, completed in resetSlots): a pending copy holds no register +// and no block barrier waits for it. Lanes without a slot do nothing. +__device__ __forceinline__ void fetchResetHeader( + CpTopkParams const& p, SmemLayout const& s, int32_t buf, int32_t offset) +{ + if (offset >= 0) + { + uint32_t const dst = static_cast( + __cvta_generic_to_shared(resetHeaderHome(s) + static_cast(threadIdx.x % 32))); + uint32_t const* const src = p.ucWorkspace + int64_t{buf} * kBufferWords + offset; + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;" ::"r"(dst), "l"(src) : "memory"); + } +} + +// Reset this reset warp's slots of buffer `buf` (headers fetched by fetchResetHeaders); other warps return. +__device__ void resetSlots(CpTopkParams const& p, SmemLayout const& s, int32_t buf) +{ + if (static_cast(threadIdx.x) / 32 < kFirstResetWarp) + { + return; + } + int32_t const clearer = resetClearer(); + int32_t const numClearers = resetClearers(); + // Each lane completes its own header copy; the warp barrier then makes every lane's header visible to the warp. + asm volatile("cp.async.wait_all;" ::: "memory"); + __syncwarp(); + int32_t const per = resetSlotsPerWarp(numClearers); + int32_t const lane = threadIdx.x % 32; + int32_t const first = resetFirstSlot(clearer, numClearers); + uint4 const* const headers = resetHeaderHome(s); + uint4 const empty = make_uint4(kEmptyWord, kEmptyWord, kEmptyWord, kEmptyWord); + for (int32_t i = 0; i < per && first + i < kResetSlots; ++i) + { + uint4 const h = headers[i]; + if (h.x == kEmptyWord && h.y == kEmptyWord) + { + continue; + } + int32_t const slot = first + i; + uint32_t* base = resetSlotBase(p, buf, slot); + uint32_t const capacity = slot >= kMaxSources * kMaxRows ? kRound2Capacity : kSourceCapacity; + uint32_t const count = h.x == kEmptyWord ? capacity : min(h.x, capacity); + int32_t const vectors = static_cast(count + 1u) / 2; + for (int32_t v = lane; v < vectors; v += 32) + { + store4(base + kSlotHeaderWords + 4 * static_cast(v), empty); + } + if (lane < 2) + { + store4(base + 4 * lane, empty); + } + } +} + +// Merge of one row on this rank (one CTA): gathers every source's round-1 candidates, selects the exact +// top-kTopK, runs the second round when a source withheld keys above the threshold, and writes the row and its bridge +// outputs. +__device__ void mergeRow(CpTopkParams const& p, SmemLayout const& s, uint32_t const** segSlots, uint32_t* segBegin, + int32_t buf, int32_t row, int32_t rowLen, int32_t localLen, float const* rowLogits, int32_t* outRow, + WaitCtx const& wait, BridgeShared const& bs) +{ + int32_t const world = p.worldSize; + int32_t const rank = p.rank; + int32_t const sources = numSources(world); + uint32_t const cap = static_cast(roundCap(world)); + __syncthreads(); + if (threadIdx.x == 0) + { + s.scal[kSlotCounter2] = 0u; + } + if (static_cast(threadIdx.x) < sources) + { + // Keys each source holds, computed while its header is in flight. + int32_t const q = threadIdx.x; + s.scal[kSlotSrcLen + q] = static_cast( + sourceLength(min(localLength(rowLen, q / kSourcesPerRank, world), p.localCapacity), q % kSourcesPerRank)); + uint2 const h = pollPair(p.ucWorkspace + round1Offset(buf, threadIdx.x, row), wait, kWaitRound1Header); + s.scal[kSlotSrcCount + threadIdx.x] = min(h.x, cap); + s.scal[kSlotSrcBound + threadIdx.x] = h.y; + } + __syncthreads(); + if (threadIdx.x < 32) + { + // Exclusive prefix of the source counts (at most 64 sources: two per lane). + uint32_t const a = static_cast(threadIdx.x) < sources ? s.scal[kSlotSrcCount + threadIdx.x] : 0u; + uint32_t const b + = static_cast(threadIdx.x) + 32 < sources ? s.scal[kSlotSrcCount + threadIdx.x + 32] : 0u; + uint32_t x = a; + uint32_t z = b; +#pragma unroll + for (int32_t o = 1; o < 32; o <<= 1) + { + uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, x, o); + uint32_t const w = __shfl_up_sync(0xFFFFFFFFu, z, o); + if (static_cast(threadIdx.x) >= o) + { + x += y; + z += w; + } + } + uint32_t const totalA = __shfl_sync(0xFFFFFFFFu, x, 31); + s.scal[kSlotSrcOffset + threadIdx.x] = x - a; + s.scal[kSlotSrcOffset + 32 + threadIdx.x] = totalA + z - b; + if (threadIdx.x == 31) + { + s.scal[kSlotSrcOffset + kMaxSources] = totalA + z; + } + } + __syncthreads(); + int32_t const total = static_cast(s.scal[kSlotSrcOffset + kMaxSources]); + markPhase(p, 4); + + // Candidates: padded (source, vector) positions; issue a batch of independent 16-byte loads, then re-poll the + // ones not yet written. The key statistics are tracked for the merge select. + int32_t const vectorsPerSource = static_cast(cap) / 2; + int32_t const padded = sources * vectorsPerSource; + KeyStats st; +#pragma unroll 1 + for (int32_t v0 = 0; v0 < padded; v0 += kThreads * kPollBatch) + { + uint4 v[kPollBatch]; +#pragma unroll + for (int32_t i = 0; i < kPollBatch; ++i) + { + int32_t const jv = v0 + i * kThreads + static_cast(threadIdx.x); + int32_t const q = jv / vectorsPerSource; + int32_t const kv = jv % vectorsPerSource; + v[i] = make_uint4(0u, 0u, 0u, 0u); + if (jv < padded && 2 * kv < static_cast(s.scal[kSlotSrcCount + q])) + { + v[i] = loadVolatile4(p.ucWorkspace + round1Offset(buf, q, row) + kSlotHeaderWords + 4 * kv); + } + } +#pragma unroll + for (int32_t i = 0; i < kPollBatch; ++i) + { + int32_t const jv = v0 + i * kThreads + static_cast(threadIdx.x); + int32_t const q = jv / vectorsPerSource; + int32_t const kv = jv % vectorsPerSource; + if (jv < padded && 2 * kv < static_cast(s.scal[kSlotSrcCount + q])) + { + bool const second = 2 * kv + 1 < static_cast(s.scal[kSlotSrcCount + q]); + if (!vectorReady(v[i], second)) + { + v[i] = pollVector(p.ucWorkspace + round1Offset(buf, q, row) + kSlotHeaderWords + 4 * kv, second, + wait, kWaitRound1Candidates); + } + int32_t const j = static_cast(s.scal[kSlotSrcOffset + q]) + 2 * kv; + s.mkey[j] = v[i].x; + s.midx[j] = v[i].y; + st.add(v[i].x); + if (second) + { + s.mkey[j + 1] = v[i].z; + s.midx[j + 1] = v[i].w; + st.add(v[i].z); + } + } + } + } + stamp(debugSlots(p, 0), 45); + storeWarpStats(st, s.stats); + __syncthreads(); + markPhase(p, 5); + + // Digits of the merge select span the candidates of the sign the cut lies in (see kPositiveKey). + SignRange const r = total > 0 ? signRange(readStats(s.stats), static_cast(kTopK)) + : SignRange{0u, uint64_t{kEmptyWord}, 32, 0u, false}; + SmemKeys const merged{s.mkey, s.midx}; + SelectOut const m = blockSelect( + s, merged, total, r.lo, r.hi, static_cast(kTopK), 0u, r.lo, r.bits, debugSlots(p, 48), r.aboveHi); + putDebug(debugSlots(p, 0), 62, total); + putDebug(debugSlots(p, 0), 63, m.passes + (r.negative ? 10 : 0)); + bool takeAll = false; + uint32_t const kth = mergeThreshold(m, takeAll); + markPhase(p, 6); + stamp(debugSlots(p, 0), 58); + if (!emitSelection(s, merged, total, kth, takeAll, outRow, sources, p, bs, row, debugSlots(p, 0))) + { + return; + } + + // ---- Round 2 (rare; every rank takes it): withheld keys in [kth, bound) of overflowing sources ---------- + __syncthreads(); + // 1. This rank's overflowing sources send their top-kRound2Capacity withheld keys. + for (int32_t c = 0; c < kSourcesPerRank; ++c) + { + int32_t const q = rank * kSourcesPerRank + c; + if (!sourceOverflows(s, q, kth)) + { + continue; + } + uint64_t const b = s.scal[kSlotSrcBound + q]; + int32_t const nq = loadSource(s, rowLogits, localLen, c, false); + __syncthreads(); + SelectOut const sel2 = blockSelect( + s, SmemKeys{s.chunk, nullptr}, nq, kth, b, static_cast(kRound2Capacity), kRound2Capacity, 0u, 32); + (void) sendSource( + s, sel2, nq, kth, b, kRound2Capacity, p.mcWorkspace + round2Offset(buf, q, row), c, rank, world, nullptr); + __syncthreads(); + } + // 2. Gather the round-2 segments of every overflowing source (all ranks). + if (threadIdx.x == 0) + { + int32_t segs = 0; + uint32_t off = 0; + for (int32_t q = 0; q < sources; ++q) + { + if (!sourceOverflows(s, q, kth)) + { + continue; + } + uint32_t const* slot = p.ucWorkspace + round2Offset(buf, q, row); + uint2 const h = pollPair(slot, wait, kWaitRound2Header); + segSlots[segs] = slot; + segBegin[segs] = off; + off += min(h.x, static_cast(kRound2Capacity)); + ++segs; + } + segBegin[segs] = off; + s.scal[kSlotTotal] = static_cast(segs); + } + __syncthreads(); + int32_t const segs = static_cast(s.scal[kSlotTotal]); + int32_t const total2 = static_cast(segBegin[segs]); + MergeKeys const keys{s.mkey, s.midx, total, segSlots, segBegin, segs}; + for (int32_t k = threadIdx.x; k < total2; k += kThreads) + { + (void) pollPair(keys.pair(k), wait, kWaitRound2Candidates); + } + __syncthreads(); + SelectOut const m2 + = blockSelect(s, keys, total + total2, 0u, uint64_t{kEmptyWord}, static_cast(kTopK), 0u, 0u, 32); + bool takeAll2 = false; + uint32_t const kth2 = mergeThreshold(m2, takeAll2); + (void) emitSelection(s, keys, total + total2, kth2, takeAll2, outRow, 0, p, bs, row); + if (threadIdx.x == 0) + { + atomicAdd(&p.state[2], 1); + } +} + +// ---- Thread-block cluster primitives (one cluster of kClusterCtas CTAs per row; cluster rank = blockIdx.x) ---------- + +__device__ __forceinline__ uint32_t smemAddr(void const* ptr) +{ + return static_cast(__cvta_generic_to_shared(ptr)); +} + +// Shared-memory window address of `ptr` (this CTA's shared memory) in CTA `rank` of the cluster. +__device__ __forceinline__ uint32_t clusterAddr(void const* ptr, int32_t rank) +{ + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(remote) : "r"(smemAddr(ptr)), "r"(rank)); + return remote; +} + +// Initialize the exchange barriers and announce it to the cluster (relaxed arrive; clusterInitWait completes it +// before the first exchange). Touches shared memory only. +__device__ __forceinline__ void clusterInit(ClusterRecords& rec) +{ + if (threadIdx.x == 0) + { + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(smemAddr(&rec.bar[0])) : "memory"); + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(smemAddr(&rec.bar[1])) : "memory"); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + asm volatile("barrier.cluster.arrive.relaxed.aligned;" ::: "memory"); +} + +__device__ __forceinline__ void clusterInitWait() +{ + asm volatile("barrier.cluster.wait.aligned;" ::: "memory"); +} + +// This CTA's single arrival on exchange barrier `bar`, expecting `bytes` from the cluster's st.async stores. +__device__ __forceinline__ void expectBytes(uint64_t* bar, uint32_t bytes) +{ + asm volatile( + "{\n\t.reg .b64 st;\n\tmbarrier.arrive.expect_tx.shared::cta.b64 st, [%0], %1;\n\t}" ::"r"(smemAddr(bar)), + "r"(bytes) + : "memory"); +} + +// Wait until the phase of exchange barrier `bar` with parity `parity` completed (every expected byte arrived). +__device__ __forceinline__ void waitExchange(uint64_t* bar, uint32_t parity) +{ + uint32_t done; + do + { + asm volatile( + "{\n\t.reg .pred p;\n\tmbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 p, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, p;\n\t}" + : "=r"(done) + : "r"(smemAddr(bar)), "r"(parity) + : "memory"); + } while (done == 0u); +} + +// Store 16 / 8 bytes at `ptr` (a record of this CTA's layout) in CTA `rank`, completing them on its barrier `bar`. +__device__ __forceinline__ void sendTo(int32_t rank, void const* ptr, uint64_t const* bar, uint4 v) +{ + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [%0], {%1, %2, %3, %4}, [%5];" ::"r"( + clusterAddr(ptr, rank)), + "r"(v.x), "r"(v.y), "r"(v.z), "r"(v.w), "r"(clusterAddr(bar, rank)) + : "memory"); +} + +__device__ __forceinline__ void sendTo(int32_t rank, void const* ptr, uint64_t const* bar, uint2 v) +{ + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.v2.b32 [%0], {%1, %2}, [%3];" ::"r"( + clusterAddr(ptr, rank)), + "r"(v.x), "r"(v.y), "r"(clusterAddr(bar, rank)) + : "memory"); +} + +// Exchange sequence of a row: per pass x = 2 * pass group totals and x = 2 * pass + 1 the selected group's bins. +// Every CTA sends the records of exchange x + 1 only after its own exchange x completed, so a barrier phase (and an +// alternating record set) is reused only after every CTA consumed the previous one. +__device__ __forceinline__ uint64_t* exchangeBar(ClusterRecords& rec, uint32_t x) +{ + return &rec.bar[x & 1u]; +} + +__device__ __forceinline__ uint32_t exchangeParity(uint32_t x) +{ + return (x >> 1) & 1u; +} + +// Inclusive scan over the lanes of a warp. +__device__ __forceinline__ uint32_t warpInclusiveScan(uint32_t v) +{ + int32_t const lane = threadIdx.x % 32; +#pragma unroll + for (int32_t o = 1; o < 32; o <<= 1) + { + uint32_t const y = __shfl_up_sync(0xFFFFFFFFu, v, o); + if (lane >= o) + { + v += y; + } + } + return v; +} + +// Tallies of the cluster select over the candidates of the cluster CTAs before this one (in candidate order), kept +// by warp 0 and published in shared memory: those strictly above the selected bins of all passes so far, and those +// in the selected bin of the last pass. +struct CtaTally +{ + uint32_t aboveBefore; + uint32_t binBefore; +}; + +// Cluster radix select of the need-th largest key among the candidates of every CTA of the cluster (this thread +// holds k0 and k1; kEmptyWord is no candidate, as it lies above every range). It computes exactly what blockSelect +// with cap = 0 computes over the union of the candidates: the same digits (from `base` and topBits), the same bins +// and counts from the cluster-summed histograms, and the same stop rules. Per pass: every CTA histograms its own +// candidates and sends its totals per group of kGroupBins bins to every CTA (st.async, two groups per warp); warp 0 +// of every CTA finds, from the top, the group holding the target, sends its own counts of that group's bins to every +// CTA, and finds the bin from the eight CTAs' counts. From the same records, warp 0 also tallies the candidates of +// the CTAs before this one above the selected bins and in the last selected bin (`tally`), which give this CTA's +// emission offsets without another exchange. +__device__ SelectOut clusterSelect(SmemLayout const& s, ClusterRecords& rec, CtaTally& tally, uint32_t k0, uint32_t k1, + uint32_t lo, uint64_t hi, uint32_t need, uint32_t base, int32_t topBits, uint32_t aboveHi, int64_t* dbg) +{ + int32_t const me = blockIdx.x; + int32_t const lane = threadIdx.x % 32; + int32_t const warp = threadIdx.x / 32; + SelectOut out{lo, 0u, 0u, 0u, false, false, 0}; + uint32_t prefMask = 0; + uint32_t prefVal = 0; + uint32_t above = aboveHi; + int32_t remaining = max(topBits, 1); + uint32_t* hist = s.hist; + uint32_t* histNext = s.hist + kBins; + // Warp 0: the running tallies of the CTAs before this one. + uint32_t aboveBefore = 0u; + uint32_t binBefore = 0u; + for (int32_t b = threadIdx.x; b < (1 << min(kDigitBits, remaining)); b += kThreads) + { + hist[b] = 0; + } + __syncthreads(); +#pragma unroll 1 + for (int32_t pass = 0; remaining > 0; ++pass) + { + int32_t const width = min(kDigitBits, remaining); + int32_t const shift = remaining - width; + int32_t const nbins = 1 << width; + uint32_t const binMask = static_cast(nbins - 1); + int32_t const groups = (nbins + kGroupBins - 1) / kGroupBins; + int32_t const pairs = (groups + 1) / 2; + uint32_t const xg = 2u * pass; + uint32_t const xb = xg + 1u; + out.passes = pass + 1; + if (dbg != nullptr && threadIdx.x == 0 && pass < 3) + { + dbg[1 + 3 * pass] = clock64(); + } + if (threadIdx.x == 0) + { + expectBytes(exchangeBar(rec, xg), kClusterCtas * pairs * 8); + } + // The next pass's histogram: its last reads (the previous pass's bins sent by warp 0) are complete. + if (shift > 0) + { + for (int32_t b = threadIdx.x; b < (1 << min(kDigitBits, shift)); b += kThreads) + { + histNext[b] = 0; + } + } + histogramKey(hist, true, k0, lo, hi, base, prefMask, prefVal, shift, binMask); + histogramKey(hist, true, k1, lo, hi, base, prefMask, prefVal, shift, binMask); + __syncthreads(); + uint32_t(*const groupRec)[kMaxGroups] = rec.groups[pass & 1]; + uint32_t(*const binRec)[kGroupBins] = rec.bins[pass & 1]; + if (warp < pairs) + { + int32_t const b0 = 2 * warp * kGroupBins + lane; + int32_t const b1 = b0 + kGroupBins; + uint32_t const t0 = __reduce_add_sync(0xFFFFFFFFu, b0 < nbins ? hist[b0] : 0u); + uint32_t const t1 = __reduce_add_sync(0xFFFFFFFFu, b1 < nbins ? hist[b1] : 0u); + if (lane < kClusterCtas) + { + sendTo(lane, &groupRec[me][2 * warp], exchangeBar(rec, xg), make_uint2(t0, t1)); + } + } + uint32_t const target = need - above; + if (warp == 0) + { + waitExchange(exchangeBar(rec, xg), exchangeParity(xg)); + if (dbg != nullptr && lane == 0 && pass < 3) + { + dbg[2 + 3 * pass] = clock64(); + } + // Lane l sums the groups at positions 2l and 2l + 1 from the top over the cluster. + int32_t const ga = groups - 1 - 2 * lane; + int32_t const gb = ga - 1; + uint32_t a = 0; + uint32_t b = 0; + uint32_t beforeA = 0; + uint32_t beforeB = 0; +#pragma unroll + for (int32_t c = 0; c < kClusterCtas; ++c) + { + uint32_t const x = ga >= 0 ? groupRec[c][ga] : 0u; + uint32_t const y = gb >= 0 ? groupRec[c][gb] : 0u; + a += x; + b += y; + beforeA += c < me ? x : 0u; + beforeB += c < me ? y : 0u; + } + uint32_t const incl = warpInclusiveScan(a + b); + uint32_t const before = incl - a - b; + uint32_t const total = __shfl_sync(0xFFFFFFFFu, incl, 31); + bool const holds = before < target && target <= incl; + uint32_t const holder = __ballot_sync(0xFFFFFFFFu, holds); + if (lane == 0) + { + s.scal[kSlotTotal] = total; + } + if (holder != 0u) + { + int32_t const hl = __ffs(holder) - 1; + bool const first = target <= before + a; + int32_t const gSel = __shfl_sync(0xFFFFFFFFu, first ? ga : gb, hl); + uint32_t const gAbove = __shfl_sync(0xFFFFFFFFu, first ? before : before + a, hl); + if (lane == 0) + { + expectBytes(exchangeBar(rec, xb), kClusterCtas * kGroupBins * 4); + } + // This CTA's counts of the selected group's bins to every CTA: lane l sends vector l % 8 to CTAs + // l / 8 and l / 8 + 4. + int32_t const v = lane % 8; + int32_t const b0 = gSel * kGroupBins + 4 * v; + uint4 const q = make_uint4(b0 < nbins ? hist[b0] : 0u, b0 + 1 < nbins ? hist[b0 + 1] : 0u, + b0 + 2 < nbins ? hist[b0 + 2] : 0u, b0 + 3 < nbins ? hist[b0 + 3] : 0u); + sendTo(lane / 8, &binRec[me][4 * v], exchangeBar(rec, xb), q); + sendTo(lane / 8 + 4, &binRec[me][4 * v], exchangeBar(rec, xb), q); + // While the bins travel: the earlier CTAs' candidates in the groups above the selected one (both + // groups of the lanes before the holder, and the holder's first group when it selected the second). + bool const selFirst = __shfl_sync(0xFFFFFFFFu, first, hl); + aboveBefore += __reduce_add_sync( + 0xFFFFFFFFu, lane < hl ? beforeA + beforeB : (lane == hl && !selFirst ? beforeA : 0u)); + waitExchange(exchangeBar(rec, xb), exchangeParity(xb)); + // Bins of the selected group from the top: lane l holds bin 32 * gSel + 31 - l. + int32_t const bin = gSel * kGroupBins + kGroupBins - 1 - lane; + uint32_t c = 0; + uint32_t cBefore = 0; +#pragma unroll + for (int32_t d = 0; d < kClusterCtas; ++d) + { + uint32_t const x = binRec[d][kGroupBins - 1 - lane]; + c += x; + cBefore += d < me ? x : 0u; + } + uint32_t const y = warpInclusiveScan(c); + uint32_t const acc = gAbove + y - c; + bool const found = acc < target && target <= acc + c; + if (found) + { + s.scal[kSlotBin] = static_cast(bin); + s.scal[kSlotAbove] = acc; + s.scal[kSlotBinCount] = c; + } + // The earlier CTAs' candidates in the bins above the selected one (lanes before the finder) and in + // the selected bin (the finder's). + int32_t const finder = __ffs(__ballot_sync(0xFFFFFFFFu, found)) - 1; + aboveBefore += __reduce_add_sync(0xFFFFFFFFu, lane < finder ? cBefore : 0u); + binBefore = __shfl_sync(0xFFFFFFFFu, cBefore, finder); + } + else if (lane == 0) + { + // Every key is taken (pass 0 only, the same decision in every CTA): the bin exchange still + // completes, with no bytes, so that the exchange sequence stays aligned. + expectBytes(exchangeBar(rec, xb), 0u); + waitExchange(exchangeBar(rec, xb), exchangeParity(xb)); + } + if (lane == 0) + { + tally.aboveBefore = aboveBefore; + tally.binBefore = binBefore; + } + } + __syncthreads(); + if (dbg != nullptr && threadIdx.x == 0 && pass < 3) + { + dbg[3 + 3 * pass] = clock64(); + } + uint32_t const total = s.scal[kSlotTotal]; + if (pass == 0) + { + out.total = total + aboveHi; + if (total + aboveHi == 0u || target > total) + { + out.all = true; + out.above = total + aboveHi; + return out; + } + } + // The scalars are next written by warp 0 after the next histogram's block barrier. + uint32_t const bin = s.scal[kSlotBin]; + above += s.scal[kSlotAbove]; + out.binCount = s.scal[kSlotBinCount]; + prefVal |= bin << shift; + prefMask |= binMask << shift; + out.edge = base + prefVal; + out.above = above; + remaining = shift; + uint32_t* const t = hist; + hist = histNext; + histNext = t; + if (remaining > 0 && above + out.binCount == need) + { + out.full = true; + return out; + } + } + return out; +} + +// Spin until the header of a round-1 slot was written and, when the slot's count covers candidate vector kv, that +// vector too; returns the header's count (clamped to `cap`) and bound. Both are polled together, so neither waits +// for the other; aborts the kernel after the timeout. +__device__ uint2 pollSlotVector(uint32_t const* slot, int32_t kv, uint32_t cap, uint4& v, WaitCtx const& w) +{ + uint32_t const* vec = slot + kSlotHeaderWords + 4 * kv; + uint2 h = loadVolatile2(slot); + v = loadVolatile4(vec); + uint64_t start = 0; + for (int32_t polls = 1;; ++polls) + { + bool const headerReady = !isEmpty(h); + uint32_t const cnt = headerReady ? min(h.x, cap) : 0u; + if (headerReady && (2u * kv >= cnt || vectorReady(v, 2u * kv + 1u < cnt))) + { + return make_uint2(cnt, h.y); + } + if (polls % kPollsPerTimerRead == 0) + { + uint64_t const now = globalTimerNs(); + if (start == 0) + { + start = now; + } + else if (now - start > w.timeoutNs) + { + abortOnTimeout(w, headerReady ? kWaitRound1Candidates : kWaitRound1Header); + } + } + if (!headerReady) + { + h = loadVolatile2(slot); + } + v = loadVolatile4(vec); + } +} + +// Spin until the whole 32-byte round-1 header of a slot was written; aborts the kernel after the timeout. +__device__ void pollHeader(uint32_t const* slot, uint4& a, uint4& b, WaitCtx const& w) +{ + a = loadVolatile4(slot); + b = loadVolatile4(slot + 4); + uint64_t start = 0; + for (int32_t polls = 1; !(vectorReady(a, true) && vectorReady(b, true)); ++polls) + { + if (polls % kPollsPerTimerRead == 0) + { + uint64_t const now = globalTimerNs(); + if (start == 0) + { + start = now; + } + else if (now - start > w.timeoutNs) + { + abortOnTimeout(w, kWaitRound1Header); + } + } + a = loadVolatile4(slot); + b = loadVolatile4(slot + 4); + } +} + +// Merge of one row by its cluster. CTA c gathers the round-1 candidates of sources [c * world, (c + 1) * world) (one +// 16-byte vector, two candidate slots, per thread of the poll warps, kept in registers: thread order is crossRankTopk's +// candidate order) while its header warps read every source's header (counts, bounds and the key statistics of the keys +// sent, so every CTA has the row's statistics and can decide every source's overflow without an exchange). The cluster +// runs crossRankTopk's merge select on the summed histograms (clusterSelect), whose records also give every CTA's +// number of keys above the threshold and of ties, and every CTA writes its selected candidates at crossRankTopk's +// positions: keys above the threshold in candidate order, then ties in candidate order while they fit, then -1, each +// position with its bridge outputs (pool bases fetched per candidate during the cluster select). Returns true, writing +// nothing, when a source withheld keys that may lie above the threshold (the same decision on every CTA and rank); the +// row is then merged by mergeRow. +__device__ bool clusterMergeRow(CpTopkParams const& p, SmemLayout const& s, ClusterRecords& rec, CtaTally& tally, + int32_t buf, int32_t row, int32_t rowLen, int32_t* outRow, WaitCtx const& wait, BridgeShared const& bs) +{ + int32_t const world = p.worldSize; + int32_t const sources = numSources(world); + int32_t const me = blockIdx.x; + int32_t const tid = threadIdx.x; + int32_t const lane = tid % 32; + int32_t const warp = tid / 32; + uint32_t const cap = static_cast(roundCap(world)); + int32_t const vectorsPerSource = static_cast(cap) / 2; + int32_t const q0 = me * world; + int64_t* const dbg = debugSlots(p, 0); + // This thread's candidates (kEmptyWord: none; never a key). + uint32_t k0 = kEmptyWord; + uint32_t k1 = kEmptyWord; + if (tid == 0) + { + s.scal[kSlotCounter2] = 0u; + } + if (tid < world * vectorsPerSource) + { + int32_t const q = q0 + tid / vectorsPerSource; + int32_t const kv = tid % vectorsPerSource; + uint4 v; + uint2 const h = pollSlotVector(p.ucWorkspace + round1Offset(buf, q, row), kv, cap, v, wait); + if (2u * kv < h.x) + { + k0 = v.x; + if (2u * kv + 1u < h.x) + { + k1 = v.z; + } + } + // The tokens wait in shared memory until the emission (s.midx is free in the cluster merge). Their pool bases + // are fetched now, while the cluster selects: copies of the candidates' block-table entries into this + // thread's words beyond kBaseOffsetWords, completed before the block barrier of the emission prefix. + *reinterpret_cast(s.midx + 2 * tid) = make_uint2(v.y, v.w); + if (k0 != kEmptyWord) + { + int32_t const req = sharedRequest(p.bridge, bs); + uint32_t* const bases = s.midx + kBaseOffsetWords + 2 * tid; + fetchBase(p.bridge, req, static_cast(v.y), bases); + if (k1 != kEmptyWord) + { + fetchBase(p.bridge, req, static_cast(v.w), bases + 1); + } + } + } + else if (warp >= kPollWarps && warp < kFirstResetWarp) + { + // Header warps: thread h reads source h's header; per source its count (clamped to the round-1 capacity), + // bound and length, and per header warp the reduced key statistics. + int32_t const q = tid - kPollWarps * 32; + KeyStats st; + if (q < sources) + { + uint4 a; + uint4 b; + pollHeader(p.ucWorkspace + round1Offset(buf, q, row), a, b, wait); + s.scal[kSlotSrcCount + q] = min(a.x, cap); + s.scal[kSlotSrcBound + q] = a.y; + s.scal[kSlotSrcLen + q] = static_cast(sourceLength( + min(localLength(rowLen, q / kSourcesPerRank, world), p.localCapacity), q % kSourcesPerRank)); + st.maxKey = a.z; + st.minKey = a.w - 1u; + st.positives = b.x; + st.maxNeg = b.y; + st.minPos = b.z - 1u; + } + uint32_t const cnt = q < sources ? s.scal[kSlotSrcCount + q] : 0u; + // Candidates and keys >= +0.0 of the cluster CTAs before this one (CTA c merges sources + // [c * world, (c + 1) * world)), of this one and of the row. + int32_t const c = q < sources ? q / world : kClusterCtas; + uint32_t const nBefore = __reduce_add_sync(0xFFFFFFFFu, c < me ? cnt : 0u); + uint32_t const nMine = __reduce_add_sync(0xFFFFFFFFu, c == me ? cnt : 0u); + uint32_t const nAll = __reduce_add_sync(0xFFFFFFFFu, cnt); + uint32_t const pBefore = __reduce_add_sync(0xFFFFFFFFu, c < me ? st.positives : 0u); + if (lane == 0) + { + atomicAdd(&s.scal[kSlotCountBefore], nBefore); + atomicAdd(&s.scal[kSlotCountMine], nMine); + atomicAdd(&s.scal[kSlotCountTotal], nAll); + atomicAdd(&s.scal[kSlotPositivesBefore], pBefore); + } + uint32_t const v0 = __reduce_max_sync(0xFFFFFFFFu, st.maxKey); + uint32_t const v1 = __reduce_min_sync(0xFFFFFFFFu, st.minKey); + uint32_t const v2 = __reduce_add_sync(0xFFFFFFFFu, st.positives); + uint32_t const v3 = __reduce_max_sync(0xFFFFFFFFu, st.maxNeg); + uint32_t const v4 = __reduce_min_sync(0xFFFFFFFFu, st.minPos); + if (lane == 0) + { + // In the scan scratch (s.stats may still be read by the send's last warp; s.warpSums is next written + // by the emission scan, after the reads below). + uint32_t* const out = s.warpSums + kNumKeyStats * (warp - kPollWarps); + out[0] = v0; + out[1] = v1; + out[2] = v2; + out[3] = v3; + out[4] = v4; + } + } + stamp(dbg, 45); + __syncthreads(); + markPhase(p, 4); + + // Statistics of the row's candidates from every source's header (identical in every CTA and rank). + KeyStats g; +#pragma unroll + for (int32_t w = 0; w < kHeaderWarps; ++w) + { + uint32_t const* const in = s.warpSums + kNumKeyStats * w; + g.maxKey = max(g.maxKey, in[0]); + g.minKey = min(g.minKey, in[1]); + g.positives += in[2]; + g.maxNeg = max(g.maxNeg, in[3]); + g.minPos = min(g.minPos, in[4]); + } + uint32_t const total = s.scal[kSlotCountTotal]; + stamp(dbg, 61); + markPhase(p, 5); + + SignRange const r + = total > 0u ? signRange(g, static_cast(kTopK)) : SignRange{0u, uint64_t{kEmptyWord}, 32, 0u, false}; + clusterInitWait(); + SelectOut const m = clusterSelect( + s, rec, tally, k0, k1, r.lo, r.hi, static_cast(kTopK), r.lo, r.bits, r.aboveHi, debugSlots(p, 48)); + putDebug(debugSlots(p, 48), 0, s.scal[kSlotCountMine]); + putDebug(dbg, 62, total); + putDebug(dbg, 63, m.passes + (r.negative ? 10 : 0)); + bool takeAll = false; + uint32_t const kth = mergeThreshold(m, takeAll); + markPhase(p, 6); + + // Keys above the threshold and ties of the cluster CTAs before this one and of the row (identical in every CTA), + // from the select's tallies; with the threshold below a wholly taken bin the ties fall beyond kTopK and are taken + // as keys above it. + uint32_t gtBefore; + uint32_t eqBefore; + uint32_t allGt; + uint32_t allEq; + if (takeAll) + { + gtBefore = s.scal[kSlotCountBefore]; + eqBefore = 0u; + allGt = total; + allEq = 0u; + } + else + { + uint32_t const aboveBefore = tally.aboveBefore + (r.aboveHi > 0u ? s.scal[kSlotPositivesBefore] : 0u); + gtBefore = m.full ? aboveBefore + tally.binBefore : aboveBefore; + eqBefore = m.full ? 0u : tally.binBefore; + allGt = m.full ? m.above + m.binCount : m.above; + allEq = m.full ? 0u : m.binCount; + } + // This CTA's positions, scanned in thread order; a source of the row that withheld keys above the threshold + // sends the row to the second round. + bool const gt0 = k0 != kEmptyWord && (takeAll || k0 > kth); + bool const eq0 = k0 != kEmptyWord && !takeAll && k0 == kth; + bool const gt1 = k1 != kEmptyWord && (takeAll || k1 > kth); + bool const eq1 = k1 != kEmptyWord && !takeAll && k1 == kth; + if (tid < sources && sourceOverflows(s, tid, kth)) + { + atomicOr(&s.scal[kSlotCounter2], 1u); + } + // Exclusive prefix in thread order (k0 before k1): within the warp from ballots, across warps from the warp + // totals (one block barrier). + uint32_t const below = (1u << lane) - 1u; + uint32_t const bGt0 = __ballot_sync(0xFFFFFFFFu, gt0); + uint32_t const bGt1 = __ballot_sync(0xFFFFFFFFu, gt1); + uint32_t const bEq0 = __ballot_sync(0xFFFFFFFFu, eq0); + uint32_t const bEq1 = __ballot_sync(0xFFFFFFFFu, eq1); + if (lane == 0) + { + s.warpSums[warp] = static_cast(__popc(bGt0) + __popc(bGt1)) + + (static_cast(__popc(bEq0) + __popc(bEq1)) << 16); + } + // The pool-base copies are complete before the barrier (a second-round row reuses s.midx). + asm volatile("cp.async.wait_all;" ::: "memory"); + __syncthreads(); + uint32_t const warpsBefore = __reduce_add_sync(0xFFFFFFFFu, lane < warp ? s.warpSums[lane] : 0u); + uint32_t const before = warpsBefore + static_cast(__popc(bGt0 & below) + __popc(bGt1 & below)) + + (static_cast(__popc(bEq0 & below) + __popc(bEq1 & below)) << 16); + stamp(dbg, 58); + if (s.scal[kSlotCounter2] != 0u) + { + return true; + } + uint32_t const want = static_cast(kTopK); + uint32_t posGt = gtBefore + (before & 0xFFFFu); + uint32_t posEq = allGt + eqBefore + (before >> 16); + uint2 tokens = make_uint2(0u, 0u); + uint2 tokenBases = make_uint2(0u, 0u); + if (k0 != kEmptyWord) + { + tokens = *reinterpret_cast(s.midx + 2 * tid); + tokenBases = *reinterpret_cast(s.midx + kBaseOffsetWords + 2 * tid); + } + // Output positions of this thread's two candidates (want: not written) and of its share of the -1 padding (at + // most one position per thread of the cluster). + static_assert(kClusterCtas * kThreads >= kTopK, "one padding position per thread"); + uint32_t const pos0 = gt0 ? posGt : (eq0 ? posEq : want); + posGt += gt0 ? 1u : 0u; + posEq += eq0 ? 1u : 0u; + uint32_t const pos1 = gt1 ? posGt : (eq1 ? posEq : want); + uint32_t const selected = min(allGt + allEq, want); + uint32_t const posPad = selected + static_cast(me * kThreads + tid); + emitCandidatesFor(p.bridge.groupSize, outRow, bs.poolRow, bs.copyRow, p.bridge.outStrideG, + static_cast(p.bridge.strideFactor), bs.layer, pos0, pos1, posPad, tokens, tokenBases); + stamp(dbg, 60); + return false; +} + +__global__ void __launch_bounds__(kThreads, 1) cpTopkZ11Kernel(CpTopkParams p) +{ + markPhase(p, 0); + extern __shared__ __align__(16) uint32_t smemRaw[]; + __shared__ uint32_t const* segSlots[kMaxSources]; + __shared__ uint32_t segBegin[kMaxSources + 1]; + __shared__ __align__(16) ClusterRecords rec; + __shared__ CtaTally tally; + __shared__ int32_t epochShared; + __shared__ __align__(16) BridgeShared bridgeRow; + __shared__ int32_t emitScratch[kWarmScratchWords]; + int32_t const world = p.worldSize; + int32_t const rank = p.rank; + int32_t const src = blockIdx.x; + int32_t const row = blockIdx.y; + SmemLayout const s = carveSmem(smemRaw, p.localCapacity, world); + int32_t* outRow = p.out + static_cast(row) * kTopK; + clusterInit(rec); + if (threadIdx.x <= kSlotPositivesBefore - kSlotCountBefore) + { + s.scal[kSlotCountBefore + threadIdx.x] = 0u; + } + int32_t const resetOffset = resetHeaderOffset(); + // Every global read (lengths, logits, the epoch word, the workspace, the bridge inputs) comes after the grid + // dependency wait, so a programmatic early launch is safe whatever the preceding kernels write. + asm volatile("griddepcontrol.wait;" ::: "memory"); + markPhase(p, 10); + // Thread 0 and the reset warps read this CTA's epoch word, written only by the previous exchange launch of this + // op, before the key loads; nothing waits for the value until the key loads are done. Thread 0 publishes it to the + // CTA at the block barrier that follows the key loads, and the reset warps start the copies of the previous + // buffer's headers right before that barrier (slot offsets computed before the wait, from the grid only). + int32_t* const epochWord = p.state + kEpochBase + row * kSourcesPerRank + src; + int32_t epochRead = 0; + if (threadIdx.x == 0 || static_cast(threadIdx.x) >= kFirstResetWarp * 32) + { + epochRead = *reinterpret_cast(epochWord); + } + int32_t const prevBufRead = (epochRead + kNumLamportBuffers - 1) % kNumLamportBuffers; + int32_t const kvLen = p.kvLens[0]; + int32_t const rowLen = max(0, kvLen - p.nextN + row + 1); + // Identity launch (every row has at most kTopK tokens; the same decision on every rank): no exchange, and + // neither the workspace nor the epoch words are touched, so identity launches never shorten the reuse chain of + // the Lamport buffers (see crossRankTopkGlobalEmit.h). + bool const exchange = kvLen - p.nextN + p.numRows > kTopK; + // Thread 0 advances this CTA's epoch word (the CTAs of the last row also those of the rows this launch does not + // have, so every word counts every exchange launch), after a block barrier that follows every read of the word in + // this CTA. Nothing else of this launch reads the epoch words, and the next launch reads them only after its grid + // dependency wait. + auto advanceEpoch = [&]() + { + if (exchange && threadIdx.x == 0) + { + int32_t const rows = row == static_cast(gridDim.y) - 1 ? kMaxRows - row : 1; + for (int32_t r = 0; r < rows; ++r) + { + epochWord[r * kSourcesPerRank] = epochRead + 1; + } + } + }; + + if (rowLen <= kTopK) + { + // Every rank agrees on this branch: identity selection, no exchange for this row. Each CTA of the row writes + // kTopK / kClusterCtas positions with their bridge outputs (bridge inputs read directly). + constexpr int32_t kIdentityPerCta = kTopK / kClusterCtas; + if (static_cast(threadIdx.x) < kIdentityPerCta) + { + CpTopkBridgeParams const& b = p.bridge; + int32_t const pos = src * kIdentityPerCta + static_cast(threadIdx.x); + int32_t const tok = pos < rowLen ? pos : -1; + outRow[pos] = tok; + int32_t const req = b.reqIdx[row] - b.reqIdxBias; + BridgeRowOut const bridgeOut = bridgeRowOut(b, row, bridgeCopyRow(b, row, loadAccepted(b, row))); + writeBridge(b, bridgeOut, static_cast(pos), tok, directBase(b, req, tok), + [&](int32_t g) { return b.layerIds != nullptr ? b.layerIds[g] : b.layerId; }); + } + if (exchange) + { + fetchResetHeader(p, s, prevBufRead, resetOffset); + resetSlots(p, s, prevBufRead); + markPhase(p, 9); + } + __syncthreads(); + advanceEpoch(); + clusterInitWait(); + } + else + { + int32_t const localLen = min(localLength(rowLen, rank, world), p.localCapacity); + float const* rowLogits = p.logits + static_cast(row) * p.logitsStride; + // The first reset warp starts the copies of the row's request index, layer ids and accepted count (completed + // with the reset). + if (static_cast(threadIdx.x) / 32 == kFirstResetWarp) + { + fetchRowScalars(p, bridgeRow, row); + } + + // ---- Round 1: this source's local top keys --------------------------------------------------------- + int32_t const n = loadSource(s, rowLogits, localLen, src, true); + if (threadIdx.x == 0) + { + epochShared = epochRead; + } + // The request index copy is complete before the barrier (only its group is committed at this point). + if (static_cast(threadIdx.x) == kFirstResetWarp * 32 + kReqLane) + { + asm volatile("cp.async.wait_group 0;" ::: "memory"); + } + // The reset warps start fetching the headers of their share of the previous exchange launch's buffer (an + // exchange row implies an exchange launch); the values are first used by the reset after the send. + fetchResetHeader(p, s, prevBufRead, resetOffset); + __syncthreads(); + markPhase(p, 1); + advanceEpoch(); + uint32_t const cap = static_cast(roundCap(world)); + uint32_t const need = static_cast(roundNeed(world)); + SelectOut sel{0u, 0u, 0u, static_cast(n), true, false, 0}; + uint32_t selLo = 0u; + int32_t passInfo = 0; + if (n > static_cast(cap)) + { + // The cut lies among the keys of one sign (see kPositiveKey); the digits span that sign's key range. + SignRange const r = signRange(readStats(s.stats), need); + selLo = r.lo; + SmemKeys const keys{s.chunk, nullptr}; + sel = blockSelect(s, keys, n, r.lo, r.hi, need, cap, r.lo, r.bits, debugSlots(p, 32), r.aboveHi); + passInfo = sel.passes + (r.negative ? 10 : 0); + } + markPhase(p, 2); + int32_t const sendBuf = epochShared % kNumLamportBuffers; + uint32_t const sent = sendSource(s, sel, n, selLo, uint64_t{kEmptyWord}, cap, + p.mcWorkspace + round1Offset(sendBuf, rank * kSourcesPerRank + src, row), src, rank, world, + reinterpret_cast(s.mkey), debugSlots(p, 0)); + putDebug(debugSlots(p, 0), 46, sent); + putDebug(debugSlots(p, 0), 47, passInfo); + markPhase(p, 3); + + // The reset warps reset the previous exchange launch's buffer, whose candidates every rank has consumed (a + // peer writes it again only two exchange launches later, after this launch completed on this rank), while + // the poll and header warps wait for this launch's candidates. + resetSlots(p, s, (epochShared + kNumLamportBuffers - 1) % kNumLamportBuffers); + int32_t const buf = epochShared % kNumLamportBuffers; + if (threadIdx.x == kFirstResetWarp * 32) + { + int64_t* const d = debugSlots(p, 0); + if (d != nullptr) + { + d[9] = static_cast(globalTimerNs()); + } + // The row's pool and copy rows (the accepted count is visible after resetSlots' warp barrier), read by + // the emission after the block barrier that follows the candidate poll. + CpTopkBridgeParams const& b = p.bridge; + int32_t const copyRow = bridgeCopyRow(b, row, sharedAccepted(b, bridgeRow)); + bridgeRow.poolRow = b.out + row * b.outStride0; + bridgeRow.copyRow = copyRow >= 0 ? b.copyDst + copyRow * b.copyStride0 : nullptr; + // The emission's instructions, run once now while this warp waits for the poll (see emitCandidates); + // the layer ids are complete after resetSlots. + warmEmission(b.groupSize, emitScratch, bridgeRow.layer); + } + + // ---- Merge: the row's cluster ----------------------------------------------------------------------- + WaitCtx const wait{p.state, p.timeoutNs, rank}; + if (clusterMergeRow(p, s, rec, tally, buf, row, rowLen, outRow, wait, bridgeRow) && src == 0) + { + // A source withheld keys above the threshold: crossRankTopk's merge, with its second round, on CTA 0. + mergeRow(p, s, segSlots, segBegin, buf, row, rowLen, localLen, rowLogits, outRow, wait, bridgeRow); + int64_t* const dbg = debugSlots(p, 0); + putDebug(dbg, 63, dbg != nullptr ? dbg[63] + 100 : 0); + } + markPhase(p, 7); + } + markPhase(p, 8); +} + +#endif + +} // namespace + +size_t cpTopkSmemBytes(int32_t localCapacity, int32_t worldSize) +{ + return smemBytes(localCapacity, worldSize); +} + +int32_t cpTopkMaxLocalCapacity(int32_t worldSize) +{ + static int32_t cached[kMaxRanks + 1] = {}; + if (worldSize < 1 || worldSize > kMaxRanks) + { + return 0; + } + if (cached[worldSize] == 0) + { + int device = 0; + cudaGetDevice(&device); + int maxSmemOptin = 0; + cudaDeviceGetAttribute(&maxSmemOptin, cudaDevAttrMaxSharedMemoryPerBlockOptin, device); + int32_t cap = 0; + for (int32_t c = 1 << 22; c >= kPageTokens * kSourcesPerRank; c -= kPageTokens * kSourcesPerRank) + { + if (smemBytes(c, worldSize) + kStaticSmemBytes <= static_cast(maxSmemOptin)) + { + cap = c; + break; + } + } + cached[worldSize] = cap; + } + return cached[worldSize]; +} + +int32_t cpTopkMaxActiveClusters(int32_t worldSize) +{ + // Queried once per world size and device (the value + 1; 0 = not queried yet). + constexpr int32_t kMaxDevices = 64; + static int32_t cached[kMaxDevices][kMaxRanks + 1] = {}; + if (worldSize < 1 || worldSize > kMaxRanks) + { + return 0; + } + int device = 0; + cudaGetDevice(&device); + if (device >= 0 && device < kMaxDevices && cached[device][worldSize] > 0) + { + return cached[device][worldSize] - 1; + } + int32_t const capacity = cpTopkMaxLocalCapacity(worldSize); + if (capacity == 0) + { + return 0; + } + size_t const smem = smemBytes(capacity, worldSize); + if (cudaFuncSetAttribute(cpTopkZ11Kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(smem)) + != cudaSuccess) + { + (void) cudaGetLastError(); + return 0; + } + cudaLaunchConfig_t config{}; + config.gridDim = dim3(kSourcesPerRank, kMaxRows); + config.blockDim = dim3(kThreads); + config.dynamicSmemBytes = smem; + cudaLaunchAttribute attr{}; + attr.id = cudaLaunchAttributeClusterDimension; + attr.val.clusterDim.x = kSourcesPerRank; + attr.val.clusterDim.y = 1; + attr.val.clusterDim.z = 1; + config.attrs = &attr; + config.numAttrs = 1; + int clusters = 0; + if (cudaOccupancyMaxActiveClusters(&clusters, cpTopkZ11Kernel, &config) != cudaSuccess) + { + (void) cudaGetLastError(); + return 0; + } + if (device >= 0 && device < kMaxDevices) + { + cached[device][worldSize] = clusters + 1; + } + return clusters; +} + +int64_t cpTopkDefaultTimeoutMs() +{ + static int64_t const timeoutMs = []() + { + char const* env = std::getenv("TRTLLM_STAIRCASE_GLM_LL_CROSS_RANK_TOPK_TIMEOUT_MS"); + if (env == nullptr || env[0] == '\0') + { + return kDefaultTimeoutMs; + } + char* end = nullptr; + long long const v = std::strtoll(env, &end, 10); + if (end == env || *end != '\0' || v <= 0) + { + throw std::invalid_argument( + std::string( + "glm_moe_dsa_cross_rank_topk_global_emit: bad TRTLLM_STAIRCASE_GLM_LL_CROSS_RANK_TOPK_TIMEOUT_MS=") + + env); + } + return static_cast(v); + }(); + return timeoutMs; +} + +void invokeCpTopk(CpTopkParams const& inParams, cudaStream_t stream) +{ + if (inParams.numRows <= 0) + { + return; + } + CpTopkParams params = inParams; + if (params.timeoutNs == 0) + { + params.timeoutNs = static_cast(cpTopkDefaultTimeoutMs()) * 1000000ull; + } + if (params.numRows > kMaxRows || params.worldSize <= 0 || params.worldSize > kMaxRanks || params.rank < 0 + || params.rank >= params.worldSize) + { + throw std::invalid_argument("glm_moe_dsa_cross_rank_topk_global_emit: unsupported rows/world/rank"); + } + CpTopkBridgeParams const& b = params.bridge; + if (b.out == nullptr || b.reqIdx == nullptr || b.blockTable == nullptr || b.groupSize < 1 + || b.groupSize > kMaxGroupSize || b.maxBlocksPerReq < 0 || (b.copyMode != 0 && b.copyMode != 1) + || b.copyNextN < 1 || b.copyRows < 0 + || (b.copyDst != nullptr + && (b.copyMode == 0 ? b.copyRows > params.numRows + : static_cast(b.copyRows) * b.copyNextN > params.numRows))) + { + throw std::invalid_argument("glm_moe_dsa_cross_rank_topk_global_emit: bad bridge parameters"); + } + // The cluster merge stages two tokens per polled candidate vector below kBaseOffsetWords and their pool bases + // beyond it. + int32_t const pollThreads = params.worldSize * (roundCap(params.worldSize) / 2); + if (2 * pollThreads > kBaseOffsetWords || kBaseOffsetWords + 2 * pollThreads > mergeCapacity(params.worldSize)) + { + throw std::invalid_argument( + "glm_moe_dsa_cross_rank_topk_global_emit: the candidate staging exceeds the merge area"); + } + size_t const smem = smemBytes(params.localCapacity, params.worldSize); + // The attribute is per function and device; raise it only when a larger footprint is requested. + static thread_local size_t configuredSmem[64] = {}; + int device = 0; + cudaGetDevice(&device); + cudaError_t err = cudaSuccess; + if (device >= 64 || configuredSmem[device] < smem) + { + err = cudaFuncSetAttribute( + cpTopkZ11Kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(smem)); + if (err != cudaSuccess) + { + throw std::runtime_error( + std::string("glm_moe_dsa_cross_rank_topk_global_emit: smem attribute: ") + cudaGetErrorString(err)); + } + if (device < 64) + { + configuredSmem[device] = smem; + } + } + // Programmatic dependent launch as for the other TensorRT-LLM kernels (on unless TRTLLM_ENABLE_PDL=0). + static bool const enablePdl = []() + { + char const* env = std::getenv("TRTLLM_ENABLE_PDL"); + return !(env != nullptr && env[0] == '0' && env[1] == '\0'); + }(); + cudaLaunchConfig_t config{}; + config.gridDim = dim3(kSourcesPerRank, params.numRows); + config.blockDim = dim3(kThreads); + config.dynamicSmemBytes = smem; + config.stream = stream; + cudaLaunchAttribute attrs[2]; + attrs[0].id = cudaLaunchAttributeClusterDimension; + attrs[0].val.clusterDim.x = kSourcesPerRank; + attrs[0].val.clusterDim.y = 1; + attrs[0].val.clusterDim.z = 1; + attrs[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[1].val.programmaticStreamSerializationAllowed = enablePdl ? 1 : 0; + config.attrs = attrs; + config.numAttrs = 2; + err = cudaLaunchKernelEx(&config, cpTopkZ11Kernel, params); + if (err == cudaSuccess) + { + err = cudaGetLastError(); + } + if (err != cudaSuccess) + { + throw std::runtime_error( + std::string("glm_moe_dsa_cross_rank_topk_global_emit: launch: ") + cudaGetErrorString(err)); + } +} + +} // namespace kernels::glm_moe_dsa::cross_rank_topk_global_emit + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopkGlobalEmit.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopkGlobalEmit.h new file mode 100644 index 000000000000..c51b6fff00fe --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/crossRankTopkGlobalEmit.h @@ -0,0 +1,252 @@ +/* + * 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 +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa::cross_rank_topk_global_emit +{ + +// Context-partitioned DSA indexer top-k with a cluster merge that also writes the top-k bridge's outputs. +// +// It is the cluster-merge kernel (glmMoeDsa/crossRankTopkClusterMerge.h; the selection, exchange, workspace, epoch, +// reset and timeout below are that kernel's, unchanged) whose emission also writes, at every output position, what the +// top-k index bridge (glmMoeDsa/ropeKvAppend.h, topkBridgeKernel) writes from the selected row: +// - the global KV-pool rows of every layer of the remap group: pool[g][row][pos] = remapOne(tok, layer g), i.e. +// -1 for tok < 0 or a page beyond the request's block table, else +// blockTable[req, tok / 64] * strideFactor + tok % 64 + layerId[g] * 64 (int32 arithmetic, as the bridge); +// - the native top-k copy row of this row, if any: copy mode 0 copies rows [0, copyRows) to copy rows +// copyRowStart + row; copy mode 1 copies, for request i < copyRows, its last accepted row +// i * copyNextN + clamp(numAccepted[i] - 1, 0, copyNextN - 1) (i * copyNextN + copyNextN - 1 without +// numAccepted) to copy row copyRowStart + i. +// In the cluster merge, the pool bases come off the critical path: right after the candidate poll, every poll thread +// starts 4-byte cp.async copies of its (at most two) candidates' block-table entries blockTable[req, tok / 64] into +// its own shared-memory words, which complete while the cluster selects the threshold; the emission then only adds +// and stores. The emission is a called function that one thread of the CTA first runs into shared scratch while it +// waits for the poll, so its instructions are cached on the SM when the emission runs. The row's request index, layer +// ids and accepted count are copied into shared memory by one reset warp right after the grid dependency wait. +// Identity rows and the CTA-0 second round read the block table directly. Every bridge input (request indices, block +// table, layer ids, accepted counts) is read after the grid dependency wait, like every other global access of the +// kernel. +// +// The merge selects exactly what the kernel of glmMoeDsa/crossRankTopk.h selects, in the same output order, with the +// same context partition, exchange sources, round-1 / round-2 protocol and Lamport workspace layout: +// - global 64-token page p belongs to rank p % worldSize; within a rank, each verify row is handled by +// kSourcesPerRank CTAs, CTA s owning the 8-token block (s + page) % kSourcesPerRank of every local page +// (exchange source q = rank * kSourcesPerRank + s); +// - each source selects its local top keys with crossRankTopk's rule and broadcasts them as (order-preserving key, +// global token) pairs in source order, plus a header (count, exclusive bound of the withheld keys; the cluster +// merge adds the key statistics of the keys sent), through the multicast Lamport workspace. +// What differs is how the row is merged and how a launch ends: +// - the kSourcesPerRank CTAs of one row form a thread-block cluster. CTA s merges the candidates of sources +// [s * worldSize, (s + 1) * worldSize) (at 8 ranks: the 8 sources of rank s). The cluster runs crossRankTopk's +// merge select (the same digits and early-stop rules) on the cluster-summed histograms, and every CTA writes its +// selected candidates directly into the output row at crossRankTopk's positions (keys above the threshold in +// candidate order, then ties in candidate order, then -1). Each round-1 slot header carries its source's key +// statistics, so every CTA derives the row's statistics from the headers it polls; the CTAs exchange only the +// select's histogram totals, with st.async stores into each other's shared memory completed on transaction +// barriers, and derive every CTA's output counts from the per-CTA tallies of the select. The only cluster-wide +// barrier orders the barriers' initialization before the first exchange. A row in which a source withheld keys +// above the threshold (identical decision on every rank) is merged instead by CTA 0 with crossRankTopk's single-CTA +// merge, including its second round; +// - every CTA keeps its own epoch word (one per (row, source CTA) slot): its thread 0 reads the word after the grid +// dependency wait (only the previous exchange launch of this op writes it), the value is published to the CTA at +// the block barrier that follows the key loads, and the CTA advances the word with a plain store after that +// barrier (nothing else in the launch reads it; the next launch reads it after its own grid dependency wait); the +// CTAs of the last row also advance the words of the rows the launch does not have, so every word counts the +// exchange launches. There is no arrival counter and no end-of-kernel fence. +// +// Rows of at most kTopK valid tokens produce the identity selection followed by -1, without an exchange. A launch whose +// rows all have at most kTopK valid tokens (the same on every rank: the lengths come from the shared KV length) is an +// identity launch: it writes neither the workspace nor the epoch words. Every other launch is an exchange launch. +// Exchange launches rotate kNumLamportBuffers buffers by the epoch, which counts exchange launches; in exchange launch +// x every CTA resets its share of its own rank's copy of the buffer of exchange launch x - 1 between its send and its +// merge (the headers of its share are fetched once the epoch is known, after the grid dependency wait). Every rank's +// merges of x - 1 have read their copy of it, and a peer writes it again only in exchange launch x + 2, which starts +// after the peer's merges of x + 1 received this rank's sources of x + 1, i.e. after this rank completed launch x +// (identity launches in between neither write nor wait, so they do not shorten this chain). Every global memory +// access of the kernel follows its grid dependency wait, which orders it after the completion (and memory +// visibility) of the preceding kernels of the stream, including the previous launch of this op; the kernel never +// triggers its dependents early. +// +// A merge that waits longer than the timeout (CpTopkParams::timeoutNs, %globaltimer) for a peer's words records the +// wait in the state's error word and aborts the kernel with a trap: the exchange cannot complete correctly any more, +// and a late peer write could land in a buffer this rank reuses, so the run fails instead of continuing. + +//! Selected indices per verify row (DSA index_topk). +constexpr int32_t kTopK = 2048; +//! Ranks of the context-parallel group. +constexpr int32_t kMaxRanks = 8; +//! Verify rows per call. +constexpr int32_t kMaxRows = 8; +//! Tokens per indexer K page; the context is partitioned by whole pages. +constexpr int32_t kPageTokens = 64; +//! CTAs (exchange sources) per rank and verify row; also the cluster size (one cluster per row). +constexpr int32_t kSourcesPerRank = 8; +//! Exchange sources per row. +constexpr int32_t kMaxSources = kMaxRanks * kSourcesPerRank; +//! Round-1 slots per (source, row). +constexpr int32_t kSourceCapacity = 256; +//! Keys each saturated source aims to send in round 1 at 8 ranks (crossRankTopk's rule: three times the even share of +//! kTopK; refining stops at the first bin edge with at most four times the even share of keys at or above it). +constexpr int32_t kSourceBudget = 96; +//! Round-2 slots per (source, row). +constexpr int32_t kRound2Capacity = 2048; +//! Threads per CTA. +constexpr int32_t kThreads = 1024; +//! Lamport buffers rotated by the epoch. +constexpr int32_t kNumLamportBuffers = 3; +//! 32-bit word that marks an empty Lamport slot; never a valid key, token index, count or bound. +constexpr uint32_t kEmptyWord = 0xFFFFFFFFu; + +//! 32-bit words of one (source, row) slot: a 32-byte header then two words per candidate, so candidates can be written +//! and read with 16-byte vector accesses. Round-1 header: count, exclusive bound of the withheld keys, then the key +//! statistics of the keys sent (largest key, smallest key + 1, keys >= +0.0, largest key below +0.0, smallest key +//! >= +0.0 plus 1; the "+ 1" maps the "none" value 0xFFFFFFFF to 0), then 0. Round-2 header: count, bound. +constexpr int64_t kSlotHeaderWords = 8; +constexpr int64_t kRound1SlotWords = kSlotHeaderWords + 2 * int64_t{kSourceCapacity}; +constexpr int64_t kRound2SlotWords = kSlotHeaderWords + 2 * int64_t{kRound2Capacity}; +//! 32-bit words of one Lamport buffer (round-1 region, then round-2 region). +constexpr int64_t kBufferWords = int64_t{kMaxSources} * kMaxRows * (kRound1SlotWords + kRound2SlotWords); +//! 32-bit words of the whole exchange workspace; every word must start as kEmptyWord. +constexpr int64_t kWorkspaceWords = kNumLamportBuffers * kBufferWords; +//! int32 words of the per-rank device state, zero-initialized: [0, 2) unused, [2] second rounds taken, [3] error word +//! (nonzero = the wait that timed out, written right before the kernel aborts), then one epoch word per +//! (row, source CTA) slot: state[kEpochBase + row * kSourcesPerRank + s]. +constexpr int32_t kEpochBase = 4; +constexpr int32_t kStateWords = kEpochBase + kMaxRows * kSourcesPerRank; +//! Default wait for a peer's Lamport words before the kernel aborts +//! (TRTLLM_STAIRCASE_GLM_LL_CROSS_RANK_TOPK_TIMEOUT_MS overrides it). +constexpr int64_t kDefaultTimeoutMs = 120000; + +//! Values of the error word. +enum CpTopkWait : int32_t +{ + kWaitRound1Header = 1, + kWaitRound1Candidates = 2, + kWaitRound2Header = 3, + kWaitRound2Candidates = 4 +}; + +//! Most layers of one remap group (the bridge's kMaxGroupSize). +constexpr int32_t kMaxGroupSize = 8; +//! Words of the merge area (the token staging of the cluster merge) its candidates' tokens use: at most 520 16-byte +//! candidate vectors per CTA (world sizes 5 to 7), two tokens each. The candidates' pool bases follow them, at the +//! same per-thread offsets. +constexpr int32_t kBaseOffsetWords = 1040; + +//! Outputs of the top-k bridge written at the emission (see the comment at the top). +struct CpTopkBridgeParams +{ + //! [groupSize, numRows, kTopK] int32 pool rows with strides (outStrideG, outStride0, 1). + int32_t* out = nullptr; + int64_t outStrideG = 0; + int64_t outStride0 = 0; + int32_t groupSize = 1; + //! Optional [groupSize] int32 layer ids; when null every member uses layerId. + int32_t const* layerIds = nullptr; + int32_t layerId = 0; + //! [numRows] int32 request index per row, before subtracting reqIdxBias. + int32_t const* reqIdx = nullptr; + int32_t reqIdxBias = 0; + //! int32 block table [requests, maxBlocksPerReq] with strides (btStride0, btStride1), block size kPageTokens. + int32_t const* blockTable = nullptr; + int64_t btStride0 = 0; + int64_t btStride1 = 1; + int32_t maxBlocksPerReq = 0; + int32_t strideFactor = 0; + //! Optional copy destination [>= copyRowStart + copyRows, >= kTopK] int32 with row stride copyStride0. + int32_t* copyDst = nullptr; + int64_t copyStride0 = 0; + int32_t copyRows = 0; + int32_t copyRowStart = 0; + int32_t copyMode = 0; + int32_t copyNextN = 1; + //! Optional [copyRows] accepted counts of copy mode 1, int32 or int64. + void const* numAccepted = nullptr; + bool numAcceptedIs64 = false; +}; + +struct CpTopkParams +{ + //! [numRows, logitsStride] fp32 logits of this rank's local columns. + float const* logits; + int64_t logitsStride; + //! Local columns the logits view holds (each row's local length is clamped to it). + int32_t localCapacity; + //! [1] int32 global indexer KV length of the request; row r has kvLen - nextN + r + 1 valid tokens. + int32_t const* kvLens; + //! [numRows, kTopK] int32 output: global token indices, -1 for unused slots. + int32_t* out; + int32_t numRows; + int32_t nextN; + int32_t rank; + int32_t worldSize; + //! Multicast address of the exchange workspace (kWorkspaceWords words, identical layout on every rank). + uint32_t* mcWorkspace; + //! This rank's unicast address of the same workspace. + uint32_t* ucWorkspace; + //! [kStateWords] int32 per-rank device state, zero-initialized. + int32_t* state; + //! Optional [numRows, kSourcesPerRank, kDebugSlots] int64 phase timestamps (diagnosis only; nullptr = off). + int64_t* debugTimes = nullptr; + //! Longest wait for a peer's words before the kernel aborts (ns); 0 = the default (kDefaultTimeoutMs, or + //! TRTLLM_STAIRCASE_GLM_LL_CROSS_RANK_TOPK_TIMEOUT_MS). + uint64_t timeoutNs = 0; + //! The bridge outputs (out, reqIdx and blockTable are required). + CpTopkBridgeParams bridge; +}; + +//! Timestamps per CTA in CpTopkParams::debugTimes: [0, 16) %globaltimer (ns) and [16, 32) clock64 per phase (0 start, +//! 10 grid dependency wait returned, 1 keys loaded, 2 local select, 3 sent, 9 previous buffer reset (%globaltimer of +//! the first reset warp only), 4 candidates polled, 5 cluster statistics agreed, 6 merge select, 7 row written, 8 end), +//! [32, 42) local select (keys, then clock64 at the start of / after the histogram of / after the bin search of each +//! pass), [42, 45) send (start, keys taken, stores issued), 46 round-1 keys sent by this source, 47 local select passes +//! (+10 when the cut was searched among the negative keys), [48, 58) merge select (keys = this CTA's candidates; per +//! pass clock64 at its start / after the group totals arrived / after the bin was decided), clock64 of the merge: 45 +//! candidates polled, 61 statistics combined, 58 emission positions scanned, 60 row written; 62 candidates of the row +//! (all CTAs), 63 merge select passes (+10 likewise; +100 when the row took crossRankTopk's single-CTA merge, which +//! then reuses 45 and 58-61 as crossRankTopk does); clock64 and values of thread 0. +constexpr int32_t kDebugSlots = 64; + +//! Effective wait timeout of a launch with params.timeoutNs == 0 (ms). +int64_t cpTopkDefaultTimeoutMs(); + +//! Dynamic shared memory the kernel needs (bytes). +size_t cpTopkSmemBytes(int32_t localCapacity, int32_t worldSize); + +//! Largest local capacity the kernel accepts for a world size (shared memory; the cluster merge's limit). +int32_t cpTopkMaxLocalCapacity(int32_t worldSize); + +//! Clusters of the kernel's launch configuration (largest local capacity) that can be resident at once on this +//! device (cudaOccupancyMaxActiveClusters, queried once per world size and device); 0 when the query fails. No row +//! waits for another row of its rank, but a row's cluster waits for the same row's sources on every peer rank, and +//! the ranks may make their rows resident in different orders; the caller therefore requires kMaxRows co-resident +//! clusters, so every cluster of a launch can be resident at once (other kernels sharing the SMs are finite and do +//! not wait for this one). +int32_t cpTopkMaxActiveClusters(int32_t worldSize); + +void invokeCpTopk(CpTopkParams const& params, cudaStream_t stream); + +} // namespace kernels::glm_moe_dsa::cross_rank_topk_global_emit + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/denseMlp.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/denseMlp.cu new file mode 100644 index 000000000000..7837384be9b1 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/denseMlp.cu @@ -0,0 +1,453 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Weight-streaming BF16 kernels for the dense (non-MoE) MLP at decode row counts (M <= 8). +// +// Both projections are bound by reading the BF16 weight matrix once from HBM. Each CTA owns a +// 16-row weight tile and its warps split the contraction dimension. Every thread loads 16 bytes of +// two weight rows and of one activation row per step and feeds them to two m16n8k16 tensor-core +// MMAs with the weight on the MMA M side and the (at most 8) activation rows on the MMA N side. +// The contraction index assigned to an MMA k-slot depends only on the thread's position inside its +// quad, so weight and activation fragments use the same permutation of k and the sum is exact up to +// FP32 summation order. Each quad therefore reads 64 contiguous bytes of one weight row per step. +// +// The gate/up kernel places 8 gate rows and the matching 8 up rows in one tile, so each thread ends +// with gate and up of the same output column and applies SwiGLU in registers. +// +// Bytes in flight: a thread holds one chunk of steps in registers, and one lane per quad prefetches +// the quad's row segment of the next chunk into L2 with cp.async.bulk.prefetch, which keeps the HBM +// queue full without spending registers. +// +// Programmatic dependent launch: weights never depend on the preceding kernel, so the first chunk +// (registers) and the L2 prefetch of the next one are issued before griddepcontrol.wait; dependents +// are released right after the wait. The tilings keep each grid within one wave of resident CTAs, +// so early-launched dependent CTAs that wait in griddepcontrol.wait cannot starve the primary grid. + +#include "tensorrt_llm/kernels/glmMoeDsa/denseMlp.h" + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa::dense_mlp +{ + +namespace +{ + +using Bf16 = __nv_bfloat16; + +constexpr int kWarpSize = 32; +//! Weight rows per CTA tile (MMA M). +constexpr int kTileRows = 16; +//! Output columns per gate/up CTA: 8 gate rows + 8 up rows form one tile. +constexpr int kGateUpCols = kTileRows / 2; +//! BF16 values per 16-byte vector load. +constexpr int kVecElems = 8; +//! Threads per MMA quad; a quad covers kQuadThreads * kVecElems contraction elements per step. +constexpr int kQuadThreads = 4; +//! Contraction elements consumed by one warp step (two m16n8k16 MMAs). +constexpr int kKPerStep = kQuadThreads * kVecElems; + +//! Compile-time tiling of one kernel instance. +//! \tparam kWarps warps per CTA; each warp owns a contiguous 1/kWarps slice of the contraction. +//! \tparam kSteps steps per register chunk. +//! \tparam kMinBlocks minimum resident CTAs per SM requested from the compiler (register budget). +//! \tparam kCachedX load activations with ld.global.ca instead of the non-coherent path. +//! \tparam kL2Ahead chunks of each weight row segment prefetched into L2 ahead of the register loads. +template +struct Tiling +{ + static constexpr int kWarps = kWarps_; + static constexpr int kSteps = kSteps_; + static constexpr int kMinBlocks = kMinBlocks_; + static constexpr bool kCachedX = kCachedX_; + static constexpr int kL2Ahead = kL2Ahead_; + static constexpr int kKPerChunk = kKPerStep * kSteps; +}; + +__device__ __forceinline__ void pdlWait() +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif +} + +__device__ __forceinline__ void pdlLaunchDependents() +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); +#endif +} + +//! Streaming weight load: not kept in L1, 256-byte L2 fetch granularity. +__device__ __forceinline__ uint4 loadWeight(Bf16 const* ptr) +{ + uint4 r; + asm volatile("ld.global.nc.L1::no_allocate.L2::256B.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(r.x), "=r"(r.y), "=r"(r.z), "=r"(r.w) + : "l"(ptr)); + return r; +} + +//! Activation load: every CTA reads the same few rows, so keep them in L1. +template +__device__ __forceinline__ uint4 loadActivation(Bf16 const* ptr) +{ + if constexpr (kCached) + { + uint4 r; + asm volatile("ld.global.ca.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(r.x), "=r"(r.y), "=r"(r.z), "=r"(r.w) + : "l"(ptr)); + return r; + } + else + { + return __ldg(reinterpret_cast(ptr)); + } +} + +__device__ __forceinline__ void prefetchL2(void const* ptr, uint32_t bytes) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;" ::"l"(ptr), "r"(bytes) : "memory"); +#endif +} + +__device__ __forceinline__ void mmaBf16( + float (&d)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t b0, uint32_t b1) +{ + asm("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, " + "{%0, %1, %2, %3};" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)); +} + +//! Registers of one chunk: kSteps steps of two weight rows and one activation row. +template +struct Chunk +{ + uint4 w0[kSteps]; + uint4 w1[kSteps]; + uint4 x[kSteps]; +}; + +template +__device__ __forceinline__ void loadChunkWeights(Chunk& chunk, Bf16 const* w0, Bf16 const* w1, int chunkIdx) +{ + int const offset = chunkIdx * kKPerStep * kSteps; +#pragma unroll + for (int u = 0; u < kSteps; ++u) + { + chunk.w0[u] = loadWeight(w0 + offset + u * kKPerStep); + chunk.w1[u] = loadWeight(w1 + offset + u * kKPerStep); + } +} + +template +__device__ __forceinline__ void loadChunkActivations(Chunk& chunk, Bf16 const* x, bool hasRow, int chunkIdx) +{ + int const offset = chunkIdx * kKPerStep * kSteps; +#pragma unroll + for (int u = 0; u < kSteps; ++u) + { + chunk.x[u] = hasRow ? loadActivation(x + offset + u * kKPerStep) : make_uint4(0U, 0U, 0U, 0U); + } +} + +//! Two MMAs per step. MMA A uses contraction elements 0..3 of the thread's 8-element vector, MMA B +//! elements 4..7. Fragment registers: a0/a2 = weight row "group", a1/a3 = weight row "group + 8", +//! b0/b1 = activation row "group". +template +__device__ __forceinline__ void computeChunk(float (&accA)[4], float (&accB)[4], Chunk const& chunk) +{ +#pragma unroll + for (int u = 0; u < kSteps; ++u) + { + mmaBf16(accA, chunk.w0[u].x, chunk.w1[u].x, chunk.w0[u].y, chunk.w1[u].y, chunk.x[u].x, chunk.x[u].y); + mmaBf16(accB, chunk.w0[u].z, chunk.w1[u].z, chunk.w0[u].w, chunk.w1[u].w, chunk.x[u].z, chunk.x[u].w); + } +} + +//! Streams the CTA's two weight rows per thread over the warp's contraction slice and returns the +//! CTA-reduced accumulator fragment in warp 0 (other warps return false). +//! acc[0..1] = tile row "group", activation rows 2 * tig + {0, 1}; acc[2..3] = tile row "group + 8". +template +__device__ __forceinline__ bool streamTile(float (&acc)[4], Bf16 const* __restrict__ x, Bf16 const* __restrict__ weight, + int numRows, int kDim, int row0, int row1) +{ + constexpr int kWarps = T::kWarps; + constexpr int kSteps = T::kSteps; + __shared__ float reduceBuf[kWarps][4][kWarpSize]; + + int const lane = threadIdx.x % kWarpSize; + int const warp = threadIdx.x / kWarpSize; + int const group = lane / kQuadThreads; + int const tig = lane % kQuadThreads; + + int const kPerWarp = kDim / kWarps; + int const kBase = warp * kPerWarp + tig * kVecElems; + Bf16 const* w0 = weight + static_cast(row0) * kDim + kBase; + Bf16 const* w1 = weight + static_cast(row1) * kDim + kBase; + bool const hasRow = group < numRows; + Bf16 const* xRow = x + static_cast(hasRow ? group : 0) * kDim + kBase; + int const numChunks = kPerWarp / T::kKPerChunk; + + // One lane per quad prefetches the quad's contiguous row segments of a chunk into L2. + auto prefetchChunkL2 = [&](int chunkIdx) + { + if constexpr (T::kL2Ahead > 0) + { + if (tig == 0 && chunkIdx < numChunks) + { + uint32_t const bytes = T::kKPerChunk * static_cast(sizeof(Bf16)); + prefetchL2(w0 + chunkIdx * T::kKPerChunk, bytes); + prefetchL2(w1 + chunkIdx * T::kKPerChunk, bytes); + } + } + }; + + float accA[4] = {0.F, 0.F, 0.F, 0.F}; + float accB[4] = {0.F, 0.F, 0.F, 0.F}; + + // Weights do not depend on the preceding kernel: request them before the dependency wait. + Chunk chunk; + loadChunkWeights(chunk, w0, w1, 0); +#pragma unroll + for (int a = 1; a <= T::kL2Ahead; ++a) + { + prefetchChunkL2(a); + } + pdlWait(); + pdlLaunchDependents(); + loadChunkActivations(chunk, xRow, hasRow, 0); + +#pragma unroll 1 + for (int c = 0; c < numChunks; ++c) + { + if (c > 0) + { + loadChunkWeights(chunk, w0, w1, c); + loadChunkActivations(chunk, xRow, hasRow, c); + } + prefetchChunkL2(c + 1 + T::kL2Ahead); + computeChunk(accA, accB, chunk); + } + +#pragma unroll + for (int e = 0; e < 4; ++e) + { + acc[e] = accA[e] + accB[e]; + } + if constexpr (kWarps > 1) + { +#pragma unroll + for (int e = 0; e < 4; ++e) + { + reduceBuf[warp][e][lane] = acc[e]; + } + __syncthreads(); + if (warp != 0) + { + return false; + } + // Fixed order over warps keeps the result deterministic. +#pragma unroll + for (int w = 1; w < kWarps; ++w) + { +#pragma unroll + for (int e = 0; e < 4; ++e) + { + acc[e] += reduceBuf[w][e][lane]; + } + } + } + return true; +} + +__device__ __forceinline__ float roundToBf16(float v) +{ + return __bfloat162float(__float2bfloat16_rn(v)); +} + +template +__global__ void __launch_bounds__(T::kWarps* kWarpSize, T::kMinBlocks) gateUpSwigluKernel(Bf16 const* __restrict__ x, + Bf16 const* __restrict__ weight, Bf16* __restrict__ h, int numRows, int hidden, int inter) +{ + int const lane = threadIdx.x % kWarpSize; + int const group = lane / kQuadThreads; + int const tig = lane % kQuadThreads; + int const col = blockIdx.x * kGateUpCols + group; + + float acc[4]; + if (!streamTile(acc, x, weight, numRows, hidden, col, inter + col)) + { + return; + } +#pragma unroll + for (int b = 0; b < 2; ++b) + { + int const row = 2 * tig + b; + if (row < numRows) + { + // gate and up are rounded to BF16 first, matching a BF16 GEMM output feeding SwiGLU. + float const gate = roundToBf16(acc[b]); + float const up = roundToBf16(acc[2 + b]); + float const sig = 1.F / (1.F + __expf(-gate)); + h[static_cast(row) * inter + col] = __float2bfloat16_rn(gate * sig * up); + } + } +} + +template +__global__ void __launch_bounds__(T::kWarps* kWarpSize, T::kMinBlocks) downKernel(Bf16 const* __restrict__ h, + Bf16 const* __restrict__ weight, Bf16* __restrict__ out, int numRows, int inter, int hidden) +{ + int const lane = threadIdx.x % kWarpSize; + int const group = lane / kQuadThreads; + int const tig = lane % kQuadThreads; + int const n0 = blockIdx.x * kTileRows + group; + + float acc[4]; + if (!streamTile(acc, h, weight, numRows, inter, n0, n0 + kTileRows / 2)) + { + return; + } +#pragma unroll + for (int b = 0; b < 2; ++b) + { + int const row = 2 * tig + b; + if (row < numRows) + { + out[static_cast(row) * hidden + n0] = __float2bfloat16_rn(acc[b]); + out[static_cast(row) * hidden + n0 + kTileRows / 2] = __float2bfloat16_rn(acc[2 + b]); + } + } +} + +template +cudaError_t launchKernel(void (*kernel)(Bf16 const*, Bf16 const*, Bf16*, int, int, int), int gridSize, Bf16 const* in, + Bf16 const* weight, Bf16* out, int numRows, int kDim, int outDim, bool enablePdl, cudaStream_t stream) +{ + cudaLaunchConfig_t config{}; + config.gridDim = dim3(gridSize); + config.blockDim = dim3(T::kWarps * kWarpSize); + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = enablePdl ? 1 : 0; + config.attrs = attrs; + config.numAttrs = 1; + return cudaLaunchKernelEx(&config, kernel, in, weight, out, numRows, kDim, outDim); +} + +//! Variant table (warps, steps); the template arguments of each index are in launchVariant. +struct VariantDesc +{ + int warps; + int steps; +}; + +constexpr VariantDesc kVariants[] = { + {4, 8}, // 0: down default. L2 prefetch one chunk ahead + {8, 4}, // 1: gate/up default. Three CTAs/SM register budget, L2 prefetch one chunk ahead + {8, 4}, // 2: as 1 with ld.global.ca activation loads + {4, 8}, // 3: as 0 without L2 prefetch + {8, 8}, // 4: two CTAs/SM, no L2 prefetch +}; +constexpr int kNumVariants = sizeof(kVariants) / sizeof(kVariants[0]); + +bool isContractionSupported(int kDim, int variant) +{ + VariantDesc const& v = kVariants[variant]; + return kDim > 0 && kDim % (v.warps * kKPerStep * v.steps) == 0; +} + +//! Dispatches a variant index to its kernel instance; kGateUp selects the gate/up or the down kernel. +template +cudaError_t launchVariant(int variant, Bf16 const* in, Bf16 const* weight, Bf16* out, int numRows, int kDim, int outDim, + bool enablePdl, cudaStream_t stream) +{ + auto launch = [&](auto tiling) -> cudaError_t + { + using T = decltype(tiling); + if constexpr (kGateUp) + { + return launchKernel( + gateUpSwigluKernel, outDim / kGateUpCols, in, weight, out, numRows, kDim, outDim, enablePdl, stream); + } + else + { + return launchKernel( + downKernel, outDim / kTileRows, in, weight, out, numRows, kDim, outDim, enablePdl, stream); + } + }; + switch (variant) + { + case 0: return launch(Tiling<4, 8, 1, false, 1>{}); + case 1: return launch(Tiling<8, 4, 3, false, 1>{}); + case 2: return launch(Tiling<8, 4, 3, true, 1>{}); + case 3: return launch(Tiling<4, 8, 1, false, 0>{}); + case 4: return launch(Tiling<8, 8, 2, false, 0>{}); + default: return cudaErrorInvalidValue; + } +} + +} // namespace + +int getNumVariants() +{ + return kNumVariants; +} + +bool isGateUpSupported(int numRows, int hidden, int inter, int variant) +{ + return variant >= 0 && variant < kNumVariants && numRows >= 1 && numRows <= kMaxRows && inter > 0 + && inter % kGateUpCols == 0 && isContractionSupported(hidden, variant); +} + +bool isDownSupported(int numRows, int inter, int hidden, int variant) +{ + return variant >= 0 && variant < kNumVariants && numRows >= 1 && numRows <= kMaxRows && hidden > 0 + && hidden % kTileRows == 0 && isContractionSupported(inter, variant); +} + +cudaError_t invokeGateUpSwiglu(Bf16 const* x, Bf16 const* gateUpWeight, Bf16* h, int numRows, int hidden, int inter, + int variant, bool enablePdl, cudaStream_t stream) +{ + if (!isGateUpSupported(numRows, hidden, inter, variant)) + { + return cudaErrorInvalidValue; + } + return launchVariant(variant, x, gateUpWeight, h, numRows, hidden, inter, enablePdl, stream); +} + +cudaError_t invokeDown(Bf16 const* h, Bf16 const* downWeight, Bf16* out, int numRows, int inter, int hidden, + int variant, bool enablePdl, cudaStream_t stream) +{ + if (!isDownSupported(numRows, inter, hidden, variant)) + { + return cudaErrorInvalidValue; + } + return launchVariant(variant, h, downWeight, out, numRows, inter, hidden, enablePdl, stream); +} + +} // namespace kernels::glm_moe_dsa::dense_mlp + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/denseMlp.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/denseMlp.h new file mode 100644 index 000000000000..3d82801ffca0 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/denseMlp.h @@ -0,0 +1,62 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa::dense_mlp +{ + +//! Largest number of activation rows (tokens) the small-M dense MLP kernels accept. +constexpr int kMaxRows = 8; + +//! Number of tiling variants compiled for each kernel. Variant 0 is the default. +int getNumVariants(); + +//! Host-side eligibility checks (no device work). +//! \param numRows activation rows M. +//! \param hidden model hidden size (gate/up contraction size, down output size). +//! \param inter rank-local intermediate size (gate/up output size, down contraction size). +bool isGateUpSupported(int numRows, int hidden, int inter, int variant); +bool isDownSupported(int numRows, int inter, int hidden, int variant); + +//! Gate/up projection fused with SwiGLU. +//! For every row m < numRows and column j < inter: +//! gate = bf16(sum_k x[m, k] * W[j, k]), up = bf16(sum_k x[m, k] * W[inter + j, k]), +//! h[m, j] = bf16((gate * sigmoid(gate)) * up), all sums in FP32. +//! \param x activations [numRows, hidden], row-major, BF16. +//! \param gateUpWeight [2 * inter, hidden], row-major, BF16 (rows [0, inter) gate, [inter, 2 * inter) up). +//! \param h output [numRows, inter], row-major, BF16. +cudaError_t invokeGateUpSwiglu(__nv_bfloat16 const* x, __nv_bfloat16 const* gateUpWeight, __nv_bfloat16* h, int numRows, + int hidden, int inter, int variant, bool enablePdl, cudaStream_t stream); + +//! Down projection: out[m, n] = bf16(sum_k h[m, k] * W[n, k]) with FP32 sums. +//! \param h activations [numRows, inter], row-major, BF16. +//! \param downWeight [hidden, inter], row-major, BF16. +//! \param out output [numRows, hidden], row-major, BF16. +cudaError_t invokeDown(__nv_bfloat16 const* h, __nv_bfloat16 const* downWeight, __nv_bfloat16* out, int numRows, + int inter, int hidden, int variant, bool enablePdl, cudaStream_t stream); + +} // namespace kernels::glm_moe_dsa::dense_mlp + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/draftLiveRowHop.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/draftLiveRowHop.cu new file mode 100644 index 000000000000..33bfd92f7cef --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/draftLiveRowHop.cu @@ -0,0 +1,586 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/glmMoeDsa/draftLiveRowHop.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa::draft_live_rows +{ +namespace +{ +namespace cg = cooperative_groups; + +// Work decomposition (the post-attention router hop's, on the compact rows). +// +// Every cluster has kSlices CTAs; CTA c owns slice [c * kSlice, (c + 1) * kSlice) of the hidden dimension (96 threads +// x 8 elements), the slice a CTA of the native one-shot kernel owns, so the RMSNorm reduction reproduces the native +// tree. There are kG x kRowSplit clusters (kG = 4 expert groups, kRowSplit = min(R, 3) row groups, R = live rows). +// Cluster (g, h) collects and normalizes the compact rows r with r % kRowSplit == h and computes their logits for +// experts [64 g, 64 g + 64). Its router weights live in registers (loaded before the PDL wait). Each CTA multiplies +// its rows of slice c by the experts' weights of slice c with BF16 MMAs (experts on m16, rows on n8) and pushes the +// partial logits to the CTA that owns the expert; owners add the 8 slices in a fixed order, so the logits are +// identical on every rank and do not depend on R or on the row's position. Clusters never synchronize with each +// other; inside a cluster, slice sums and partial logits travel as st.async stores that complete transactions on the +// consumer's mbarrier. +constexpr int kSlices = 8; +constexpr int kSlice = kHiddenDim / kSlices; +constexpr int kEltsPerVec = 8; +constexpr int kVecPerSlice = kSlice / kEltsPerVec; +constexpr int kNativeWarpsPerSlice = kVecPerSlice / 32; +// Row stride of the normalized-slice tile in 16-byte vectors; 100 = 96 + 4 keeps the 16-byte B-fragment loads of two +// rows in one quarter-warp phase on disjoint banks. +constexpr int kNormedStride = kVecPerSlice + 4; +constexpr int kMmaChains = 4; +constexpr int kExpertGroups = 4; +constexpr uint32_t kLamportSentinel = 0x80000000U; + +static_assert(kHiddenDim % (kSlices * kEltsPerVec * 32) == 0); +static_assert(kVecPerSlice == kNativeWarpsPerSlice * 32); +static_assert(kMaxRows == 8, "the router MMA maps rows onto the n8 dimension"); + +struct HopConfig +{ + static constexpr int kThreads = 384; + static constexpr int kWarps = kThreads / 32; + static constexpr int kExpertsPerCluster = kNumExperts / kExpertGroups; + static constexpr int kExpertTiles = kExpertsPerCluster / 16; + static constexpr int kKGroups = kWarps / kExpertTiles; + static constexpr int kKPerGroup = kSlice / kKGroups; + static constexpr int kPairsPerGroup = kKPerGroup / 32; + static constexpr int kExpertsPerOwner = kExpertsPerCluster / kSlices; + static_assert(kExpertsPerCluster % 16 == 0 && kWarps % kExpertTiles == 0 && kSlice % (kKGroups * 32) == 0); + static_assert(kPairsPerGroup % (kMmaChains / 2) == 0); +}; + +constexpr int kMaxRowsPerCtaAll = kMaxRows / 2; + +struct alignas(16) HopStorage +{ + uint4 gamma[kVecPerSlice]; + float vecSq[kMaxRowsPerCtaAll][kVecPerSlice]; + float warpSq[kMaxRowsPerCtaAll][kNativeWarpsPerSlice]; + float blockSq[kMaxRowsPerCtaAll][kSlices]; + // Completes when all slice sums of this CTA's rows have arrived (st.async complete_tx). + alignas(8) uint64_t normBar; + // Completes when all partial logits of the experts this CTA owns have arrived. + alignas(8) uint64_t routerBar; +}; + +struct alignas(16) RouterStorage +{ + uint4 normed[kMaxRows][kNormedStride]; // local row on the MMA n dimension; unused rows stay zero + float partial[kSlices][HopConfig::kKGroups][HopConfig::kExpertsPerOwner][kMaxRows]; +}; + +union BfVec +{ + uint4 packed; + __nv_bfloat16 elts[kEltsPerVec]; +}; + +__device__ __forceinline__ uint4 loadVolatile(void const* ptr) +{ + uint4 v; + asm volatile("ld.volatile.global.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(ptr) + : "memory"); + return v; +} + +__device__ __forceinline__ uint4 loadConstant(void const* ptr) +{ + uint4 v; + asm volatile("ld.global.nc.v4.u32 {%0, %1, %2, %3}, [%4];" : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) : "l"(ptr)); + return v; +} + +__device__ __forceinline__ bool isDirty(uint4 const& v) +{ + return v.x == kLamportSentinel || v.y == kLamportSentinel || v.z == kLamportSentinel || v.w == kLamportSentinel; +} + +// Lamport payloads must never contain the sentinel word: map bf16 -0.0 to +0.0 (the reduction starts from +0.0, so +// the sum is unchanged). +__device__ __forceinline__ uint4 sanitize(uint4 v) +{ + BfVec b; + b.packed = v; +#pragma unroll + for (int i = 0; i < kEltsPerVec; ++i) + { + if (__bfloat16_as_ushort(b.elts[i]) == 0x8000U) + { + b.elts[i] = __ushort_as_bfloat16(0); + } + } + return b.packed; +} + +__device__ __forceinline__ uint32_t sharedAddr(void const* ptr) +{ + return static_cast(__cvta_generic_to_shared(ptr)); +} + +// Address of the same shared-memory object in CTA `rank` of the cluster. +__device__ __forceinline__ uint32_t clusterAddr(uint32_t addr, uint32_t rank) +{ + uint32_t out; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(out) : "r"(addr), "r"(rank)); + return out; +} + +__device__ __forceinline__ void mbarrierInitExpectTx(uint64_t* bar, uint32_t bytes) +{ + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(sharedAddr(bar)) : "memory"); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(sharedAddr(bar)), "r"(bytes) + : "memory"); +} + +__device__ __forceinline__ void mbarrierWaitPhase0(uint64_t* bar) +{ + asm volatile( + "{\n" + ".reg .pred done;\n" + "LIVE_ROW_MBAR_WAIT_%=:\n" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 done, [%0], 0;\n" + "@!done bra LIVE_ROW_MBAR_WAIT_%=;\n" + "}\n" ::"r"(sharedAddr(bar)) + : "memory"); +} + +// Remote shared-memory stores that complete transactions on the destination CTA's mbarrier. +__device__ __forceinline__ void storeAsync(uint32_t remoteAddr, float v, uint32_t remoteBar) +{ + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.f32 [%0], %1, [%2];" ::"r"(remoteAddr), "f"(v), + "r"(remoteBar) + : "memory"); +} + +__device__ __forceinline__ void storeAsync(uint32_t remoteAddr, float a, float b, uint32_t remoteBar) +{ + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.v2.f32 [%0], {%1, %2}, [%3];" ::"r"(remoteAddr), + "f"(a), "f"(b), "r"(remoteBar) + : "memory"); +} + +__device__ __forceinline__ void mmaBf16( + float (&d)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t b0, uint32_t b1) +{ + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, " + "{%0, %1, %2, %3};" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)); +} + +// R = live (compact) rows; kRowSplit = row groups. +template +__global__ void __launch_bounds__(HopConfig::kThreads, 1) liveRowHopKernel(LiveRowHopParams const params) +{ + static_assert(R >= 1 && R <= kMaxRows); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + using Cfg = HopConfig; + constexpr int kClusters = kExpertGroups * kRowSplit; + constexpr int kGridCtas = kClusters * kSlices; + // Rows one CTA collects, normalizes and multiplies. + constexpr int kRowsPerCta = (R + kRowSplit - 1) / kRowSplit; + static_assert(kRowsPerCta <= kMaxRowsPerCtaAll); + static_assert(kRowsPerCta * kVecPerSlice <= Cfg::kThreads && kRowsPerCta * kNativeWarpsPerSlice <= Cfg::kWarps); + static_assert(Cfg::kExpertsPerOwner * kRowsPerCta * kSlices <= Cfg::kThreads); + + __shared__ HopStorage hop; + __shared__ RouterStorage rs; + + cg::cluster_group cluster = cg::this_cluster(); + int const slice = static_cast(cluster.block_rank()); + int const clusterId = blockIdx.y; + int const expertGroup = clusterId % kExpertGroups; + int const rowGroup = clusterId / kExpertGroups; + int const tid = threadIdx.x; + int const lane = tid & 31; + int const warp = tid >> 5; + int const sliceBase = slice * kSlice; + + // Compact rows this CTA collects and normalizes: local row j is compact row rowGroup + j * kRowSplit. + int const numRowsCta = (R - rowGroup + kRowSplit - 1) / kRowSplit; + // This thread's item: one 16-byte vector of one row slice. + int const itemRow = tid / kVecPerSlice; + int const itemVec = tid % kVecPerSlice; + bool const hasItem = itemRow < numRowsCta; + int const row = rowGroup + itemRow * kRowSplit; + // The rank's own partial of a row slice is published by exactly one CTA of the grid. + bool const sends = hasItem && itemRow % kExpertGroups == expertGroup; + size_t const colOff = static_cast(sliceBase) + itemVec * kEltsPerVec; + size_t const rowOff = static_cast(row) * kHiddenDim + colOff; + // The cluster that writes the normalized rows and the residual. + bool const writesOutputs = expertGroup == 0; + + // ---- Prologue: constant model weights only; may overlap the producer of `input` under PDL. ---- + if (tid < kVecPerSlice) + { + hop.gamma[tid] = loadConstant(reinterpret_cast(params.gamma) + slice * kVecPerSlice + tid); + } + int const gid = lane >> 2; + int const tig = lane & 3; + uint4 wLo[Cfg::kPairsPerGroup]; + uint4 wHi[Cfg::kPairsPerGroup]; + { + int const tile = warp % Cfg::kExpertTiles; + int const kGroup = warp / Cfg::kExpertTiles; + auto const* w = reinterpret_cast<__nv_bfloat16 const*>(params.routerWeight); + int const expert = expertGroup * Cfg::kExpertsPerCluster + tile * 16 + gid; +#pragma unroll + for (int p = 0; p < Cfg::kPairsPerGroup; ++p) + { + int const k = sliceBase + kGroup * Cfg::kKPerGroup + p * 32 + tig * kEltsPerVec; + wLo[p] = loadConstant(w + static_cast(expert) * kHiddenDim + k); + wHi[p] = loadConstant(w + static_cast(expert + 8) * kHiddenDim + k); + } + // MMA columns of rows this CTA does not own stay zero. + for (int i = tid; i < kMaxRows * kNormedStride; i += Cfg::kThreads) + { + rs.normed[i / kNormedStride][i % kNormedStride] = make_uint4(0, 0, 0, 0); + } + } + + // Cluster-local completion barriers: every CTA of a cluster has the same rows, so each expects the 8 slice sums + // of its rows and the 8 x kKGroups partial tiles of the experts it owns. A split cluster barrier publishes the + // initialization; its wait precedes the first remote store, long after all CTAs arrived. + if (tid == 0) + { + mbarrierInitExpectTx(&hop.normBar, static_cast(numRowsCta * kSlices * sizeof(float))); + mbarrierInitExpectTx(&hop.routerBar, numRowsCta > 0 ? static_cast(sizeof(rs.partial)) : 0U); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + + // ---- Wait for the producer of `input` (and transitively of `residualIn`, `rowIds` and the protocol state). ---- + asm volatile("griddepcontrol.wait;" ::: "memory"); + // Dependents only start their own prologue; they still wait for this grid to complete. + asm volatile("griddepcontrol.launch_dependents;"); + + // The live row's source row (written by an earlier kernel of the stream: a plain load after the wait). + int64_t srcRow = 0; + if (hasItem) + { + srcRow = params.rowIds[row]; + } + uint4 const flags = loadVolatile(params.flags); + srcRow = srcRow < 0 ? 0 : (srcRow >= params.numInputRows ? params.numInputRows - 1 : srcRow); + size_t const srcOff = static_cast(srcRow) * kHiddenDim + colOff; + auto const* input = reinterpret_cast<__nv_bfloat16 const*>(params.input); + auto const* residualIn = reinterpret_cast<__nv_bfloat16 const*>(params.residualIn); + uint4 res = make_uint4(0, 0, 0, 0); + // This rank's own partial, as published (sanitized); it is used from registers, not from the multicast copy. + uint4 own = make_uint4(0, 0, 0, 0); + if (hasItem) + { + res = *reinterpret_cast(residualIn + srcOff); + own = sanitize(*reinterpret_cast(input + srcOff)); + } + uint32_t const cur = flags.x; + uint32_t const clearBytes = flags.y; + char* ucBase = reinterpret_cast(params.ucBuffer); + auto* curUc = reinterpret_cast<__nv_bfloat16*>(ucBase + cur * kLamportBufferBytes); + auto* curMc + = reinterpret_cast<__nv_bfloat16*>(reinterpret_cast(params.mcBuffer) + cur * kLamportBufferBytes); + // Lamport layout of one buffer: [compact row][source rank][kHiddenDim]. + size_t const slotBase = static_cast(row) * kNumRanks * kHiddenDim + colOff; + + // ---- Publish this rank's partial to every rank (one multicast store per vector). ---- + if (sends) + { + *reinterpret_cast(curMc + slotBase + static_cast(params.rank) * kHiddenDim) = own; + } + + // ---- Collect the other ranks' partials. ---- + uint4 vals[kNumRanks]; + if (hasItem) + { +#pragma unroll + for (int s = 0; s < kNumRanks; ++s) + { + vals[s] = s == params.rank ? own : loadVolatile(curUc + slotBase + static_cast(s) * kHiddenDim); + } + } + + // ---- Reset the previous call's buffer to the sentinel (reused two calls from now), off the poll's path. ---- + { + uint32_t const dirty = (cur + kNumLamportBuffers - 1) % kNumLamportBuffers; + auto* dirtyBuf = reinterpret_cast(ucBase + dirty * kLamportBufferBytes); + uint32_t const numVec = clearBytes / sizeof(uint4); + uint4 const sentinel = make_uint4(kLamportSentinel, kLamportSentinel, kLamportSentinel, kLamportSentinel); + uint32_t const ctaLinear = clusterId * kSlices + slice; + for (uint32_t j = ctaLinear * Cfg::kThreads + tid; j < numVec; j += kGridCtas * Cfg::kThreads) + { + dirtyBuf[j] = sentinel; + } + } + + if (hasItem) + { + bool pending = true; + while (pending) + { + pending = false; +#pragma unroll + for (int s = 0; s < kNumRanks; ++s) + { + if (isDirty(vals[s])) + { + pending = true; + vals[s] = loadVolatile(curUc + slotBase + static_cast(s) * kHiddenDim); + } + } + } + } + + // ---- Reduce in rank order, add the residual, and form per-vector sums of squares (native arithmetic). ---- + uint4 resOut = make_uint4(0, 0, 0, 0); + if (hasItem) + { + float acc[kEltsPerVec]; +#pragma unroll + for (int e = 0; e < kEltsPerVec; ++e) + { + acc[e] = 0.F; + } +#pragma unroll + for (int s = 0; s < kNumRanks; ++s) + { + BfVec x; + x.packed = vals[s]; +#pragma unroll + for (int e = 0; e < kEltsPerVec; ++e) + { + acc[e] += __bfloat162float(x.elts[e]); + } + } + BfVec rv; + rv.packed = res; + BfVec y; + float sq = 0.F; +#pragma unroll + for (int e = 0; e < kEltsPerVec; ++e) + { + y.elts[e] = __hadd(__float2bfloat16(acc[e]), rv.elts[e]); + sq += __bfloat162float(__hmul(y.elts[e], y.elts[e])); + } + resOut = y.packed; + hop.vecSq[itemRow][itemVec] = sq; + } + __syncthreads(); + if (tid == 0) + { + // Every thread of this CTA has read the protocol state; CTA 0 advances it only after all CTAs report. + asm volatile("red.release.gpu.global.add.u32 [%0], %1;" ::"l"(params.flags + 2), "r"(1U) : "memory"); + } + + // ---- RMS: warp butterfly over 32 vectors, then (w0 + w2) + w1 per slice, then the 8 slices in order. ---- + if (warp < numRowsCta * kNativeWarpsPerSlice) + { + int const j = warp / kNativeWarpsPerSlice; + int const q = warp % kNativeWarpsPerSlice; + float val = hop.vecSq[j][q * 32 + lane]; +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) + { + val = val + __shfl_xor_sync(0xffffffffU, val, mask, 32); + } + if (lane == 0) + { + hop.warpSq[j][q] = val; + } + } + __syncthreads(); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + if (tid < numRowsCta * kSlices) + { + int const j = tid / kSlices; + int const dst = tid % kSlices; + float const blockSum = (hop.warpSq[j][0] + hop.warpSq[j][2]) + hop.warpSq[j][1]; + storeAsync( + clusterAddr(sharedAddr(&hop.blockSq[j][slice]), dst), blockSum, clusterAddr(sharedAddr(&hop.normBar), dst)); + } + mbarrierWaitPhase0(&hop.normBar); + + if (hasItem) + { + float fullSum = 0.F; +#pragma unroll + for (int c = 0; c < kSlices; ++c) + { + fullSum += hop.blockSq[itemRow][c]; + } + float const rcpRms = rsqrtf(fullSum / static_cast(kHiddenDim) + params.eps); + BfVec y; + y.packed = resOut; + BfVec g; + g.packed = hop.gamma[itemVec]; + BfVec n; +#pragma unroll + for (int e = 0; e < kEltsPerVec; ++e) + { + n.elts[e] = __float2bfloat16(__bfloat162float(y.elts[e]) * rcpRms * __bfloat162float(g.elts[e])); + } + if (writesOutputs) + { + *reinterpret_cast(reinterpret_cast<__nv_bfloat16*>(params.normOut) + rowOff) = n.packed; + *reinterpret_cast(reinterpret_cast<__nv_bfloat16*>(params.residualOut) + rowOff) = y.packed; + } + rs.normed[itemRow][itemVec] = n.packed; + } + + if (numRowsCta > 0) + { + __syncthreads(); + int const tile = warp % Cfg::kExpertTiles; + int const kGroup = warp / Cfg::kExpertTiles; + // Independent accumulator chains keep the MMA latency off a single dependency chain. + float acc[kMmaChains][4] = {}; +#pragma unroll + for (int p = 0; p < Cfg::kPairsPerGroup; ++p) + { + // One 16-byte load covers two k16 steps; the k-permutation inside a step is the same for both operands. + uint4 const b = rs.normed[gid][kGroup * (Cfg::kKPerGroup / kEltsPerVec) + p * 4 + tig]; + int const chain = (2 * p) % kMmaChains; + mmaBf16(acc[chain], wLo[p].x, wHi[p].x, wLo[p].y, wHi[p].y, b.x, b.y); + mmaBf16(acc[chain + 1], wLo[p].z, wHi[p].z, wLo[p].w, wHi[p].w, b.z, b.w); + } + float sum4[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + sum4[i] = (acc[0][i] + acc[1][i]) + (acc[2][i] + acc[3][i]); + } + // Expert e (cluster-local) is owned by CTA e / kExpertsPerOwner. + int const e0 = tile * 16 + gid; + int const e1 = e0 + 8; + uint32_t const owner0 = e0 / Cfg::kExpertsPerOwner; + uint32_t const owner1 = e1 / Cfg::kExpertsPerOwner; + storeAsync(clusterAddr(sharedAddr(&rs.partial[slice][kGroup][e0 % Cfg::kExpertsPerOwner][tig * 2]), owner0), + sum4[0], sum4[1], clusterAddr(sharedAddr(&hop.routerBar), owner0)); + storeAsync(clusterAddr(sharedAddr(&rs.partial[slice][kGroup][e1 % Cfg::kExpertsPerOwner][tig * 2]), owner1), + sum4[2], sum4[3], clusterAddr(sharedAddr(&hop.routerBar), owner1)); + // Also keeps this CTA resident until every partial addressed to it has landed. + mbarrierWaitPhase0(&hop.routerBar); + + // This CTA owns experts [slice * kExpertsPerOwner, (slice + 1) * kExpertsPerOwner) of the cluster. Lane groups + // of 8 sum over the 8 slices. + if (tid < Cfg::kExpertsPerOwner * numRowsCta * kSlices) + { + int const c = tid % kSlices; + int const ij = tid / kSlices; + int const ei = ij / numRowsCta; + int const j = ij % numRowsCta; + float sum = 0.F; +#pragma unroll + for (int kg = 0; kg < Cfg::kKGroups; ++kg) + { + sum += rs.partial[c][kg][ei][j]; + } +#pragma unroll + for (int mask = kSlices / 2; mask > 0; mask >>= 1) + { + sum = sum + __shfl_xor_sync(0xffffffffU, sum, mask, kSlices); + } + if (c == 0) + { + int const r = rowGroup + j * kRowSplit; + params.routerLogits[r * kNumExperts + expertGroup * Cfg::kExpertsPerCluster + + slice * Cfg::kExpertsPerOwner + ei] + = sum; + } + } + } + + // ---- CTA 0 advances the protocol state once every CTA has read it. ---- + if (tid == 0 && slice == 0 && clusterId == 0) + { + uint32_t arrived = 0; + do + { + asm volatile("ld.acquire.gpu.global.u32 %0, [%1];" : "=r"(arrived) : "l"(params.flags + 2) : "memory"); + } while (arrived < static_cast(kGridCtas)); + // The next call clears exactly the slots this call wrote: compact rows 0..R-1 of every rank. + uint4 const next = make_uint4((cur + 1) % kNumLamportBuffers, + static_cast(R * kNumRanks * kHiddenDim * sizeof(__nv_bfloat16)), 0U, 0U); + asm volatile("st.relaxed.gpu.global.v4.u32 [%0], {%1, %2, %3, %4};" ::"l"(params.flags), "r"(next.x), + "r"(next.y), "r"(next.z), "r"(next.w) + : "memory"); + } +#else + // Clusters, distributed shared memory and PDL need SM90+; the launcher is only reached on supported devices. + __trap(); +#endif +} + +template +cudaError_t launchTyped(LiveRowHopParams const& params) +{ + cudaLaunchAttribute attrs[2]; + attrs[0].id = cudaLaunchAttributeClusterDimension; + attrs[0].val.clusterDim.x = kSlices; + attrs[0].val.clusterDim.y = 1; + attrs[0].val.clusterDim.z = 1; + attrs[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[1].val.programmaticStreamSerializationAllowed = params.enablePdl ? 1 : 0; + + cudaLaunchConfig_t config{}; + config.gridDim = dim3(kSlices, kExpertGroups * kRowSplit, 1); + config.blockDim = dim3(HopConfig::kThreads, 1, 1); + config.dynamicSmemBytes = 0; + config.stream = params.stream; + config.attrs = attrs; + config.numAttrs = 2; + return cudaLaunchKernelEx(&config, liveRowHopKernel, params); +} + +} // namespace + +cudaError_t invokeDraftLiveRowHop(LiveRowHopParams const& params) +{ + if (params.rank < 0 || params.rank >= kNumRanks || params.flags == nullptr || params.ucBuffer == nullptr + || params.mcBuffer == nullptr || params.rowIds == nullptr || params.routerWeight == nullptr + || params.routerLogits == nullptr || params.numInputRows < 1 || params.numInputRows > kMaxRows + || params.numRows < 1 || params.numRows > params.numInputRows) + { + return cudaErrorInvalidValue; + } + // 4 expert groups x min(R, 3) row groups (at most 12 clusters of 8 CTAs, each CTA <= 3 row slices). + switch (params.numRows) + { + case 1: return launchTyped<1, 1>(params); + case 2: return launchTyped<2, 2>(params); + case 3: return launchTyped<3, 3>(params); + case 4: return launchTyped<4, 3>(params); + case 5: return launchTyped<5, 3>(params); + case 6: return launchTyped<6, 3>(params); + case 7: return launchTyped<7, 3>(params); + case 8: return launchTyped<8, 3>(params); + default: return cudaErrorInvalidValue; + } +} + +} // namespace kernels::glm_moe_dsa::draft_live_rows + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/draftLiveRowHop.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/draftLiveRowHop.h new file mode 100644 index 000000000000..a6c6f40b40f1 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/draftLiveRowHop.h @@ -0,0 +1,88 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa::draft_live_rows +{ + +//! Live-row post-attention hop (MTP draft step 0). +//! +//! A fork of the router hop of postAttentionHop.h that serves only the ``numRows`` live rows of a call with +//! ``numInputRows`` rows. Compact row ``r`` (0 <= r < numRows) is source row ``rowIds[r]`` of ``input`` and +//! ``residualIn``; for it the kernel computes, on every rank and with identical bits on every rank: +//! sum = bf16(sum over ranks 0..7 of float(partial_rank)) (one-shot Lamport exchange, fixed order) +//! residual = sum + residualIn (bf16 add) +//! normed = bf16(float(residual) * rsqrt(mean(bf16(residual^2)) + eps) * float(gamma)) +//! logits = normed @ routerWeight^T in FP32 ([numRows, kNumExperts]) +//! Every quantity of a row is computed by the same per-row arithmetic as the post-attention hop, so compact row ``r`` +//! equals row ``rowIds[r]`` of a post-attention hop over all ``numInputRows`` rows. +//! +//! ``rowIds`` is read on the device after the kernel's programmatic dependency wait, never before it; it may be +//! written by any earlier kernel of the stream. Ids outside ``[0, numInputRows)`` are clamped into that range. +//! +//! The exchange runs on the post-attention hop's multicast workspace and protocol state, with the same layout +//! (``[row][source rank][kHiddenDim]`` bf16 per Lamport buffer, ``kNumLamportBuffers`` rotating buffers, flags +//! ``{current buffer, bytes to clear in the previous buffer, arrival counter, unused}``), so calls of both kernels +//! may be issued in any order on the stream that issues the group's hops. +struct LiveRowHopParams +{ + void const* input = nullptr; //!< [numInputRows, kHiddenDim] bf16, rank-local partial (o_proj output) + void const* residualIn = nullptr; //!< [numInputRows, kHiddenDim] bf16 + int64_t const* rowIds = nullptr; //!< [numRows] device ids of the live rows in [0, numInputRows) + void const* gamma = nullptr; //!< [kHiddenDim] bf16 RMSNorm weight + void const* routerWeight = nullptr; //!< [kNumExperts, kHiddenDim] bf16 + void* normOut = nullptr; //!< [numRows, kHiddenDim] bf16 + void* residualOut = nullptr; //!< [numRows, kHiddenDim] bf16 + float* routerLogits = nullptr; //!< [numRows, kNumExperts] fp32 + void* ucBuffer = nullptr; //!< local unicast view of the multicast workspace + void* mcBuffer = nullptr; //!< multicast address of the same workspace + uint32_t* flags = nullptr; //!< [4] device-side protocol state of this workspace + int numRows = 0; //!< live rows, 1 <= numRows <= numInputRows + int numInputRows = 0; //!< rows of ``input`` / ``residualIn``, <= kMaxRows + int rank = 0; + float eps = 1e-5F; + bool enablePdl = true; + cudaStream_t stream = nullptr; +}; + +// The post-attention hop's workspace geometry (the kernel runs on that workspace). +constexpr int kHiddenDim = 6144; +constexpr int kNumRanks = 8; +constexpr int kNumExperts = 256; +constexpr int kMaxRows = 8; +constexpr int kNumLamportBuffers = 3; + +//! Bytes of one Lamport buffer: [kMaxRows][kNumRanks][kHiddenDim] bf16. +constexpr size_t kLamportBufferBytes = size_t(kMaxRows) * kNumRanks * kHiddenDim * 2; +//! Bytes of the whole multicast workspace. +constexpr size_t kWorkspaceBytes = kLamportBufferBytes * kNumLamportBuffers; + +//! Launch the live-row hop on ``params.stream``. Returns the launch status. +cudaError_t invokeDraftLiveRowHop(LiveRowHopParams const& params); + +} // namespace kernels::glm_moe_dsa::draft_live_rows + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/draftRoutedMoe.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/draftRoutedMoe.cu new file mode 100644 index 000000000000..a4ca68a4859d --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/draftRoutedMoe.cu @@ -0,0 +1,2628 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/glmMoeDsa/draftRoutedMoe.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa::draft_routed_moe +{ + +namespace +{ + +// Ring depths (2 KB chunks per warp; the deeper pair with kFlagDeepRing) and the FC2 intermediate ring (one entry +// per expert of a tile). +constexpr int kFc1Stages = 12; +constexpr int kFc1StagesDeep = 20; +constexpr int kFc2Stages = 8; +constexpr int kFc2StagesDeep = 12; +constexpr int kFc2HStages = 6; +constexpr int kChunkBytes = 16 * kKChunk * 2; // 16 rows x 128 bytes +constexpr int kXRowBytes = kHidden * 2; // 12288 +constexpr int kHEntryBytes = kMaxRows * kInter * 2; // 4096 +constexpr int kStageBytes = kMaxRows * kTopK * 16 * 2; +constexpr float kSumEpsilon = 1e-20f; +constexpr float kBf16MaxFinite = 3.38953139e38f; +// Bytes of one FC2 weight block of an expert (one K chunk of all 6144 rows) and of one L2 prefetch piece. +constexpr int kFc2BlockBytes = kHidden * kKChunk * 2; +constexpr int kPrefetchPieceBytes = 16384; +constexpr int kPrefetchPiecesPerBlock = kFc2BlockBytes / kPrefetchPieceBytes; +static_assert(kFc2BlockBytes % kPrefetchPieceBytes == 0, "prefetch pieces must tile an FC2 block"); +// FC2 intermediate ring: expert s (in the warp's order over its items) has its K chunks at stream positions +// 4s..4s+3 (alone or as the first of a pair) or 4s-3, 4s-1, 4s+1, 4s+3 (second of a pair), so it is last read at +// iteration 4s+7 or earlier, and expert s + kFc2HStages loads into the same entry together with its first chunk, at +// position >= 4(s + kFc2HStages) - 3, issued at the end of iteration (position - kStages + 1). The entry is free when +// that iteration is not earlier than the last read: kStages <= 4 * kFc2HStages - 9. +static_assert(kFc2Stages <= 4 * kFc2HStages - 9 && kFc2StagesDeep <= 4 * kFc2HStages - 9, + "an FC2 intermediate entry is refilled before its last read"); + +// --------------------------------------------------------------------------------------------------------------- +// Arithmetic with the exact instructions of the native kernels. +// --------------------------------------------------------------------------------------------------------------- + +__device__ __forceinline__ float mulFtz(float a, float b) +{ + float r; + asm("mul.ftz.f32 %0, %1, %2;" : "=f"(r) : "f"(a), "f"(b)); + return r; +} + +__device__ __forceinline__ float addFtz(float a, float b) +{ + float r; + asm("add.ftz.f32 %0, %1, %2;" : "=f"(r) : "f"(a), "f"(b)); + return r; +} + +__device__ __forceinline__ float fmaFtz(float a, float b, float c) +{ + float r; + asm("fma.rn.ftz.f32 %0, %1, %2, %3;" : "=f"(r) : "f"(a), "f"(b), "f"(c)); + return r; +} + +__device__ __forceinline__ float tanhApprox(float x) +{ + float r; + asm("tanh.approx.f32 %0, %1;" : "=f"(r) : "f"(x)); + return r; +} + +__device__ __forceinline__ float divApproxFtz(float a, float b) +{ + float r; + asm("div.approx.ftz.f32 %0, %1, %2;" : "=f"(r) : "f"(a), "f"(b)); + return r; +} + +__device__ __forceinline__ float ex2ApproxFtz(float x) +{ + float r; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(r) : "f"(x)); + return r; +} + +__device__ __forceinline__ float rcpApproxFtz(float x) +{ + float r; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(r) : "f"(x)); + return r; +} + +__device__ __forceinline__ float mulRn(float a, float b) +{ + float r; + asm("mul.rn.f32 %0, %1, %2;" : "=f"(r) : "f"(a), "f"(b)); + return r; +} + +__device__ __forceinline__ float addRn(float a, float b) +{ + float r; + asm("add.rn.f32 %0, %1, %2;" : "=f"(r) : "f"(a), "f"(b)); + return r; +} + +//! FlashInfer ``sigmoid_accurate`` as compiled with ``-use_fast_math``. +__device__ __forceinline__ float routingSigmoid(float x) +{ + return fmaFtz(tanhApprox(mulFtz(x, 0.5f)), 0.5f, 0.5f); +} + +//! Round to BF16 like ``F2FP.SATFINITE.BF16.F32``: round to nearest even, overflow saturates, NaN stays NaN. +__device__ __forceinline__ __nv_bfloat16 bf16SatFinite(float x) +{ + __nv_bfloat16 r = __float2bfloat16_rn(x); + float back = __bfloat162float(r); + if (back == INFINITY || back == -INFINITY) + { + r = __float2bfloat16_rn(x > 0.f ? kBf16MaxFinite : -kBf16MaxFinite); + } + return r; +} + +//! The native gated epilogue (``swiGlu``, no clamp limit, alpha 1, beta 0): ``(up + 0) * (gate * sigmoid(gate))`` +//! with ``sigmoid(g) = rcp.approx(1 + ex2.approx(-(g * log2 e)))``; both inputs are clamped to +-FLT_MAX first. +__device__ __forceinline__ float swiGlu(float up, float gate) +{ + float const g = gate > FLT_MAX ? FLT_MAX : (gate < -FLT_MAX ? -FLT_MAX : gate); + float const u = up > FLT_MAX ? FLT_MAX : (up < -FLT_MAX ? -FLT_MAX : up); + float const t = mulRn(mulRn(g, 1.4426950216293334961f), -1.f); + float const s = rcpApproxFtz(addRn(ex2ApproxFtz(t), 1.f)); + return mulRn(addRn(u, 0.f), mulRn(g, s)); +} + +// --------------------------------------------------------------------------------------------------------------- +// Data movement and MMA. +// --------------------------------------------------------------------------------------------------------------- + +__device__ __forceinline__ uint32_t smemAddr(void const* p) +{ + return static_cast(__cvta_generic_to_shared(p)); +} + +__device__ __forceinline__ void cpAsync16(uint32_t dst, void const* src) +{ + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;" ::"r"(dst), "l"(__cvta_generic_to_global(src)) : "memory"); +} + +__device__ __forceinline__ void cpAsyncCommit() +{ + asm volatile("cp.async.commit_group;" ::: "memory"); +} + +template +__device__ __forceinline__ void cpAsyncWait() +{ + asm volatile("cp.async.wait_group %0;" ::"n"(N) : "memory"); +} + +__device__ __forceinline__ void ldmatrixX4(uint32_t (&r)[4], uint32_t addr) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(addr)); +} + +__device__ __forceinline__ void mmaBf16(float (&c)[4], uint32_t const (&a)[4], uint32_t b0, uint32_t b1) +{ + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, " + "{%0, %1, %2, %3};" + : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1)); +} + +__device__ __forceinline__ void gridDependencyWait() +{ + asm volatile("griddepcontrol.wait;" ::: "memory"); +} + +__device__ __forceinline__ void gridDependencyTrigger() +{ + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); +} + +//! Byte offset of 16-byte unit ``u`` of a 128-byte row segment for row ``r`` (XOR swizzle, conflict-free ldmatrix). +__device__ __forceinline__ uint32_t swizzled(int r, int u) +{ + return static_cast(((u & ~7) | ((u & 7) ^ (r & 7))) << 4); +} + +//! Copies one 16-row x 64-element weight chunk (2 KB, contiguous in the BlockMajorK layout) into a ring slot. +__device__ __forceinline__ void issueWeightChunk(uint32_t slot, __nv_bfloat16 const* src, int lane) +{ +#pragma unroll + for (int j = 0; j < 4; ++j) + { + int const idx = lane + 32 * j; + int const r = idx >> 3; + int const u = idx & 7; + cpAsync16(slot + r * 128 + swizzled(r, u), src + r * kKChunk + u * 8); + } +} + +//! A fragments of k-step ``q`` (16 elements) of a chunk: rows 0..15, natural K order (m16n8k16 row-major A). +__device__ __forceinline__ void loadA(uint32_t (&a)[4], uint32_t slot, int q, int lane) +{ + int const m = lane >> 3; + int const r = (lane & 7) | ((m & 1) << 3); + int const u = 2 * q + (m >> 1); + ldmatrixX4(a, slot + r * 128 + swizzled(r, u)); +} + +//! B fragments of two consecutive k-steps from a [token][k] BF16 tile (rows ``rowBytes`` apart, units swizzled +//! by token); tokens at or above ``rows`` read token 0 (their columns are discarded). +__device__ __forceinline__ void loadB2(uint32_t (&b)[4], uint32_t base, int rowBytes, int unit, int rows, int lane) +{ + int tok = lane & 7; + tok = tok < rows ? tok : 0; + int const u = unit + (lane >> 3); + ldmatrixX4(b, base + tok * rowBytes + swizzled(tok, u)); +} + +// --------------------------------------------------------------------------------------------------------------- +// Routing (FlashInfer routingCustom: SigmoidBias + ScaledSumNormalize, warp per token). +// --------------------------------------------------------------------------------------------------------------- + +//! Selection key: the score's bits mapped to an unsigned order (a NaN score ranks above every number), then the +//! expert id so that ties go to the lower index. +__device__ __forceinline__ uint64_t packKey(float key, int expert) +{ + uint32_t bits = __float_as_uint(key); + bits = (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u); + bits = key != key ? 0xFFFFFFFFu : bits; + return (static_cast(bits) << 32) | static_cast(0xFFFF & (65535 - expert)); +} + +//! Lexicographic warp maximum of 64-bit keys with two ``redux.sync.max.u32`` (as the native top-k reduction). +__device__ __forceinline__ uint64_t warpMax64(uint64_t v) +{ + uint32_t const hi = static_cast(v >> 32); + uint32_t const lo = static_cast(v); + uint32_t maxHi; + asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;" : "=r"(maxHi) : "r"(hi)); + uint32_t const loContrib = hi == maxHi ? lo : 0u; + uint32_t maxLo; + asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;" : "=r"(maxLo) : "r"(loContrib)); + return (static_cast(maxHi) << 32) | maxLo; +} + +struct RouteSmem +{ + int topExpert[kMaxRows][kTopK]; + float topWeight[kMaxRows][kTopK]; //!< Value of the BF16 weight. + int slotDistinct[kMaxRows][kTopK]; + int distinct[kMaxDistinct]; + uint32_t mask[kNumExperts / 32]; //!< Experts selected by any row: bit e % 32 of word e / 32. + int numDistinct; +}; + +//! A lane's routing inputs of one token: logits of experts 8l .. 8l + 7 (two 16-byte loads). +struct RouteLogits +{ + float4 l0; + float4 l1; +}; + +__device__ __forceinline__ RouteLogits loadRouteLogits(DraftRoutedMoeParams const& p, int token, int lane) +{ + float4 const* logits = reinterpret_cast(p.logits + token * kNumExperts + lane * (kNumExperts / 32)); + return RouteLogits{logits[0], logits[1]}; +} + +//! Compare-exchange for a descending order: a[i] >= a[j] afterwards. +__device__ __forceinline__ void keyExchange(uint64_t& hi, uint64_t& lo) +{ + uint64_t const a = hi; + uint64_t const b = lo; + hi = a > b ? a : b; + lo = a > b ? b : a; +} + +//! Sorts 8 keys in descending order (Batcher's odd-even merge sort, 19 compare-exchanges). +__device__ __forceinline__ void sortKeysDescending(uint64_t (&k)[8]) +{ + keyExchange(k[0], k[1]); + keyExchange(k[2], k[3]); + keyExchange(k[4], k[5]); + keyExchange(k[6], k[7]); + keyExchange(k[0], k[2]); + keyExchange(k[1], k[3]); + keyExchange(k[4], k[6]); + keyExchange(k[5], k[7]); + keyExchange(k[1], k[2]); + keyExchange(k[5], k[6]); + keyExchange(k[0], k[4]); + keyExchange(k[1], k[5]); + keyExchange(k[2], k[6]); + keyExchange(k[3], k[7]); + keyExchange(k[2], k[4]); + keyExchange(k[3], k[5]); + keyExchange(k[1], k[2]); + keyExchange(k[3], k[4]); + keyExchange(k[5], k[6]); +} + +//! Top-8 of one token by one warp; lane k < 8 receives slot k (expert, weight). Lane l holds experts 8l .. 8l + 7 +//! (its logits and its 16-byte slice of the BF16 bias, ``braw``); the selection key carries the expert id, so the +//! result does not depend on this placement. +__device__ void routeToken( + DraftRoutedMoeParams const& p, int token, int lane, RouteLogits const& in, uint4 braw, RouteSmem& rs) +{ + constexpr int kPerLane = kNumExperts / 32; + float const logit[kPerLane] = {in.l0.x, in.l0.y, in.l0.z, in.l0.w, in.l1.x, in.l1.y, in.l1.z, in.l1.w}; + uint32_t const biasPairs[kPerLane / 2] = {braw.x, braw.y, braw.z, braw.w}; + float sig[kPerLane]; + uint64_t key[kPerLane]; +#pragma unroll + for (int i = 0; i < kPerLane; ++i) + { + uint32_t const pair = biasPairs[i / 2]; + float const bias = __uint_as_float((i % 2 == 0 ? pair << 16 : pair & 0xFFFF0000u)); + sig[i] = routingSigmoid(logit[i]); + key[i] = packKey(addFtz(sig[i], bias), lane * kPerLane + i); + } + // The lane's keys in descending order; then every round takes the warp maximum of the lanes' leading keys and + // the owning lane drops its leading key (keys are unique: they carry the expert id). Every round thus selects + // the largest key not selected yet, as the native top-k reduction does. + sortKeysDescending(key); + float mySig = 0.f; + int myExpert = 0; +#pragma unroll + for (int k = 0; k < kTopK; ++k) + { + uint64_t const best = warpMax64(key[0]); + bool const pop = key[0] == best; +#pragma unroll + for (int i = 0; i + 1 < kPerLane; ++i) + { + key[i] = pop ? key[i + 1] : key[i]; + } + key[kPerLane - 1] = pop ? 0 : key[kPerLane - 1]; + int const e = 65535 - static_cast(best & 0xFFFF); + int const owner = e / kPerLane; + int const slot = e % kPerLane; + float s = 0.f; +#pragma unroll + for (int i = 0; i < kPerLane; ++i) + { + s = i == slot ? sig[i] : s; + } + s = __shfl_sync(0xffffffffu, s, owner); + if (lane == k) + { + mySig = s; + myExpert = e; + } + } + // cg::reduce(warp, sigmoid, plus): XOR butterfly 16, 8, 4, 2, 1; lanes >= 8 hold 0. + float sum = lane < kTopK ? mySig : 0.f; +#pragma unroll + for (int off = 16; off > 0; off >>= 1) + { + sum = addFtz(sum, __shfl_xor_sync(0xffffffffu, sum, off)); + } + if (lane < kTopK) + { + float const w = divApproxFtz(mulFtz(mySig, p.routedScaling), addFtz(sum, kSumEpsilon)); + rs.topExpert[token][lane] = myExpert; + rs.topWeight[token][lane] = __bfloat162float(__float2bfloat16_rn(w)); + atomicOr(&rs.mask[myExpert >> 5], 1u << (myExpert & 31)); + } +} + +//! Routing of all rows plus the distinct-expert list (ascending ids), by the whole CTA. ``tokensDone`` (optional) +//! receives clock64 when this thread's warp has routed its tokens. +template +__device__ void routeAll( + DraftRoutedMoeParams const& p, RouteSmem& rs, int warp, int lane, int64_t* tokensDone = nullptr) +{ + if (threadIdx.x < kNumExperts / 32) + { + rs.mask[threadIdx.x] = 0; + } + // A warp routes rows warp, warp + kWarps, ...: their logit loads are all issued before the first selection. + constexpr int kPerWarp = (kMaxRows + kWarps - 1) / kWarps; + uint4 const braw = *reinterpret_cast(p.bias + lane * (kNumExperts / 32)); + RouteLogits in[kPerWarp]; +#pragma unroll + for (int i = 0; i < kPerWarp; ++i) + { + if (warp + i * kWarps < p.rows) + { + in[i] = loadRouteLogits(p, warp + i * kWarps, lane); + } + } + __syncthreads(); // the mask is cleared before any row sets its bits +#pragma unroll + for (int i = 0; i < kPerWarp; ++i) + { + if (warp + i * kWarps < p.rows) + { + routeToken(p, warp + i * kWarps, lane, in[i], braw, rs); + } + } + if (tokensDone != nullptr) + { + *tokensDone = clock64(); + } + __syncthreads(); + // Distinct experts in ascending order: the index of expert e is the number of selected experts below it. + if (warp == 0 && lane < kNumExperts / 32) + { + int idx = 0; + for (int w = 0; w < lane; ++w) + { + idx += __popc(rs.mask[w]); + } + uint32_t m = rs.mask[lane]; + if (lane == kNumExperts / 32 - 1) + { + rs.numDistinct = idx + __popc(m); + } + while (m != 0) + { + rs.distinct[idx++] = lane * 32 + __ffs(m) - 1; + m &= m - 1; + } + } + for (int i = threadIdx.x; i < p.rows * kTopK; i += kWarps * 32) + { + int const e = rs.topExpert[i / kTopK][i % kTopK]; + int idx = __popc(rs.mask[e >> 5] & ((1u << (e & 31)) - 1u)); + for (int w = 0; w < (e >> 5); ++w) + { + idx += __popc(rs.mask[w]); + } + rs.slotDistinct[i / kTopK][i % kTopK] = idx; + } + __syncthreads(); +} + +// --------------------------------------------------------------------------------------------------------------- +// FC1: routing + gate/up GEMM + SwiGLU. +// --------------------------------------------------------------------------------------------------------------- + +template +__global__ void __launch_bounds__(kWarps * 32) mtpMoeFc1Kernel(DraftRoutedMoeParams const p) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + extern __shared__ __align__(128) uint8_t smem[]; + __shared__ RouteSmem rs; + int const warp = threadIdx.x >> 5; + int const lane = threadIdx.x & 31; + int const rows = p.rows; + uint8_t* xs = smem; + uint8_t* ring = smem + rows * kXRowBytes + warp * kStages * kChunkBytes; + uint32_t const xsAddr = smemAddr(xs); + uint32_t const ringAddr = smemAddr(ring); + + // Every input below (logits, activation) is produced by the predecessor grid. + gridDependencyWait(); + + // The activation rows, swizzled per 128-byte segment by token. + for (int i = threadIdx.x; i < rows * (kXRowBytes / 16); i += kWarps * 32) + { + int const t = i / (kXRowBytes / 16); + int const u = i % (kXRowBytes / 16); + cpAsync16(xsAddr + t * kXRowBytes + swizzled(t, u), p.x + t * kHidden + u * 8); + } + cpAsyncCommit(); + + routeAll(p, rs, warp, lane); + int const numDistinct = rs.numDistinct; + if (blockIdx.x == 0) + { + for (int i = threadIdx.x; i < kRouteWords; i += kWarps * 32) + { + int v = 0; + if (i == 0) + { + v = numDistinct; + } + else if (i < kRouteSlotOffset) + { + v = i - kRouteDistinctOffset < numDistinct ? rs.distinct[i - kRouteDistinctOffset] : -1; + } + else if (i < kRouteWeightOffset) + { + int const j = i - kRouteSlotOffset; + v = j / kTopK < rows ? rs.slotDistinct[j / kTopK][j % kTopK] : -1; + } + else + { + int const j = i - kRouteWeightOffset; + v = j / kTopK < rows ? __float_as_int(rs.topWeight[j / kTopK][j % kTopK]) : 0; + } + p.route[i] = v; + } + } + if (p.flags & kFlagEarlyTrigger) + { + gridDependencyTrigger(); + } + cpAsyncWait<0>(); + __syncthreads(); + + int const totalWarps = gridDim.x * kWarps; + int const gw = blockIdx.x * kWarps + warp; + int const items = numDistinct * kFc1Tiles; + int const myItems = items > gw ? (items - gw + totalWarps - 1) / totalWarps : 0; + int const total = myItems * kFc1Chunks; + // With kPairs, the warp's tiles are taken two at a time and their chunks interleave (A0 B0 A1 B1 ...): two + // independent accumulation chains per warp, each still in natural K order. An odd last tile runs alone. + int const pairChunks = kPairs ? (myItems / 2) * 2 * kFc1Chunks : 0; + + // Chunk n of the warp's stream -> (local tile index, K chunk). + auto chunkOf = [&](int n, int& tile, int& kc) + { + if (n < pairChunks) + { + int const r = n % (2 * kFc1Chunks); + tile = 2 * (n / (2 * kFc1Chunks)) + (r & 1); + kc = r >> 1; + } + else + { + tile = (pairChunks / kFc1Chunks) + (n - pairChunks) / kFc1Chunks; + kc = (n - pairChunks) % kFc1Chunks; + } + }; + + auto issue = [&](int n) + { + int tile, kc; + chunkOf(n, tile, kc); + int const item = gw + tile * totalWarps; + int const e = rs.distinct[item / kFc1Tiles]; + int const a = item % kFc1Tiles; + __nv_bfloat16 const* src = p.w13 + ((static_cast(e) * kFc1Chunks + kc) * kFc1Rows + a * 16) * kKChunk; + issueWeightChunk(ringAddr + (n % kStages) * kChunkBytes, src, lane); + }; + + auto epilogue = [&](int tile, float const(&acc)[4]) + { + int const item = gw + tile * totalWarps; + int const d = item / kFc1Tiles; + int const a = item % kFc1Tiles; + int const g = lane >> 2; + int const t = lane & 3; + // Rows 0..7 of the tile are up rows, rows 8..15 the gate rows of the same intermediate columns. + int const j = 16 * (a >> 1) + 2 * g + (a & 1); +#pragma unroll + for (int c = 0; c < 2; ++c) + { + int const tok = 2 * t + c; + if (tok < rows) + { + p.h[(static_cast(d) * kMaxRows + tok) * kInter + j] + = bf16SatFinite(swiGlu(acc[c], acc[2 + c])); + } + } + }; + +#pragma unroll 1 + for (int n = 0; n < kStages - 1; ++n) + { + if (n < total) + { + issue(n); + } + cpAsyncCommit(); + } + + float acc[4] = {-0.f, -0.f, -0.f, -0.f}; + float acc2[4] = {-0.f, -0.f, -0.f, -0.f}; +#pragma unroll 1 + for (int n = 0; n < total; ++n) + { + // Chunk n has landed once at most kStages - 2 newer groups are pending. + cpAsyncWait(); + __syncwarp(); + int tile, kc; + chunkOf(n, tile, kc); + uint32_t const slot = ringAddr + (n % kStages) * kChunkBytes; + if (n >= pairChunks) + { +#pragma unroll + for (int q = 0; q < 4; q += 2) + { + uint32_t b[4]; + loadB2(b, xsAddr, kXRowBytes, kc * 8 + 2 * q, rows, lane); + uint32_t a[4]; + loadA(a, slot, q, lane); + mmaBf16(acc, a, b[0], b[1]); + loadA(a, slot, q + 1, lane); + mmaBf16(acc, a, b[2], b[3]); + } + if (kc == kFc1Chunks - 1) + { + epilogue(tile, acc); + acc[0] = acc[1] = acc[2] = acc[3] = -0.f; + } + } + else if (n & 1) + { + // Both chunks of this K step of the pair are in: chunk n - 1 (first tile) and n (second tile). + uint32_t const slot0 = ringAddr + ((n - 1) % kStages) * kChunkBytes; +#pragma unroll + for (int q = 0; q < 4; q += 2) + { + uint32_t b[4]; + loadB2(b, xsAddr, kXRowBytes, kc * 8 + 2 * q, rows, lane); + uint32_t a0[4]; + uint32_t a1[4]; + loadA(a0, slot0, q, lane); + loadA(a1, slot, q, lane); + mmaBf16(acc, a0, b[0], b[1]); + mmaBf16(acc2, a1, b[0], b[1]); + loadA(a0, slot0, q + 1, lane); + loadA(a1, slot, q + 1, lane); + mmaBf16(acc, a0, b[2], b[3]); + mmaBf16(acc2, a1, b[2], b[3]); + } + if (kc == kFc1Chunks - 1) + { + epilogue(tile - 1, acc); + epilogue(tile, acc2); + acc[0] = acc[1] = acc[2] = acc[3] = -0.f; + acc2[0] = acc2[1] = acc2[2] = acc2[3] = -0.f; + } + } + __syncwarp(); + // The slot refilled now held chunk n - 1, consumed in this or the previous iteration. + if (n + kStages - 1 < total) + { + issue(n + kStages - 1); + } + cpAsyncCommit(); + } + cpAsyncWait<0>(); + + if (p.flags & kFlagPrefetchFc2) + { + int const pieces = numDistinct * kFc2Chunks * kPrefetchPiecesPerBlock; + for (int piece = gw * 32 + lane; piece < pieces; piece += totalWarps * 32) + { + int const d = piece / (kFc2Chunks * kPrefetchPiecesPerBlock); + int const rem = piece % (kFc2Chunks * kPrefetchPiecesPerBlock); + int const kc = rem / kPrefetchPiecesPerBlock; + int const off = rem % kPrefetchPiecesPerBlock; + uint8_t const* addr = reinterpret_cast(p.w2 + + (static_cast(rs.distinct[d]) * kFc2Chunks + kc) * kHidden * kKChunk) + + static_cast(off) * kPrefetchPieceBytes; + asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;" ::"l"(__cvta_generic_to_global(addr)), + "r"(kPrefetchPieceBytes) + : "memory"); + } + } + if (!(p.flags & kFlagEarlyTrigger)) + { + gridDependencyTrigger(); + } +#endif +} + +// --------------------------------------------------------------------------------------------------------------- +// FC2: down GEMM per distinct expert + native finalize, 16 hidden rows per work item. +// --------------------------------------------------------------------------------------------------------------- + +struct Fc2Smem +{ + float weight[kMaxRows][kTopK]; + int distinct[kMaxDistinct]; + int8_t slotOf[kMaxRows][kMaxDistinct]; + int numDistinct; +}; + +__device__ __forceinline__ int hiddenOfStoredRow(int s) +{ + return (s & ~31) | ((s & 7) << 2) | ((s & 31) >> 3); +} + +template +__global__ void __launch_bounds__(kWarps * 32) mtpMoeFc2Kernel(DraftRoutedMoeParams const p) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + extern __shared__ __align__(128) uint8_t smem[]; + __shared__ Fc2Smem fs; + int const warp = threadIdx.x >> 5; + int const lane = threadIdx.x & 31; + int const rows = p.rows; + constexpr int kWarpBytes = kStages * kChunkBytes + kFc2HStages * kHEntryBytes + kStageBytes; + uint8_t* mine = smem + warp * kWarpBytes; + uint32_t const ringAddr = smemAddr(mine); + uint32_t const hAddr = ringAddr + kStages * kChunkBytes; + __nv_bfloat16* stage = reinterpret_cast<__nv_bfloat16*>(mine + kStages * kChunkBytes + kFc2HStages * kHEntryBytes); + + // The routing record and the intermediate are written by the FC1 grid. + gridDependencyWait(); + + int const numDistinct = p.route[0]; + for (int i = threadIdx.x; i < kMaxRows * kMaxDistinct; i += kWarps * 32) + { + fs.slotOf[i / kMaxDistinct][i % kMaxDistinct] = -1; + } + for (int i = threadIdx.x; i < kMaxDistinct; i += kWarps * 32) + { + fs.distinct[i] = p.route[kRouteDistinctOffset + i]; + } + __syncthreads(); + for (int i = threadIdx.x; i < rows * kTopK; i += kWarps * 32) + { + int const t = i / kTopK; + int const k = i % kTopK; + fs.slotOf[t][p.route[kRouteSlotOffset + i]] = static_cast(k); + fs.weight[t][k] = __int_as_float(p.route[kRouteWeightOffset + i]); + } + __syncthreads(); + + int const totalWarps = gridDim.x * kWarps; + int const gw = blockIdx.x * kWarps + warp; + int const perItem = numDistinct * kFc2Chunks; + int const myItems = kFc2Tiles > gw ? (kFc2Tiles - gw + totalWarps - 1) / totalWarps : 0; + int const total = myItems * perItem; + // With kPairs, the experts of an item are taken two at a time and their chunks interleave (two independent + // accumulation chains, each in natural K order); an odd last expert runs alone. + int const pairChunks = kPairs ? (numDistinct / 2) * 2 * kFc2Chunks : 0; + + // Chunk m of an item -> (distinct expert index, K chunk). + auto chunkOf = [&](int m, int& d, int& kc) + { + if (m < pairChunks) + { + int const r = m % (2 * kFc2Chunks); + d = 2 * (m / (2 * kFc2Chunks)) + (r & 1); + kc = r >> 1; + } + else + { + d = (pairChunks / kFc2Chunks) + (m - pairChunks) / kFc2Chunks; + kc = (m - pairChunks) % kFc2Chunks; + } + }; + // Intermediate entry of expert d of local item il (entries are consumed in this order). + auto hEntry = [&](int il, int d) { return hAddr + ((il * numDistinct + d) % kFc2HStages) * kHEntryBytes; }; + + auto issue = [&](int n) + { + int const il = n / perItem; + int d, kc; + chunkOf(n % perItem, d, kc); + int const tile = gw + il * totalWarps; + int const e = fs.distinct[d]; + __nv_bfloat16 const* src = p.w2 + ((static_cast(e) * kFc2Chunks + kc) * kHidden + tile * 16) * kKChunk; + issueWeightChunk(ringAddr + (n % kStages) * kChunkBytes, src, lane); + if (kc == 0) + { + uint32_t const entry = hEntry(il, d); + for (int u = lane; u < rows * (kInter / 8); u += 32) + { + int const tok = u / (kInter / 8); + int const unit = u % (kInter / 8); + cpAsync16(entry + tok * (kInter * 2) + swizzled(tok, unit), + p.h + (static_cast(d) * kMaxRows + tok) * kInter + unit * 8); + } + } + }; + + // FC2 rows of expert d (C fragment) -> BF16, into the token's slot of the staging area. + auto scatter = [&](int tile, int d, float const(&acc)[4]) + { + int const g = lane >> 2; + int const t = lane & 3; +#pragma unroll + for (int c = 0; c < 2; ++c) + { + int const tok = 2 * t + c; + int const k = tok < rows ? fs.slotOf[tok][d] : -1; + if (k >= 0) + { + __nv_bfloat16 const lo = bf16SatFinite(acc[c]); + __nv_bfloat16 const hi = bf16SatFinite(acc[2 + c]); + stage[(tok * kTopK + k) * 16 + g] = lo; + stage[(tok * kTopK + k) * 16 + g + 8] = hi; + if (p.debugFc2 != nullptr) + { + int64_t const base = (static_cast(tok) * kTopK + k) * kHidden; + p.debugFc2[base + hiddenOfStoredRow(tile * 16 + g)] = lo; + p.debugFc2[base + hiddenOfStoredRow(tile * 16 + g + 8)] = hi; + } + } + } + }; + + // Native finalize of the item's 16 hidden rows for every token: FP32 fma chain from 0 in slot order. + auto finalize = [&](int tile) + { + __syncwarp(); + for (int pr = lane; pr < rows * 16; pr += 32) + { + int const tok = pr >> 4; + int const row = pr & 15; + float sum = 0.f; +#pragma unroll + for (int k = 0; k < kTopK; ++k) + { + sum = fmaFtz(fs.weight[tok][k], __bfloat162float(stage[(tok * kTopK + k) * 16 + row]), sum); + } + p.out[static_cast(tok) * kHidden + hiddenOfStoredRow(tile * 16 + row)] = __float2bfloat16_rn(sum); + } + }; + +#pragma unroll 1 + for (int n = 0; n < kStages - 1; ++n) + { + if (n < total) + { + issue(n); + } + cpAsyncCommit(); + } + + float acc[4] = {-0.f, -0.f, -0.f, -0.f}; + float acc2[4] = {-0.f, -0.f, -0.f, -0.f}; +#pragma unroll 1 + for (int n = 0; n < total; ++n) + { + cpAsyncWait(); + __syncwarp(); + int const il = n / perItem; + int const m = n % perItem; + int const tile = gw + il * totalWarps; + int d, kc; + chunkOf(m, d, kc); + uint32_t const slot = ringAddr + (n % kStages) * kChunkBytes; + if (m >= pairChunks) + { + uint32_t const entry = hEntry(il, d); +#pragma unroll + for (int q = 0; q < 4; q += 2) + { + uint32_t b[4]; + loadB2(b, entry, kInter * 2, kc * 8 + 2 * q, rows, lane); + uint32_t a[4]; + loadA(a, slot, q, lane); + mmaBf16(acc, a, b[0], b[1]); + loadA(a, slot, q + 1, lane); + mmaBf16(acc, a, b[2], b[3]); + } + if (kc == kFc2Chunks - 1) + { + scatter(tile, d, acc); + acc[0] = acc[1] = acc[2] = acc[3] = -0.f; + } + } + else if (m & 1) + { + uint32_t const slot0 = ringAddr + ((n - 1) % kStages) * kChunkBytes; + uint32_t const entry0 = hEntry(il, d - 1); + uint32_t const entry1 = hEntry(il, d); +#pragma unroll + for (int q = 0; q < 4; q += 2) + { + uint32_t b0[4]; + uint32_t b1[4]; + loadB2(b0, entry0, kInter * 2, kc * 8 + 2 * q, rows, lane); + loadB2(b1, entry1, kInter * 2, kc * 8 + 2 * q, rows, lane); + uint32_t a0[4]; + uint32_t a1[4]; + loadA(a0, slot0, q, lane); + loadA(a1, slot, q, lane); + mmaBf16(acc, a0, b0[0], b0[1]); + mmaBf16(acc2, a1, b1[0], b1[1]); + loadA(a0, slot0, q + 1, lane); + loadA(a1, slot, q + 1, lane); + mmaBf16(acc, a0, b0[2], b0[3]); + mmaBf16(acc2, a1, b1[2], b1[3]); + } + if (kc == kFc2Chunks - 1) + { + scatter(tile, d - 1, acc); + scatter(tile, d, acc2); + acc[0] = acc[1] = acc[2] = acc[3] = -0.f; + acc2[0] = acc2[1] = acc2[2] = acc2[3] = -0.f; + } + } + if (m == perItem - 1) + { + finalize(tile); + } + __syncwarp(); + if (n + kStages - 1 < total) + { + issue(n + kStages - 1); + } + cpAsyncCommit(); + } + cpAsyncWait<0>(); +#endif +} + +// --------------------------------------------------------------------------------------------------------------- +// TMA-fed kernels (kFlagTma). Per CTA one producer thread (last warp, lane 0) streams 3D weight boxes (16 rows per +// chain x 64 elements x several K chunks of the BlockMajorK layout, 128-byte swizzle = the layout loadA reads; 32 KB +// per FC1 box, one expert's K = 256 per FC2 box) into a ring of S stages; W MMA warps each own C 16-row accumulation +// chains per tile and consume every stage. Full barriers complete on the transaction count; each MMA warp releases +// a stage (empty barrier, count W) after the MMAs that read it. The arithmetic of every chain (fragments, k-step +// order, accumulator init, epilogues) is that of the kernels above. +// --------------------------------------------------------------------------------------------------------------- + +#ifndef GLM_MOE_DSA_DRAFT_ROUTED_MOE_TIMELINE +#define GLM_MOE_DSA_DRAFT_ROUTED_MOE_TIMELINE 0 +#endif + +//! Probe builds (GLM_MOE_DSA_DRAFT_ROUTED_MOE_TIMELINE=1) record per-CTA events into DraftRoutedMoeParams::timeline; +//! production builds compile none of it. +constexpr bool kTimeline = GLM_MOE_DSA_DRAFT_ROUTED_MOE_TIMELINE != 0; + +//! Columns of a CTA's event record: %globaltimer at entry and when MMA warp 0 is done (ns), the SM id, clock64 +//! cycles after entry, and counts. +enum TimelineSlot : int +{ + kTlEntryNs = 0, + kTlDoneNs = 1, + kTlSmId = 2, + kTlWait = 3, //!< griddepcontrol.wait returned + kTlSetup = 4, //!< routing (FC1) / routing record (FC2) in shared memory, roles split + kTlFirst = 5, //!< MMA warp 0 found its first stage full + kTlLast = 6, //!< MMA warp 0 found its last stage full + kTlDone = 7, //!< MMA warp 0 done (epilogue / finalize stored) + kTlProdFirst = 8, //!< producer issued its first weight copy + kTlProdLast = 9, //!< producer issued its last weight copy + kTlFullBlocked = 10, //!< MMA warp 0: stage waits whose first poll failed + kTlEmptyBlocked = 11, //!< producer: re-arm waits whose first poll failed + kTlChunks = 12, //!< MMA warp 0: K chunks consumed + kTlRouteTokens = 13, //!< FC1 warp 0 has routed its tokens + kTlRouted = 14, //!< FC1 routing and distinct-expert list done (before the activation wait) +}; + +//! K chunks (2 KB per 16-row chain) per stage of the TMA-fed FC1 kernel: one 32 KB 3D box per stage. +static_assert(kTmaStageChunks * kChunkBytes == 32 * 1024, "32 KB FC1 stages"); + +__device__ __forceinline__ int64_t globalTimerNs() +{ + uint64_t t; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t)); + return static_cast(t); +} + +__device__ __forceinline__ int64_t smId() +{ + uint32_t id; + asm volatile("mov.u32 %0, %%smid;" : "=r"(id)); + return id; +} + +//! This CTA's event record of kernel ``kernel`` (0: FC1, 1: FC2), or null. +__device__ __forceinline__ int64_t* timelineRow(DraftRoutedMoeParams const& p, int kernel) +{ + if (!kTimeline || p.timeline == nullptr || blockIdx.x >= kTimelineCtas) + { + return nullptr; + } + return p.timeline + (static_cast(kernel) * kTimelineCtas + blockIdx.x) * kTimelineSlots; +} + +__device__ __forceinline__ void mbarInit(uint32_t bar, uint32_t count) +{ + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" ::"r"(bar), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarFenceInit() +{ + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); +} + +__device__ __forceinline__ void mbarArriveExpectTx(uint32_t bar, uint32_t bytes) +{ + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(bar), "r"(bytes) : "memory"); +} + +__device__ __forceinline__ void mbarArrive(uint32_t bar) +{ + asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0];" ::"r"(bar) : "memory"); +} + +__device__ __forceinline__ bool mbarTryWait(uint32_t bar, uint32_t parity) +{ + uint32_t done; + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64 p, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, p;\n\t" + "}" + : "=r"(done) + : "r"(bar), "r"(parity) + : "memory"); + return done != 0; +} + +//! Non-blocking probe of the phase of ``bar`` with parity ``parity`` (acquire when it returns true). +__device__ __forceinline__ bool mbarTestWait(uint32_t bar, uint32_t parity) +{ + uint32_t done; + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "mbarrier.test_wait.parity.shared::cta.b64 p, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, p;\n\t" + "}" + : "=r"(done) + : "r"(bar), "r"(parity) + : "memory"); + return done != 0; +} + +//! Blocks until the phase of ``bar`` with parity ``parity`` has completed. Every wait is between threads of one CTA +//! (never on another grid), so one that does not complete within 2^28 polls is a protocol error: trap, do not hang. +__device__ __forceinline__ void mbarWait(uint32_t bar, uint32_t parity) +{ + for (uint32_t i = 0; !mbarTryWait(bar, parity); ++i) + { + if (i == (1u << 28)) + { + __trap(); + } + } +} + +//! mbarWait; returns 1 when the first poll failed (probe builds count these), else 0. +[[maybe_unused]] __device__ __forceinline__ int mbarWaitCounted(uint32_t bar, uint32_t parity) +{ + if (mbarTryWait(bar, parity)) + { + return 0; + } + mbarWait(bar, parity); + return 1; +} + +//! Barrier among the W MMA warps (threads 0 .. 32 W - 1) of a warp-specialized kernel; the producer warp does not +//! take part. Named barrier 1 (``__syncthreads`` uses 0). +template +__device__ __forceinline__ void mmaWarpsSync() +{ + asm volatile("bar.sync 1, %0;" ::"n"(W * 32) : "memory"); +} + +__device__ __forceinline__ void tmaLoad3d(uint32_t dst, CUtensorMap const* map, int c0, int c1, int c2, uint32_t bar) +{ + asm volatile( + "cp.async.bulk.tensor.3d.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1, {%2, %3, %4}], " + "[%5];" ::"r"(dst), + "l"(reinterpret_cast(map)), "r"(c0), "r"(c1), "r"(c2), "r"(bar) + : "memory"); +} + +//! tmaLoad3d with an L2 cache policy (createpolicy). +__device__ __forceinline__ void tmaLoad3dHint( + uint32_t dst, CUtensorMap const* map, int c0, int c1, int c2, uint32_t bar, uint64_t policy) +{ + asm volatile( + "cp.async.bulk.tensor.3d.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint [%0], [%1, {%2, " + "%3, %4}], [%5], %6;" ::"r"(dst), + "l"(reinterpret_cast(map)), "r"(c0), "r"(c1), "r"(c2), "r"(bar), "l"(policy) + : "memory"); +} + +//! L2 policy that evicts the lines it is attached to first. +__device__ __forceinline__ uint64_t l2EvictFirstPolicy() +{ + uint64_t policy; + asm volatile("createpolicy.fractional.L2::evict_first.b64 %0, 1.0;" : "=l"(policy)); + return policy; +} + +__device__ __forceinline__ void tmaPrefetchL2(CUtensorMap const* map, int c0, int c1, int c2) +{ + asm volatile( + "cp.async.bulk.prefetch.tensor.3d.L2.global.tile [%0, {%1, %2, %3}];" ::"l"(reinterpret_cast(map)), + "r"(c0), "r"(c1), "r"(c2) + : "memory"); +} + +__device__ __forceinline__ void bulkLoad(uint32_t dst, void const* src, uint32_t bytes, uint32_t bar) +{ + asm volatile("cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];" ::"r"(dst), + "l"(__cvta_generic_to_global(src)), "r"(bytes), "r"(bar) + : "memory"); +} + +__device__ __forceinline__ void prefetchTensorMap(CUtensorMap const* map) +{ + asm volatile("prefetch.tensormap [%0];" ::"l"(reinterpret_cast(map)) : "memory"); +} + +//! L2 prefetch of this CTA's share of the distinct experts' FC2 weights (16 KB pieces, expert-major, CTA-strided, so +//! the experts the FC2 grid reads first are fetched first). Weights only: legal at any point of the FC1 grid. +__device__ __forceinline__ void prefetchFc2Share(DraftRoutedMoeParams const& p, int const* distinct, int numDistinct) +{ + constexpr int kPiecesPerExpert = kFc2Chunks * kPrefetchPiecesPerBlock; + int const pieces = numDistinct * kPiecesPerExpert; + for (int piece = blockIdx.x; piece < pieces; piece += gridDim.x) + { + uint8_t const* addr + = reinterpret_cast( + p.w2 + static_cast(distinct[piece / kPiecesPerExpert]) * kFc2Chunks * kHidden * kKChunk) + + static_cast(piece % kPiecesPerExpert) * kPrefetchPieceBytes; + asm volatile( + "cp.async.bulk.prefetch.L2.global [%0], %1;" ::"l"(__cvta_generic_to_global(addr)), "r"(kPrefetchPieceBytes) + : "memory"); + } +} + +//! Fragments of one 64-element K chunk for C 16-row chains: A of k-steps 0..3 per chain, B of k-steps (0,1), (2,3). +template +struct ChunkFrags +{ + uint32_t a[C][4][4]; + uint32_t b[2][4]; +}; + +//! The four k-steps of a chunk in natural K order into each chain's own accumulator (the sequence of the kernels +//! above; the chains are interleaved per k-step and never mixed). +template +__device__ __forceinline__ void mmaChunk(float (&acc)[C][4], ChunkFrags const& f) +{ +#pragma unroll + for (int q = 0; q < 4; ++q) + { +#pragma unroll + for (int c = 0; c < C; ++c) + { + mmaBf16(acc[c], f.a[c][q], f.b[q >> 1][2 * (q & 1)], f.b[q >> 1][2 * (q & 1) + 1]); + } + } +} + +//! Per-lane ldmatrix offsets inside one K chunk. ``a[q]``: the lane's row of k-step q in a 2 KB 16-row chunk (as +//! loadA). ``b[h]``: the lane's 16-byte unit of k-steps (2h, 2h + 1) in the first 128-byte K chunk of a [token][k] +//! tile with ``bRowBytes`` per token (as loadB2). The swizzle only permutes the 16-byte units inside a 128-byte K +//! chunk, so K chunk kc of a token row is the same pattern ``kc * 128`` bytes further. +struct LaneOffsets +{ + uint32_t a[4]; + uint32_t b[2]; +}; + +__device__ __forceinline__ LaneOffsets laneOffsets(int rows, int bRowBytes, int lane) +{ + LaneOffsets o; + int const m = lane >> 3; + int const r = (lane & 7) | ((m & 1) << 3); +#pragma unroll + for (int q = 0; q < 4; ++q) + { + o.a[q] = r * 128 + swizzled(r, 2 * q + (m >> 1)); + } + int tok = lane & 7; + tok = tok < rows ? tok : 0; +#pragma unroll + for (int h = 0; h < 2; ++h) + { + o.b[h] = tok * bRowBytes + swizzled(tok, 4 * h + m); + } + return o; +} + +//! Per-lane ldmatrix offsets of the B fragments in a streamed activation slice (kFlagStreamX): one TMA box of +//! ``DraftRoutedMoeTensorMaps::x`` is [8 tokens][kXBoxChunks K chunks][128 bytes] with the 128-byte swizzle, so K chunk +//! p of token t is line 2t + p of the box. ``b[p][h]``: the lane's 16-byte unit of k-steps (2h, 2h + 1) of K chunk p. +struct LaneOffsetsX +{ + uint32_t b[2][2]; +}; + +__device__ __forceinline__ LaneOffsetsX laneOffsetsX(int rows, int lane) +{ + static_assert(kXBoxChunks == 2, "two K chunks per activation box"); + LaneOffsetsX o; + int const m = lane >> 3; + int tok = lane & 7; + tok = tok < rows ? tok : 0; +#pragma unroll + for (int q = 0; q < 2; ++q) + { + int const line = kXBoxChunks * tok + q; +#pragma unroll + for (int h = 0; h < 2; ++h) + { + o.b[q][h] = line * 128 + swizzled(line, 4 * h + m); + } + } + return o; +} + +//! The fragments of one K chunk for C 16-row chains (chain c's 2 KB chunk follows chain 0's at c * kChunkBytes) +//! from per-stage bases and constant chunk offsets, so the loads of a stage need no address arithmetic. +//! ``aBase[q]``: chain 0 of the stage's first chunk + LaneOffsets::a[q]; ``bBase[h]``: the activation K chunk of the +//! stage's first chunk + LaneOffsets::b[h]. +template +__device__ __forceinline__ void loadFragsAt( + ChunkFrags& f, uint32_t const (&aBase)[4], uint32_t aOff, uint32_t const (&bBase)[2], uint32_t bOff) +{ +#pragma unroll + for (int c = 0; c < C; ++c) + { +#pragma unroll + for (int q = 0; q < 4; ++q) + { + ldmatrixX4(f.a[c][q], aBase[q] + aOff + c * kChunkBytes); + } + } +#pragma unroll + for (int h = 0; h < 2; ++h) + { + ldmatrixX4(f.b[h], bBase[h] + bOff); + } +} + +//! loadFragsAt for chains ``kChainStride`` bytes apart (the sub-tile boxes of the triples path); chains ``c >= valid`` +//! (absent sub-tiles of the grid's last CTA) are not loaded. +template +__device__ __forceinline__ void loadFragsStrided( + ChunkFrags& f, uint32_t const (&aBase)[4], uint32_t aOff, uint32_t const (&bBase)[2], uint32_t bOff, int valid) +{ +#pragma unroll + for (int c = 0; c < C; ++c) + { + if (c < valid) + { +#pragma unroll + for (int q = 0; q < 4; ++q) + { + ldmatrixX4(f.a[c][q], aBase[q] + aOff + c * kChainStride); + } + } + } +#pragma unroll + for (int h = 0; h < 2; ++h) + { + ldmatrixX4(f.b[h], bBase[h] + bOff); + } +} + +//! mmaChunk for the chains ``c < valid`` only (each chain's k-steps in natural order, as mmaChunk). +template +__device__ __forceinline__ void mmaChunkValid(float (&acc)[C][4], ChunkFrags const& f, int valid) +{ +#pragma unroll + for (int q = 0; q < 4; ++q) + { +#pragma unroll + for (int c = 0; c < C; ++c) + { + if (c < valid) + { + mmaBf16(acc[c], f.a[c][q], f.b[q >> 1][2 * (q & 1)], f.b[q >> 1][2 * (q & 1) + 1]); + } + } + } +} + +//! Fragments of one K chunk for E FC2 experts of one 16-row tile: per expert its A (weights) and B (intermediate). +template +struct ExpertFrags +{ + uint32_t a[E][4][4]; + uint32_t b[E][2][4]; +}; + +//! Byte offset of element ``j`` of token ``tok`` in an intermediate entry stored as its shared-memory image (512-byte +//! token rows, 16-byte units swizzled by token as loadB2 reads them). +__device__ __forceinline__ int hImageElem(int tok, int j) +{ + int const u = j >> 3; + return tok * kInter + (((u & ~7) | ((u & 7) ^ (tok & 7))) << 3) + (j & 7); +} + +//! TMA-fed FC1 kernel: W MMA warps with C chains each, a ring of S stages of 32 KB (one box of ``map``). With S3 > 0 +//! (4 MMA warps, kFlagFc1Triples) a call whose tiles exceed the grid but whose 64-row sub-tiles fit three per CTA +//! runs the triples path instead (decided per call from the routing, identically in every CTA): CTA b takes sub-tiles +//! 3b .. 3b + 2 (in distinct-expert order), MMA warp w owns rows 16w .. 16w + 15 of each (three chains), and a stage +//! of its ring of S3 is three boxes of ``map3`` (64 rows x kTripleStageChunks K chunks), one per sub-tile. +//! With X (kFlagStreamX) the activation rows are not kept in shared memory: every stage also carries the activation +//! K chunks it needs (boxes of ``mapX``, kXBoxBytes each, after the weight boxes), so the ring can be deeper. +template +__global__ void __launch_bounds__((W + 1) * 32) + mtpMoeFc1TmaKernel(DraftRoutedMoeParams const p, __grid_constant__ CUtensorMap const map, + __grid_constant__ CUtensorMap const map3, __grid_constant__ CUtensorMap const mapX) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + constexpr int kChains = W * C; + constexpr int kKs = kTmaStageChunks / kChains; // K chunks per stage (one TMA box) + constexpr int kChunkSet = kChains * kChunkBytes; + constexpr int kStage = kKs * kChunkSet; + constexpr int kTileRows = 16 * kChains; + constexpr int kTilesPerExpert = kFc1Rows / kTileRows; + constexpr int kStagesPerTile = kFc1Chunks / kKs; + constexpr int kL2Ahead = 2 * S; + static_assert(kFc1Chunks % kKs == 0, "a tile must be a whole number of stages"); + static_assert(kKs % 2 == 0, "the fragment double buffer restarts at f[0] on every stage"); + // Triples path (S3 > 0): sub-tiles of 16 W rows, three per CTA; a stage holds one box per sub-tile. + constexpr int kSubRows = 16 * W; + constexpr int kSubPerExpert = kFc1Rows / kSubRows; + constexpr int kChunkStride3 = kSubRows * kKChunk * 2; // next K chunk inside a sub-tile box + constexpr int kBox3 = kTripleStageChunks * kChunkStride3; + constexpr int kStage3 = kTripleSubTiles * kBox3; + constexpr int kStagesPerTile3 = kFc1Chunks / kTripleStageChunks; + static_assert(S3 == 0 || (W == 4 && kSubRows == kTripleRows), "the triples path has 4 MMA warps (64-row boxes)"); + static_assert(kTripleStageChunks == 2, "the triples consumer walks two K chunks per stage"); + // Ring slot strides: the weight boxes, then (X) the activation boxes of the stage's K chunks. + constexpr int kXBoxes = X ? kKs / kXBoxChunks : 0; + constexpr int kSlot = kStage + kXBoxes * kXBoxBytes; + constexpr int kSlot3 = kStage3 + (X ? kXBoxBytes : 0); + static_assert(!X || (kKs % kXBoxChunks == 0 && kTripleStageChunks == kXBoxChunks), "whole activation boxes"); + static_assert(kSlot % 1024 == 0 && kSlot3 % 1024 == 0, "slots stay aligned to the 128-byte swizzle pattern"); + constexpr int kRingBytes = S * kSlot > S3 * kSlot3 ? S * kSlot : S3 * kSlot3; + constexpr int kBars = S > S3 ? S : S3; + extern __shared__ uint8_t smemRaw[]; + __shared__ RouteSmem rs; + __shared__ __align__(8) uint64_t bars[2 * kBars]; + int64_t const t0 = kTimeline ? clock64() : 0; + int64_t const entryNs = kTimeline ? globalTimerNs() : 0; + int const warp = threadIdx.x >> 5; + int const lane = threadIdx.x & 31; + int const rows = p.rows; + uint32_t const ringAddr = (smemAddr(smemRaw) + 1023u) & ~1023u; + uint32_t const xsAddr = ringAddr + kRingBytes; + uint32_t const fullBar = smemAddr(&bars[0]); + uint32_t const emptyBar = smemAddr(&bars[kBars]); + if (threadIdx.x == 0) + { + for (int s = 0; s < kBars; ++s) + { + mbarInit(fullBar + 8 * s, 1); + mbarInit(emptyBar + 8 * s, W * 32); // every MMA thread releases its own reads + } + mbarFenceInit(); + prefetchTensorMap(&map); + if constexpr (S3 > 0) + { + prefetchTensorMap(&map3); + } + if constexpr (X) + { + prefetchTensorMap(&mapX); + } + } + + // Every input below (logits, activation) is produced by the predecessor grid. + gridDependencyWait(); + int64_t const tWait = kTimeline ? clock64() - t0 : 0; + + // The MMA warps stage the activation rows (they alone read them); the producer warp does not wait for them. + for (int i = threadIdx.x; i < (X || warp == W ? 0 : rows * (kXRowBytes / 16)); i += W * 32) + { + int const t = i / (kXRowBytes / 16); + int const u = i % (kXRowBytes / 16); + cpAsync16(xsAddr + t * kXRowBytes + swizzled(t, u), p.x + t * kHidden + u * 8); + } + cpAsyncCommit(); + + int64_t tTokens = 0; + // Ends with __syncthreads (also publishes the barrier initialization). + routeAll(p, rs, warp, lane, kTimeline ? &tTokens : nullptr); + int64_t const tRouted = kTimeline ? clock64() - t0 : 0; + int const numDistinct = rs.numDistinct; + if (blockIdx.x == 0) + { + for (int i = threadIdx.x; i < kRouteWords; i += (W + 1) * 32) + { + int v = 0; + if (i == 0) + { + v = numDistinct; + } + else if (i < kRouteSlotOffset) + { + v = i - kRouteDistinctOffset < numDistinct ? rs.distinct[i - kRouteDistinctOffset] : -1; + } + else if (i < kRouteWeightOffset) + { + int const j = i - kRouteSlotOffset; + v = j / kTopK < rows ? rs.slotDistinct[j / kTopK][j % kTopK] : -1; + } + else + { + int const j = i - kRouteWeightOffset; + v = j / kTopK < rows ? __float_as_int(rs.topWeight[j / kTopK][j % kTopK]) : 0; + } + p.route[i] = v; + } + } + int const tiles = numDistinct * kTilesPerExpert; + bool const noMma = kTimeline && (p.flags & kFlagProbeNoMma) != 0; + bool const resident = kTimeline && (p.flags & kFlagProbeResident) != 0; + // The path depends only on numDistinct and the launch, so every CTA takes the same one. + int const numSub = numDistinct * kSubPerExpert; + int const sub0 = kTripleSubTiles * static_cast(blockIdx.x); + bool const triples = S3 > 0 && (p.flags & kFlagFc1Triples) != 0 && !resident && tiles > static_cast(gridDim.x) + && numSub <= kTripleSubTiles * static_cast(gridDim.x); + int const valid3 = min(kTripleSubTiles, numSub - sub0); // sub-tiles of this CTA (<= 0: none) + // Early trigger (after routing): always with kFlagEarlyTrigger, or for calls on the triples path only. + bool const early = (p.flags & kFlagEarlyTrigger) != 0 || (triples && (p.flags & kFlagEarlyTriggerTriples) != 0); + if (early) + { + gridDependencyTrigger(); + } + bool const evictFirst = (p.flags & kFlagFc1EvictFirst) != 0; + + if (warp == W) + { + if constexpr (S3 > 0) + { + if (lane == 0 && triples) + { + uint64_t const policy = evictFirst ? l2EvictFirstPolicy() : 0; + if (p.flags & kFlagPrefetchFc2Early) + { + prefetchFc2Share(p, rs.distinct, numDistinct); + } + int emptyBlocked = 0; + int64_t tFirst = 0; + for (int u = 0; u < (valid3 > 0 ? kStagesPerTile3 : 0); ++u) + { + int const s = u % S3; + if (u >= S3) + { + if constexpr (kTimeline) + { + emptyBlocked += mbarWaitCounted(emptyBar + 8 * s, ((u / S3) - 1) & 1); + } + else + { + mbarWait(emptyBar + 8 * s, ((u / S3) - 1) & 1); + } + } + mbarArriveExpectTx(fullBar + 8 * s, valid3 * kBox3 + (X ? kXBoxBytes : 0)); + if constexpr (X) + { + tmaLoad3d( + ringAddr + s * kSlot3 + kStage3, &mapX, 0, u * kTripleStageChunks, 0, fullBar + 8 * s); + } + for (int c = 0; c < valid3; ++c) + { + int const sub = sub0 + c; + uint32_t const dst = ringAddr + s * kSlot3 + c * kBox3; + int const row = (sub % kSubPerExpert) * kSubRows; + int const slab = rs.distinct[sub / kSubPerExpert] * kFc1Chunks + u * kTripleStageChunks; + if (evictFirst) + { + tmaLoad3dHint(dst, &map3, 0, row, slab, fullBar + 8 * s, policy); + } + else + { + tmaLoad3d(dst, &map3, 0, row, slab, fullBar + 8 * s); + } + } + if (kTimeline && u == 0) + { + tFirst = clock64() - t0; + } + } + int64_t const tLast = kTimeline ? clock64() - t0 : 0; + if ((p.flags & kFlagPrefetchFc2) && !(p.flags & kFlagPrefetchFc2Early)) + { + prefetchFc2Share(p, rs.distinct, numDistinct); + } + if (int64_t* row = timelineRow(p, 0)) + { + row[kTlProdFirst] = tFirst; + row[kTlProdLast] = tLast; + row[kTlEmptyBlocked] = emptyBlocked; + } + } + } + if (lane == 0 && !triples) + { + int const myTiles + = static_cast(blockIdx.x) < tiles ? (tiles - 1 - static_cast(blockIdx.x)) / gridDim.x + 1 : 0; + int const total = myTiles * kStagesPerTile; + // Box coordinates of this CTA's u-th stage: row of the tile, (expert, first K chunk) slab. + auto rowOf = [&](int u) + { + int const t = blockIdx.x + (u / kStagesPerTile) * gridDim.x; + return (t % kTilesPerExpert) * kTileRows; + }; + auto slabOf = [&](int u) + { + int const t = blockIdx.x + (u / kStagesPerTile) * gridDim.x; + return rs.distinct[t / kTilesPerExpert] * kFc1Chunks + (u % kStagesPerTile) * kKs; + }; + bool const l2 = (p.flags & kFlagL2Prefetch) != 0; + uint64_t const policy = evictFirst ? l2EvictFirstPolicy() : 0; + if (p.flags & kFlagPrefetchFc2Early) + { + prefetchFc2Share(p, rs.distinct, numDistinct); + } + if (l2) + { + for (int u = 0; u < kL2Ahead && u < total; ++u) + { + tmaPrefetchL2(&map, 0, rowOf(u), slabOf(u)); + } + } + int emptyBlocked = 0; + int64_t tFirst = 0; + for (int u = 0; u < total && !(resident && u == S); ++u) + { + if (l2 && u + kL2Ahead < total) + { + tmaPrefetchL2(&map, 0, rowOf(u + kL2Ahead), slabOf(u + kL2Ahead)); + } + int const s = u % S; + if (u >= S) + { + if constexpr (kTimeline) + { + emptyBlocked += mbarWaitCounted(emptyBar + 8 * s, ((u / S) - 1) & 1); + } + else + { + mbarWait(emptyBar + 8 * s, ((u / S) - 1) & 1); + } + } + mbarArriveExpectTx(fullBar + 8 * s, kSlot); + if constexpr (X) + { +#pragma unroll + for (int b = 0; b < kXBoxes; ++b) + { + tmaLoad3d(ringAddr + s * kSlot + kStage + b * kXBoxBytes, &mapX, 0, + (u % kStagesPerTile) * kKs + b * kXBoxChunks, 0, fullBar + 8 * s); + } + } + if (evictFirst) + { + tmaLoad3dHint(ringAddr + s * kSlot, &map, 0, rowOf(u), slabOf(u), fullBar + 8 * s, policy); + } + else + { + tmaLoad3d(ringAddr + s * kSlot, &map, 0, rowOf(u), slabOf(u), fullBar + 8 * s); + } + if (kTimeline && u == 0) + { + tFirst = clock64() - t0; + } + } + int64_t const tLast = kTimeline ? clock64() - t0 : 0; + if ((p.flags & kFlagPrefetchFc2) && !(p.flags & kFlagPrefetchFc2Early)) + { + prefetchFc2Share(p, rs.distinct, numDistinct); + } + if (int64_t* row = timelineRow(p, 0)) + { + row[kTlProdFirst] = tFirst; + row[kTlProdLast] = tLast; + row[kTlEmptyBlocked] = emptyBlocked; + } + } + } + else + { + cpAsyncWait<0>(); + mmaWarpsSync(); // the activation rows of every MMA thread are in shared memory + int64_t const tSetup = kTimeline ? clock64() - t0 : 0; + int n = 0; // K chunks consumed by this warp (all tiles) + int fullBlocked = 0; + int64_t tFirst = 0; + int64_t tLast = 0; + // Stages are used in ring order; a tile starts on a stage boundary (kFc1Chunks % kKs == 0). The fragment + // bases point at the stage of ring slot ``fetchSlot`` (``aBase``) and at the K chunk of the activation rows + // the current stage starts with (``bBase``); both move by constant steps, so no lane offset stays live. + int fetchSlot = 0; + uint32_t fetchPhase = 0; + int releaseSlot = 0; + int fetchUses = 0; + uint32_t aBase[4]; + uint32_t bBase[2]; + uint32_t bX[2][2]; // X: K chunk parity p of the activation boxes of ring slot ``fetchSlot`` + { + LaneOffsets const lo = laneOffsets(rows, kXRowBytes, lane); + LaneOffsetsX const lx = laneOffsetsX(rows, lane); +#pragma unroll + for (int q = 0; q < 4; ++q) + { + aBase[q] = ringAddr + warp * C * kChunkBytes + lo.a[q]; + } +#pragma unroll + for (int h = 0; h < 2; ++h) + { + bBase[h] = xsAddr + lo.b[h]; + bX[0][h] = ringAddr + kStage + lx.b[0][h]; + bX[1][h] = ringAddr + kStage + lx.b[1][h]; + } + } + // Waits for the stage of ring slot ``fetchSlot`` unless an earlier probe found it complete (``ready``). + auto waitStage = [&](bool ready) + { + if constexpr (kTimeline) + { + bool const wait = !resident || fetchUses++ < S; + if (wait && !ready) + { + fullBlocked += mbarWaitCounted(fullBar + 8 * fetchSlot, fetchPhase); + } + tLast = clock64() - t0; + tFirst = tFirst == 0 ? tLast : tFirst; + } + else if (!ready) + { + mbarWait(fullBar + 8 * fetchSlot, fetchPhase); + } + }; + // Moves the A bases (X: and the activation bases) to the next ring slot. + auto nextSlot = [&]() + { + int const step = fetchSlot + 1 == S ? -(S - 1) * kSlot : kSlot; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + aBase[q] += step; + } + if constexpr (X) + { +#pragma unroll + for (int h = 0; h < 2; ++h) + { + bX[0][h] += step; + bX[1][h] += step; + } + } + if (++fetchSlot == S) + { + fetchSlot = 0; + fetchPhase ^= 1; + } + }; + if constexpr (S3 > 0) + { + if (triples && valid3 > 0) + { + // Chain c of this warp: rows 16 warp .. 16 warp + 15 of sub-tile sub0 + c, in box c of each stage. + uint32_t a3[4]; + uint32_t b3[2]; + uint32_t b3X[2][2]; // X: K chunk 0 / 1 of the stage's activation box + { + LaneOffsets const lo = laneOffsets(rows, kXRowBytes, lane); + LaneOffsetsX const lx = laneOffsetsX(rows, lane); +#pragma unroll + for (int q = 0; q < 4; ++q) + { + a3[q] = ringAddr + warp * kChunkBytes + lo.a[q]; + } +#pragma unroll + for (int h = 0; h < 2; ++h) + { + b3[h] = xsAddr + lo.b[h]; + b3X[0][h] = ringAddr + kStage3 + lx.b[0][h]; + b3X[1][h] = ringAddr + kStage3 + lx.b[1][h]; + } + } + // The fragments of K chunk j (0 / 1) of the current stage. + auto frags3 = [&](ChunkFrags& f, int j) + { + if constexpr (X) + { + loadFragsStrided(f, a3, j * kChunkStride3, b3X[j], 0, valid3); + } + else + { + loadFragsStrided(f, a3, j * kChunkStride3, b3, j * kKChunk * 2, valid3); + } + }; + float acc[kTripleSubTiles][4]; +#pragma unroll + for (int c = 0; c < kTripleSubTiles; ++c) + { + acc[c][0] = acc[c][1] = acc[c][2] = acc[c][3] = -0.f; + } + ChunkFrags f[2]; + waitStage(false); + if (!noMma) + { + frags3(f[0], 0); + } +#pragma unroll 1 + for (int st = 0; st < kStagesPerTile3; ++st) + { + // K chunk 0 of the stage: fetch chunk 1, MMAs of chunk 0. + if (!noMma) + { + frags3(f[1], 1); + mmaChunkValid(acc, f[0], valid3); + } + if (st + 1 < kStagesPerTile3) + { + // K chunk 1: move to the next stage, probe it, MMAs of chunk 1, then wait only if the probe + // failed and fetch the next stage's chunk 0. + int const step = fetchSlot + 1 == S3 ? -(S3 - 1) * kSlot3 : kSlot3; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + a3[q] += step; + } + if (++fetchSlot == S3) + { + fetchSlot = 0; + fetchPhase ^= 1; + } +#pragma unroll + for (int h = 0; h < 2; ++h) + { + b3[h] += kTripleStageChunks * kKChunk * 2; + b3X[0][h] += step; + b3X[1][h] += step; + } + bool const ready = mbarTestWait(fullBar + 8 * fetchSlot, fetchPhase); + if (!noMma) + { + mmaChunkValid(acc, f[1], valid3); + } + waitStage(ready); + if (!noMma) + { + frags3(f[0], 0); + } + } + else if (!noMma) + { + mmaChunkValid(acc, f[1], valid3); + } + // Every thread releases its own reads of the stage (their fragments fed the MMAs above). + mbarArrive(emptyBar + 8 * releaseSlot); + releaseSlot = releaseSlot + 1 == S3 ? 0 : releaseSlot + 1; + n += kTripleStageChunks; + } + int const g = lane >> 2; + int const tq = lane & 3; +#pragma unroll + for (int c = 0; c < kTripleSubTiles; ++c) + { + if (c < valid3) + { + int const sub = sub0 + c; + int const d = sub / kSubPerExpert; + int const a = (sub % kSubPerExpert) * W + warp; + int const j = 16 * (a >> 1) + 2 * g + (a & 1); +#pragma unroll + for (int h = 0; h < 2; ++h) + { + int const tok = 2 * tq + h; + if (tok < rows) + { + p.h[static_cast(d) * kMaxRows * kInter + hImageElem(tok, j)] + = bf16SatFinite(swiGlu(acc[c][h], acc[c][2 + h])); + } + } + } + } + } + } + // The regular path (no tile when the triples path ran). + for (int t = triples ? tiles : static_cast(blockIdx.x); t < tiles; t += gridDim.x) + { + float acc[C][4]; +#pragma unroll + for (int c = 0; c < C; ++c) + { + acc[c][0] = acc[c][1] = acc[c][2] = acc[c][3] = -0.f; + } + // The fragments of chunk j + 1 are loaded before the MMAs of chunk j, across stage boundaries too (the + // next stage is waited for at the last chunk of the current one); f[0] holds a stage's first chunk. + ChunkFrags f[2]; + // The fragments of K chunk j of the current stage (X: chunk j is in activation box j / 2, parity j % 2). + auto frags = [&](ChunkFrags& fr, int j) + { + if constexpr (X) + { + loadFragsAt(fr, aBase, j * kChunkSet, bX[j & 1], (j >> 1) * kXBoxBytes); + } + else + { + loadFragsAt(fr, aBase, j * kChunkSet, bBase, j * 128); + } + }; + if (t != static_cast(blockIdx.x)) + { + nextSlot(); +#pragma unroll + for (int h = 0; h < 2; ++h) + { + bBase[h] -= (kStagesPerTile - 1) * kKs * 128; + } + } + waitStage(false); + if (!noMma) + { + frags(f[0], 0); + } +#pragma unroll 1 + for (int st = 0; st < kStagesPerTile; ++st) + { +#pragma unroll + for (int j = 0; j < kKs; ++j) + { + if (j + 1 < kKs) + { + if (!noMma) + { + frags(f[(j + 1) & 1], j + 1); + mmaChunk(acc, f[j & 1]); + } + } + else if (st + 1 < kStagesPerTile) + { + // Last chunk of the stage: probe the next stage, issue this chunk's MMAs, then wait only if + // the probe failed, so the barrier round trip overlaps the MMAs. + nextSlot(); +#pragma unroll + for (int h = 0; h < 2; ++h) + { + bBase[h] += kKs * 128; + } + bool const ready = mbarTestWait(fullBar + 8 * fetchSlot, fetchPhase); + if (!noMma) + { + mmaChunk(acc, f[j & 1]); + } + waitStage(ready); + if (!noMma) + { + frags(f[0], 0); + } + } + else if (!noMma) + { + mmaChunk(acc, f[j & 1]); + } + } + // Every thread releases its own reads of the stage (their fragments fed the MMAs above). + mbarArrive(emptyBar + 8 * releaseSlot); + releaseSlot = releaseSlot + 1 == S ? 0 : releaseSlot + 1; + n += kKs; + } + int const d = t / kTilesPerExpert; + int const g = lane >> 2; + int const tq = lane & 3; +#pragma unroll + for (int c = 0; c < C; ++c) + { + int const a = (t % kTilesPerExpert) * kChains + warp * C + c; + int const j = 16 * (a >> 1) + 2 * g + (a & 1); +#pragma unroll + for (int h = 0; h < 2; ++h) + { + int const tok = 2 * tq + h; + if (tok < rows) + { + p.h[static_cast(d) * kMaxRows * kInter + hImageElem(tok, j)] + = bf16SatFinite(swiGlu(acc[c][h], acc[c][2 + h])); + } + } + } + } + if (warp == 0 && lane == 0) + { + if (int64_t* row = timelineRow(p, 0)) + { + int64_t const cycles = clock64() - t0; + row[kTlEntryNs] = entryNs; + row[kTlDoneNs] = globalTimerNs(); + row[kTlSmId] = smId(); + row[kTlWait] = tWait; + row[kTlSetup] = tSetup; + row[kTlFirst] = tFirst; + row[kTlLast] = tLast; + row[kTlDone] = cycles; + row[kTlFullBlocked] = fullBlocked; + row[kTlChunks] = n; + row[kTlRouteTokens] = tTokens != 0 ? tTokens - t0 : 0; + row[kTlRouted] = tRouted; + } + } + } + if (!early) + { + // A CTA counts as triggered as soon as one of its threads triggers: wait for the MMA warps, or the idle + // producer lanes would release the dependent grid onto free SMs right after routing. + __syncthreads(); + gridDependencyTrigger(); + } +#endif +} + +//! FC2 ring depth for W MMA warps (stages of W * 8 KB weights + one intermediate entry, about 140 KB in all). +template +constexpr int fc2TmaStages() +{ + return W == 4 ? 4 : (W == 3 ? 5 : (W == 2 ? 7 : 12)); +} + +template +__global__ void __launch_bounds__((W + 1) * 32) + mtpMoeFc2TmaKernel(DraftRoutedMoeParams const p, __grid_constant__ CUtensorMap const map) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + static_assert(E == 1 || E == 2, "one or two experts per step"); + static_assert(S >= 2 * E, "the ring must hold the step being consumed and the next one"); + // One stage per (group, expert): the expert's W * 16 rows x K = 256 weight box (one 3D TMA) and its + // intermediate entry (one bulk copy), both counted on the stage's full barrier. + constexpr int kWBytes = W * kFc2Chunks * kChunkBytes; + constexpr int kStage = kWBytes + kHEntryBytes; + constexpr int kGroups = kFc2Tiles / W; + static_assert(kWBytes % 1024 == 0 && kHEntryBytes % 1024 == 0, "stages must stay 1 KB aligned"); + extern __shared__ uint8_t smemRaw[]; + __shared__ Fc2Smem fs; + __shared__ __align__(8) uint64_t bars[2 * S]; + int64_t const t0 = kTimeline ? clock64() : 0; + int64_t const entryNs = kTimeline ? globalTimerNs() : 0; + int const warp = threadIdx.x >> 5; + int const lane = threadIdx.x & 31; + int const rows = p.rows; + uint32_t const ringAddr = (smemAddr(smemRaw) + 1023u) & ~1023u; + uint8_t* const stageBase = smemRaw + (ringAddr - smemAddr(smemRaw)) + S * kStage; + uint32_t const fullBar = smemAddr(&bars[0]); + uint32_t const emptyBar = smemAddr(&bars[S]); + if (threadIdx.x == 0) + { + for (int s = 0; s < S; ++s) + { + mbarInit(fullBar + 8 * s, 1); + mbarInit(emptyBar + 8 * s, W * 32); // every MMA thread releases its own reads + } + mbarFenceInit(); + prefetchTensorMap(&map); + } + + // The routing record and the intermediate are written by the FC1 grid. + gridDependencyWait(); + int64_t const tWait = kTimeline ? clock64() - t0 : 0; + bool const noMma = kTimeline && (p.flags & kFlagProbeNoMma) != 0; + // Publishes the barrier initialization. Past this point the producer only needs the distinct experts and the MMA + // warps only the slot and weight tables, so each role reads its part of the routing record by itself. + __syncthreads(); + int const numDistinct = p.route[0]; + + if (warp == W) + { + for (int i = lane; i < kMaxDistinct; i += 32) + { + fs.distinct[i] = p.route[kRouteDistinctOffset + i]; + } + __syncwarp(); + if (lane == 0) + { + // The intermediate was written with generic stores by the previous grid; bulk copies read it. + asm volatile("fence.proxy.async.global;" ::: "memory"); + int u = 0; + int emptyBlocked = 0; + int64_t tFirst = 0; + for (int gi = blockIdx.x; gi < kGroups; gi += gridDim.x) + { + for (int d = 0; d < numDistinct; ++d, ++u) + { + int const s = u % S; + if (u >= S) + { + if constexpr (kTimeline) + { + emptyBlocked += mbarWaitCounted(emptyBar + 8 * s, ((u / S) - 1) & 1); + } + else + { + mbarWait(emptyBar + 8 * s, ((u / S) - 1) & 1); + } + } + mbarArriveExpectTx(fullBar + 8 * s, kWBytes + rows * kInter * 2); + tmaLoad3d( + ringAddr + s * kStage, &map, 0, gi * 16 * W, fs.distinct[d] * kFc2Chunks, fullBar + 8 * s); + bulkLoad(ringAddr + s * kStage + kWBytes, p.h + static_cast(d) * kMaxRows * kInter, + rows * kInter * 2, fullBar + 8 * s); + if (kTimeline && u == 0) + { + tFirst = clock64() - t0; + } + } + } + if (int64_t* row = timelineRow(p, 1)) + { + row[kTlProdFirst] = tFirst; + row[kTlProdLast] = clock64() - t0; + row[kTlEmptyBlocked] = emptyBlocked; + } + } + } + else + { + // Slot and weight tables, shared by the W MMA warps only (named barrier 1); first read at the end of a step. + for (int i = threadIdx.x; i < kMaxRows * kMaxDistinct; i += W * 32) + { + fs.slotOf[i / kMaxDistinct][i % kMaxDistinct] = -1; + } + mmaWarpsSync(); + for (int i = threadIdx.x; i < rows * kTopK; i += W * 32) + { + int const t = i / kTopK; + int const k = i % kTopK; + fs.slotOf[t][p.route[kRouteSlotOffset + i]] = static_cast(k); + fs.weight[t][k] = __int_as_float(p.route[kRouteWeightOffset + i]); + } + mmaWarpsSync(); + int64_t const tSetup = kTimeline ? clock64() - t0 : 0; + __nv_bfloat16* stage = reinterpret_cast<__nv_bfloat16*>(stageBase + warp * kStageBytes); + int u0 = 0; // stage uses of the previous groups + int fullBlocked = 0; + int64_t tFirst = 0; + int64_t tLast = 0; + // A group's experts are taken E at a time (a step); step j holds experts E*j .. E*j + E - 1, the last step + // may hold fewer. Each expert has its own stage and accumulator; their MMAs are interleaved per k-step. + int const steps = (numDistinct + E - 1) / E; + int fetchSlot = 0; // stages are used in ring order, one per (group, expert) + uint32_t fetchPhase = 0; + int releaseSlot = 0; + // Fragment bases of the next ring slot to be fetched; they move by constant steps. + uint32_t slotA[4]; + uint32_t slotB[2]; + { + LaneOffsets const lo = laneOffsets(rows, kInter * 2, lane); +#pragma unroll + for (int q = 0; q < 4; ++q) + { + slotA[q] = ringAddr + warp * kChunkBytes + lo.a[q]; + } +#pragma unroll + for (int h = 0; h < 2; ++h) + { + slotB[h] = ringAddr + kWBytes + lo.b[h]; + } + } + uint32_t aBase[E][4]; + uint32_t bBase[E][2]; + // Points the fragment bases at the stages of step j (one per expert of the step, in ring order) and probes + // them; finishStep waits for those the probe did not find complete. A step past the last expert reads the + // step's first expert again (results discarded). + uint32_t pendBar[E]; + uint32_t pendParity[E]; + bool pendWait[E]; + auto enterStep = [&](int j) + { +#pragma unroll + for (int e = 0; e < E; ++e) + { + pendWait[e] = false; + if (E * j + e < numDistinct) + { + pendBar[e] = fullBar + 8 * fetchSlot; + pendParity[e] = fetchPhase; + pendWait[e] = !mbarTestWait(pendBar[e], pendParity[e]); +#pragma unroll + for (int q = 0; q < 4; ++q) + { + aBase[e][q] = slotA[q]; + } +#pragma unroll + for (int h = 0; h < 2; ++h) + { + bBase[e][h] = slotB[h]; + } + int const step = fetchSlot + 1 == S ? -(S - 1) * kStage : kStage; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + slotA[q] += step; + } +#pragma unroll + for (int h = 0; h < 2; ++h) + { + slotB[h] += step; + } + if (++fetchSlot == S) + { + fetchSlot = 0; + fetchPhase ^= 1; + } + } + else + { +#pragma unroll + for (int q = 0; q < 4; ++q) + { + aBase[e][q] = aBase[0][q]; + } +#pragma unroll + for (int h = 0; h < 2; ++h) + { + bBase[e][h] = bBase[0][h]; + } + } + } + }; + auto finishStep = [&]() + { +#pragma unroll + for (int e = 0; e < E; ++e) + { + if (pendWait[e]) + { + if constexpr (kTimeline) + { + ++fullBlocked; + } + mbarWait(pendBar[e], pendParity[e]); + } + } + if constexpr (kTimeline) + { + tLast = clock64() - t0; + tFirst = tFirst == 0 ? tLast : tFirst; + } + }; + // K chunk kc of every expert of the current step. + auto load = [&](ExpertFrags& f, int kc) + { +#pragma unroll + for (int e = 0; e < E; ++e) + { +#pragma unroll + for (int q = 0; q < 4; ++q) + { + ldmatrixX4(f.a[e][q], aBase[e][q] + kc * W * kChunkBytes); + } +#pragma unroll + for (int h = 0; h < 2; ++h) + { + ldmatrixX4(f.b[e][h], bBase[e][h] + kc * 128); + } + } + }; + for (int gi = blockIdx.x; gi < kGroups; gi += gridDim.x) + { + int const tile = gi * W + warp; + float acc[E][4]; +#pragma unroll + for (int e = 0; e < E; ++e) + { + acc[e][0] = acc[e][1] = acc[e][2] = acc[e][3] = -0.f; + } + // The fragments of chunk kc + 1 are loaded before the MMAs of chunk kc, across steps too (the next + // step's accumulators are independent); f[0] holds a step's first chunk. + ExpertFrags f[2]; + enterStep(0); + finishStep(); + if (!noMma) + { + load(f[0], 0); + } +#pragma unroll 1 + for (int j = 0; j < steps; ++j) + { +#pragma unroll + for (int kc = 0; kc < kFc2Chunks; ++kc) + { + auto mmas = [&]() + { + if (!noMma) + { +#pragma unroll + for (int q = 0; q < 4; ++q) + { +#pragma unroll + for (int e = 0; e < E; ++e) + { + ExpertFrags const& fr = f[kc & 1]; + mmaBf16(acc[e], fr.a[e][q], fr.b[e][q >> 1][2 * (q & 1)], + fr.b[e][q >> 1][2 * (q & 1) + 1]); + } + } + } + }; + if (kc + 1 < kFc2Chunks) + { + if (!noMma) + { + load(f[(kc + 1) & 1], kc + 1); + } + mmas(); + } + else if (j + 1 < steps) + { + // Last chunk of the step: probe the next step's stages, issue this chunk's MMAs, then wait + // only for stages the probe did not find complete. + enterStep(j + 1); + mmas(); + finishStep(); + if (!noMma) + { + load(f[0], 0); + } + } + else + { + mmas(); + } + } + // Every thread releases its own reads of the step's stages (their fragments fed the MMAs above). +#pragma unroll + for (int e = 0; e < E; ++e) + { + if (E * j + e < numDistinct) + { + mbarArrive(emptyBar + 8 * releaseSlot); + releaseSlot = releaseSlot + 1 == S ? 0 : releaseSlot + 1; + } + } + // FC2 rows of each expert -> BF16, into the token's slot of the warp's staging area. + int const g = lane >> 2; + int const tq = lane & 3; +#pragma unroll + for (int e = 0; e < E; ++e) + { + int const d = E * j + e; +#pragma unroll + for (int c = 0; c < 2; ++c) + { + int const tok = 2 * tq + c; + int const k = (d < numDistinct && tok < rows) ? fs.slotOf[tok][d] : -1; + if (k >= 0) + { + __nv_bfloat16 const lo16 = bf16SatFinite(acc[e][c]); + __nv_bfloat16 const hi16 = bf16SatFinite(acc[e][2 + c]); + stage[(tok * kTopK + k) * 16 + g] = lo16; + stage[(tok * kTopK + k) * 16 + g + 8] = hi16; + if (p.debugFc2 != nullptr) + { + int64_t const base = (static_cast(tok) * kTopK + k) * kHidden; + p.debugFc2[base + hiddenOfStoredRow(tile * 16 + g)] = lo16; + p.debugFc2[base + hiddenOfStoredRow(tile * 16 + g + 8)] = hi16; + } + } + } + acc[e][0] = acc[e][1] = acc[e][2] = acc[e][3] = -0.f; + } + } + u0 += numDistinct; + // Native finalize of the tile's 16 hidden rows for every token: FP32 fma chain from 0 in slot order. + __syncwarp(); + for (int pr = lane; pr < rows * 16; pr += 32) + { + int const tok = pr >> 4; + int const row = pr & 15; + float sum = 0.f; +#pragma unroll + for (int k = 0; k < kTopK; ++k) + { + sum = fmaFtz(fs.weight[tok][k], __bfloat162float(stage[(tok * kTopK + k) * 16 + row]), sum); + } + p.out[static_cast(tok) * kHidden + hiddenOfStoredRow(tile * 16 + row)] + = __float2bfloat16_rn(sum); + } + __syncwarp(); + } + if (warp == 0 && lane == 0) + { + if (int64_t* row = timelineRow(p, 1)) + { + int64_t const cycles = clock64() - t0; + row[kTlEntryNs] = entryNs; + row[kTlDoneNs] = globalTimerNs(); + row[kTlSmId] = smId(); + row[kTlWait] = tWait; + row[kTlSetup] = tSetup; + row[kTlFirst] = tFirst; + row[kTlLast] = tLast; + row[kTlDone] = cycles; + row[kTlFullBlocked] = fullBlocked; + row[kTlChunks] = u0 * kFc2Chunks; + } + } + } +#endif +} + +template +cudaError_t launchOne(Kernel kernel, int ctas, int threads, int smemBytes, bool pdl, DraftRoutedMoeParams const& params, + cudaStream_t stream) +{ + cudaError_t status = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smemBytes); + if (status != cudaSuccess) + { + return status; + } + cudaLaunchConfig_t config{}; + config.gridDim = dim3(ctas); + config.blockDim = dim3(threads); + config.dynamicSmemBytes = smemBytes; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + config.attrs = attrs; + config.numAttrs = 1; + return cudaLaunchKernelEx(&config, kernel, params); +} + +template +cudaError_t launchTma(Kernel kernel, int ctas, int threads, int smemBytes, bool pdl, DraftRoutedMoeParams const& params, + CUtensorMap const& map, cudaStream_t stream) +{ + cudaError_t status = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smemBytes); + if (status != cudaSuccess) + { + return status; + } + cudaLaunchConfig_t config{}; + config.gridDim = dim3(ctas); + config.blockDim = dim3(threads); + config.dynamicSmemBytes = smemBytes; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + config.attrs = attrs; + config.numAttrs = 1; + return cudaLaunchKernelEx(&config, kernel, params, map); +} + +template +cudaError_t launchTma3(Kernel kernel, int ctas, int threads, int smemBytes, bool pdl, + DraftRoutedMoeParams const& params, CUtensorMap const& map, CUtensorMap const& map3, CUtensorMap const& mapX, + cudaStream_t stream) +{ + cudaError_t status = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smemBytes); + if (status != cudaSuccess) + { + return status; + } + cudaLaunchConfig_t config{}; + config.gridDim = dim3(ctas); + config.blockDim = dim3(threads); + config.dynamicSmemBytes = smemBytes; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + config.attrs = attrs; + config.numAttrs = 1; + return cudaLaunchKernelEx(&config, kernel, params, map, map3, mapX); +} + +template +cudaError_t launchTma2(Kernel kernel, int ctas, int threads, int smemBytes, bool pdl, + DraftRoutedMoeParams const& params, CUtensorMap const& map, CUtensorMap const& map3, cudaStream_t stream) +{ + return launchTma3(kernel, ctas, threads, smemBytes, pdl, params, map, map3, map, stream); +} + +} // namespace + +int fc1SmemBytes(int rows, int warps, bool deep) +{ + return rows * kXRowBytes + warps * (deep ? kFc1StagesDeep : kFc1Stages) * kChunkBytes; +} + +int fc2SmemBytes(int warps, bool deep) +{ + return warps * ((deep ? kFc2StagesDeep : kFc2Stages) * kChunkBytes + kFc2HStages * kHEntryBytes + kStageBytes); +} + +namespace +{ + +template +cudaError_t launchFc1(bool pairs, int smem, bool pdl, DraftRoutedMoeParams const& params, cudaStream_t stream) +{ + return pairs ? launchOne(mtpMoeFc1Kernel, params.fc1Ctas, 32 * W, smem, pdl, params, stream) + : launchOne(mtpMoeFc1Kernel, params.fc1Ctas, 32 * W, smem, pdl, params, stream); +} + +template +cudaError_t launchFc2(bool pairs, int smem, bool pdl, DraftRoutedMoeParams const& params, cudaStream_t stream) +{ + return pairs ? launchOne(mtpMoeFc2Kernel, params.fc2Ctas, 32 * W, smem, pdl, params, stream) + : launchOne(mtpMoeFc2Kernel, params.fc2Ctas, 32 * W, smem, pdl, params, stream); +} + +template +cudaError_t launchFc1Warps(bool pairs, int smem, bool pdl, DraftRoutedMoeParams const& params, cudaStream_t stream) +{ + switch (params.fc1Warps) + { + case 1: return launchFc1<1, S>(pairs, smem, pdl, params, stream); + case 2: return launchFc1<2, S>(pairs, smem, pdl, params, stream); + case 4: return launchFc1<4, S>(pairs, smem, pdl, params, stream); + default: return cudaErrorInvalidValue; + } +} + +template +cudaError_t launchFc2Warps(bool pairs, int smem, bool pdl, DraftRoutedMoeParams const& params, cudaStream_t stream) +{ + switch (params.fc2Warps) + { + case 1: return launchFc2<1, S>(pairs, smem, pdl, params, stream); + case 2: return launchFc2<2, S>(pairs, smem, pdl, params, stream); + case 4: return launchFc2<4, S>(pairs, smem, pdl, params, stream); + default: return cudaErrorInvalidValue; + } +} + +} // namespace + +int fc1TmaSmemBytes(int rows, int stages) +{ + return 1024 + stages * kTmaStageChunks * kChunkBytes + rows * kXRowBytes; +} + +int fc1TmaTriplesSmemBytes(int rows, int stages, int tripleStages) +{ + int const ring = stages * kTmaStageChunks * kChunkBytes; + int const ring3 = tripleStages * kTripleSubTiles * kTripleStageChunks * kTripleRows * kKChunk * 2; + return 1024 + (ring > ring3 ? ring : ring3) + rows * kXRowBytes; +} + +int fc1TmaStreamXSmemBytes() +{ + int const ring = 6 * (kTmaStageChunks * kChunkBytes + kXBoxBytes); + int const ring3 = 4 * (kTripleSubTiles * kTripleRows * kTripleStageChunks * kKChunk * 2 + kXBoxBytes); + return 1024 + (ring > ring3 ? ring : ring3); +} + +int fc2TmaSmemBytes(int warps, int stages) +{ + return 1024 + stages * (warps * kFc2Chunks * kChunkBytes + kHEntryBytes) + warps * kStageBytes; +} + +namespace +{ + +//! Dynamic shared memory budget (the static routing state and barriers take about 3.5 KB of the 227 KB). +constexpr int kMaxDynamicSmem = 220 * 1024; + +//! FC1 kernel with ring depth S (and S3 for the triples path, 4 MMA warps only; 0: no triples path). The instantiated +//! (S, S3) pairs are those launchFc1TmaC chooses: (6, 4), (4, 3), (4, 2), (3, 2). +template +cudaError_t launchFc1TmaCS(int smem, int tripleStages, bool pdl, DraftRoutedMoeParams const& params, + DraftRoutedMoeTensorMaps const& maps, cudaStream_t stream) +{ + int const n = params.fc1Ctas; + CUtensorMap const& w = maps.w13; + switch (params.fc1Warps) + { + case 1: return launchTma2(mtpMoeFc1TmaKernel<1, C, S, 0>, n, 64, smem, pdl, params, w, w, stream); + case 2: return launchTma2(mtpMoeFc1TmaKernel<2, C, S, 0>, n, 96, smem, pdl, params, w, w, stream); + case 4: + { + CUtensorMap const& w3 = maps.w13Triple; + if (tripleStages == 0) + { + return launchTma2(mtpMoeFc1TmaKernel<4, C, S, 0>, n, 160, smem, pdl, params, w, w, stream); + } + if constexpr (S == 6) + { + if (tripleStages == 4) + { + return launchTma2(mtpMoeFc1TmaKernel<4, C, 6, 4>, n, 160, smem, pdl, params, w, w3, stream); + } + } + if constexpr (S == 4) + { + if (tripleStages == 3) + { + return launchTma2(mtpMoeFc1TmaKernel<4, C, 4, 3>, n, 160, smem, pdl, params, w, w3, stream); + } + if (tripleStages == 2) + { + return launchTma2(mtpMoeFc1TmaKernel<4, C, 4, 2>, n, 160, smem, pdl, params, w, w3, stream); + } + } + if constexpr (S == 3) + { + if (tripleStages == 2) + { + return launchTma2(mtpMoeFc1TmaKernel<4, C, 3, 2>, n, 160, smem, pdl, params, w, w3, stream); + } + } + return cudaErrorInvalidValue; + } + default: return cudaErrorInvalidValue; + } +} + +//! FC1 ring depth: the deepest of 6 / 4 / 3 stages of 32 KB that fits next to the activation rows; with +//! kFlagFc1Triples and 4 MMA warps also the deepest triples ring (48 KB stages) of the instantiated pairs that fits. +//! kFlagStreamX (4 MMA warps, tile pairs; ignored otherwise): 6 stages of 32 + 2 KB and 4 triples stages of 48 + 2 KB +//! for any row count. +template +cudaError_t launchFc1TmaC( + bool pdl, DraftRoutedMoeParams const& params, DraftRoutedMoeTensorMaps const& maps, cudaStream_t stream) +{ + if constexpr (C == 2) + { + if ((params.flags & kFlagStreamX) != 0 && params.fc1Warps == 4) + { + int const smem = fc1TmaStreamXSmemBytes(); + return launchTma3(mtpMoeFc1TmaKernel<4, 2, 6, 4, true>, params.fc1Ctas, 160, smem, pdl, params, maps.w13, + maps.w13Triple, maps.x, stream); + } + } + bool const triples = (params.flags & kFlagFc1Triples) != 0 && params.fc1Warps == 4; + for (int stages : {6, 4, 3}) + { + int smem = fc1TmaSmemBytes(params.rows, stages); + if (smem > kMaxDynamicSmem) + { + continue; + } + int tripleStages = 0; + if (triples) + { + for (int t : {4, 3, 2}) + { + bool const instantiated + = (stages == 6 && t == 4) || (stages == 4 && (t == 3 || t == 2)) || (stages == 3 && t == 2); + if (instantiated && fc1TmaTriplesSmemBytes(params.rows, stages, t) <= kMaxDynamicSmem) + { + tripleStages = t; + smem = fc1TmaTriplesSmemBytes(params.rows, stages, t); + break; + } + } + if (tripleStages == 0) + { + return cudaErrorInvalidValue; + } + } + switch (stages) + { + case 6: return launchFc1TmaCS(smem, tripleStages, pdl, params, maps, stream); + case 4: return launchFc1TmaCS(smem, tripleStages, pdl, params, maps, stream); + default: return launchFc1TmaCS(smem, tripleStages, pdl, params, maps, stream); + } + } + return cudaErrorInvalidValue; +} + +template +cudaError_t launchFc2TmaW(bool pdl, DraftRoutedMoeParams const& params, CUtensorMap const& map, cudaStream_t stream) +{ + constexpr int kS = fc2TmaStages(); + int const smem = fc2TmaSmemBytes(W, kS); + return (params.flags & kFlagFc2Pairs) + ? launchTma(mtpMoeFc2TmaKernel, params.fc2Ctas, (W + 1) * 32, smem, pdl, params, map, stream) + : launchTma(mtpMoeFc2TmaKernel, params.fc2Ctas, (W + 1) * 32, smem, pdl, params, map, stream); +} + +cudaError_t launchTmaPair(DraftRoutedMoeParams const& params, DraftRoutedMoeTensorMaps const& maps, cudaStream_t stream) +{ + bool const pdl = (params.flags & kFlagPdl) != 0; + cudaError_t status = (params.flags & kFlagFc1Pairs) ? launchFc1TmaC<2>(pdl, params, maps, stream) + : launchFc1TmaC<1>(pdl, params, maps, stream); + if (status != cudaSuccess) + { + return status; + } + switch (params.fc2Warps) + { + case 1: return launchFc2TmaW<1>(pdl, params, maps.w2, stream); + case 2: return launchFc2TmaW<2>(pdl, params, maps.w2, stream); + case 3: return launchFc2TmaW<3>(pdl, params, maps.w2, stream); + case 4: return launchFc2TmaW<4>(pdl, params, maps.w2, stream); + default: return cudaErrorInvalidValue; + } +} + +} // namespace + +cudaError_t invokeDraftRoutedMoe( + DraftRoutedMoeParams const& params, DraftRoutedMoeTensorMaps const* maps, cudaStream_t stream) +{ + if (params.rows < 1 || params.rows > kMaxRows || params.fc1Ctas < 1 || params.fc2Ctas < 1) + { + return cudaErrorInvalidValue; + } + if (params.flags & kFlagTma) + { + return maps != nullptr ? launchTmaPair(params, *maps, stream) : cudaErrorInvalidValue; + } + bool const pdl = (params.flags & kFlagPdl) != 0; + bool const pairs1 = (params.flags & kFlagFc1Pairs) != 0; + bool const pairs2 = (params.flags & kFlagFc2Pairs) != 0; + bool const deep = (params.flags & kFlagDeepRing) != 0; + int const smem1 = fc1SmemBytes(params.rows, params.fc1Warps, deep); + cudaError_t status = deep ? launchFc1Warps(pairs1, smem1, pdl, params, stream) + : launchFc1Warps(pairs1, smem1, pdl, params, stream); + if (status != cudaSuccess) + { + return status; + } + int const smem2 = fc2SmemBytes(params.fc2Warps, deep); + return deep ? launchFc2Warps(pairs2, smem2, pdl, params, stream) + : launchFc2Warps(pairs2, smem2, pdl, params, stream); +} + +} // namespace kernels::glm_moe_dsa::draft_routed_moe + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/draftRoutedMoe.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/draftRoutedMoe.h new file mode 100644 index 000000000000..19d8adada74b --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/draftRoutedMoe.h @@ -0,0 +1,205 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa::draft_routed_moe +{ + +//! The routed experts of the GLM-5.2 MTP draft layer (BF16 weights, TP8/EP1 rank shard) on the +//! FlashInfer TRT-LLM-Gen BF16 layouts, bit-identical to ``trtllm_bf16_moe`` (DeepSeek-V3 routing with +//! ``n_group <= 1``, SwiGLU, ``do_finalize=True``). +//! +//! Two kernels, chained with programmatic dependent launch: +//! +//! * ``fc1``: every CTA computes the native top-8 routing of all rows from the FP32 router logits (sigmoid via +//! ``tanh.approx``, selection key ``sigmoid + bias``, ties to the lower expert index, NaN keys first, weights +//! ``bf16(sigmoid * scale / (sum + 1e-20))`` with the native FP32 sum order and approximate division), then its +//! warps stream 16-row tiles of the distinct experts' gate/up weights (BlockMajorK, rows shuffled for the gated +//! epilogue) and run ``mma.sync.m16n8k16`` over K = 6144 in natural order with one FP32 accumulator per tile, +//! followed by the native SwiGLU epilogue. The BF16 intermediate goes to a scratch buffer. +//! * ``fc2``: after the FC1 grid has completed, each warp computes 16 down-projection rows for every distinct +//! expert (K = 256 in natural order), rounds each row to BF16 like the native FC2 epilogue, and finalizes the +//! rows of every token in top-k slot order (``fma`` from 0, one BF16 rounding), as ``finalizeKernel`` does. +//! +//! The weights are read in place: ``w13`` is ``[256][96][512][64]`` and ``w2`` is ``[256][4][6144][64]`` (BF16). + +constexpr int kHidden = 6144; +constexpr int kInter = 256; //!< Per-rank intermediate size (TP8). +constexpr int kFc1Rows = 512; //!< Gate and up rows per expert per rank. +constexpr int kNumExperts = 256; +constexpr int kTopK = 8; +constexpr int kMaxRows = 8; +constexpr int kMaxDistinct = kMaxRows * kTopK; +constexpr int kKChunk = 64; //!< BlockMajorK block (elements, 128 bytes). +constexpr int kFc1Chunks = kHidden / kKChunk; +constexpr int kFc2Chunks = kInter / kKChunk; +constexpr int kFc1Tiles = kFc1Rows / 16; +constexpr int kFc2Tiles = kHidden / 16; + +//! Routing record written by the FC1 grid and read by the FC2 grid (int32 words). +constexpr int kRouteDistinctOffset = 1; //!< [kMaxDistinct] expert ids, ascending. +constexpr int kRouteSlotOffset = kRouteDistinctOffset + kMaxDistinct; //!< [kMaxRows][kTopK] distinct index. +constexpr int kRouteWeightOffset = kRouteSlotOffset + kMaxRows * kTopK; //!< [kMaxRows][kTopK] float bits of bf16. +constexpr int kRouteWords = kRouteWeightOffset + kMaxRows * kTopK; +//! Intermediate scratch: [kMaxDistinct][kMaxRows][kInter] BF16. +constexpr int64_t kScratchHElems = int64_t{kMaxDistinct} * kMaxRows * kInter; + +//! Optional behaviour (``flags``). +enum DraftRoutedMoeFlags : int +{ + kFlagPdl = 1, //!< Programmatic dependent launch for both kernels (as the native kernels). + kFlagPrefetchFc2 = 2, //!< FC1 warps prefetch the distinct experts' FC2 weights into L2 after their tiles. + kFlagEarlyTrigger = 4, //!< FC1 triggers its dependents after routing instead of after its tiles. + kFlagFc1Pairs = 8, //!< FC1 warps interleave two tiles (two independent accumulation chains). + kFlagFc2Pairs = 16, //!< FC2 warps interleave two experts (two independent accumulation chains). + kFlagDeepRing = 32, //!< Deeper weight rings (FC1 20, FC2 12 chunks of 2 KB per warp instead of 12 and 8). + kFlagTma = 64, //!< TMA-fed kernels: one producer thread streams 16W-row weight boxes into a ring that W + //!< MMA warps consume (same per-chain arithmetic). Needs the tensor maps of invokeDraftRoutedMoe. + //!< With it, kFlagFc1Pairs gives every FC1 MMA warp two 16-row chains (boxes of 32W + //!< rows), kFlagFc2Pairs gives every FC2 MMA warp two experts per step (two stages, two + //!< accumulators), kFlagPrefetchFc2 is issued by the FC1 producer after its last weight + //!< copy, and kFlagDeepRing does not apply (the rings are sized to the shared memory). + kFlagL2Prefetch = 128, //!< TMA FC1: the producer also prefetches its weight boxes into L2, 2 x ring depth ahead. + kFlagPrefetchFc2Early = 256, //!< TMA FC1: the producer prefetches its share of the FC2 weights right after routing. + //! GLM_MOE_DSA_DRAFT_ROUTED_MOE_TIMELINE builds only: MMA warps wait for and release stages without reading them + //! (memory schedule alone; the output is meaningless). + kFlagProbeNoMma = 512, + //! GLM_MOE_DSA_DRAFT_ROUTED_MOE_TIMELINE builds only: the FC1 producer fills the ring once and the MMA warps cycle + //! over the resident stages without waiting (MMA warps alone; meaningless output). + kFlagProbeResident = 1024, + //! TMA FC1 with 4 MMA warps: when the call's tiles (decided on the device from the number of distinct experts) + //! would need more than one round of the grid, every CTA instead takes three consecutive 64-row sub-tiles (three + //! chains per MMA warp, three TMA boxes per stage), if those fit in one round. Needs + //! DraftRoutedMoeTensorMaps::w13Triple. + kFlagFc1Triples = 4096, + //! TMA FC1: the early trigger of kFlagEarlyTrigger for calls on the kFlagFc1Triples path only (other calls trigger + //! after their tiles). + kFlagEarlyTriggerTriples = 8192, + //! TMA FC1: the weight boxes are loaded with an L2 evict-first policy (each is read once per call), so that they + //! do not displace FC2 weights prefetched into L2. + kFlagFc1EvictFirst = 32768, + //! TMA FC1 with 4 MMA warps and tile pairs (ignored otherwise): the activation rows are not kept in shared + //! memory; every ring stage also carries the activation K chunks it needs (boxes of DraftRoutedMoeTensorMaps::x), + //! so the weight ring is 6 stages deep (4 on the triples path) for every row count. Needs + //! DraftRoutedMoeTensorMaps::x. + kFlagStreamX = 65536 +}; + +//! Event record of the probe builds (compiled with ``GLM_MOE_DSA_DRAFT_ROUTED_MOE_TIMELINE=1``): int64 +//! ``[2][kTimelineCtas][kTimelineSlots]`` (FC1, FC2), one row per CTA; see TimelineSlot in draftRoutedMoe.cu. +constexpr int kTimelineCtas = 1024; +constexpr int kTimelineSlots = 16; + +//! K chunks per FC1 stage of the TMA-fed kernel for all its 16-row chains together (32 KB per stage and box). +constexpr int kTmaStageChunks = 16; + +//! Sub-tiles per CTA, rows per sub-tile and K chunks per stage of the kFlagFc1Triples path (4 MMA warps): a stage is +//! three boxes of 64 rows x 2 K chunks (16 KB each). +constexpr int kTripleSubTiles = 3; +constexpr int kTripleRows = 64; +constexpr int kTripleStageChunks = 2; + +//! Activation box of kFlagStreamX: 8 tokens (rows beyond the call's are zero-filled) x kXBoxChunks K chunks. +constexpr int kXBoxChunks = 2; +constexpr int kXBoxBytes = 8 * kXBoxChunks * kKChunk * 2; + +//! Tensor maps of the TMA-fed kernels (kFlagTma): 3D views [slab][row][64] BF16 of ``w13`` (256 * 96 slabs of 512 +//! rows) and ``w2`` (256 * 4 slabs of 6144 rows), 128-byte swizzle; boxes of 64 x fc1TmaBoxRows x fc1TmaBoxSlabs +//! (FC1) and 64 x 16 * fc2Warps x 4 (FC2). +struct DraftRoutedMoeTensorMaps +{ + CUtensorMap w13; + CUtensorMap w2; + CUtensorMap w13Triple; //!< kFlagFc1Triples: view of ``w13`` with boxes of 64 x kTripleRows x kTripleStageChunks. + //! kFlagStreamX: 3D view [token][K chunk][64] of the call's ``x`` (dims 64, 96, rows), boxes of 64 x kXBoxChunks x + //! 8, 128-byte swizzle. + CUtensorMap x; +}; + +struct DraftRoutedMoeParams +{ + float const* logits; //!< [rows, 256] FP32. + __nv_bfloat16 const* bias; //!< [256] router bias. + __nv_bfloat16 const* x; //!< [rows, 6144] activation (16-byte aligned). + __nv_bfloat16 const* w13; //!< [256][96][512][64]. + __nv_bfloat16 const* w2; //!< [256][4][6144][64]. + __nv_bfloat16* out; //!< [rows, 6144] finalized routed output. + __nv_bfloat16* h; //!< Scratch, kScratchHElems. + int32_t* route; //!< Scratch, kRouteWords. + __nv_bfloat16* debugFc2; //!< Optional [rows, 8, 6144]: FC2 rows per (token, slot) in hidden order. + int64_t* timeline; //!< Optional event record (GLM_MOE_DSA_DRAFT_ROUTED_MOE_TIMELINE builds; ignored otherwise). + int rows; + float routedScaling; + int fc1Ctas; + int fc1Warps; //!< 1, 2 or 4 warps per FC1 CTA. + int fc2Ctas; + int fc2Warps; //!< 1, 2 or 4 warps per FC2 CTA (also 3 with kFlagTma: 128 groups of 48 rows). + int flags; +}; + +//! Dynamic shared memory of the FC1 kernel for ``rows`` and ``warps`` (bytes); ``deep``: kFlagDeepRing. +int fc1SmemBytes(int rows, int warps, bool deep); + +//! Dynamic shared memory of the FC2 kernel for ``warps`` (bytes); ``deep``: kFlagDeepRing. +int fc2SmemBytes(int warps, bool deep); + +//! Dynamic shared memory of the TMA-fed FC1 kernel (bytes, incl. 1 KB alignment slack); stages: ring depth. +int fc1TmaSmemBytes(int rows, int stages); + +//! Dynamic shared memory of the TMA-fed FC1 kernel with the kFlagFc1Triples path (bytes): the larger of the two +//! rings (``stages`` of 32 KB, ``tripleStages`` of three sub-tile boxes) next to the activation rows. +int fc1TmaTriplesSmemBytes(int rows, int stages, int tripleStages); + +//! Dynamic shared memory of the TMA-fed FC1 kernel with kFlagStreamX (bytes): the larger of the 6-stage regular and +//! the 4-stage triples rings, activation boxes included. +int fc1TmaStreamXSmemBytes(); + +//! Dynamic shared memory of the TMA-fed FC2 kernel (bytes, incl. 1 KB alignment slack); stages: ring depth. +int fc2TmaSmemBytes(int warps, int stages); + +//! Launches the FC1 and FC2 kernels on ``stream``; ``fc1Warps``/``fc2Warps`` must be 1, 2 or 4 (``fc2Warps`` also 3 +//! with kFlagTma). ``maps`` is required +//! with kFlagTma (see DraftRoutedMoeTensorMaps) and ignored otherwise. +cudaError_t invokeDraftRoutedMoe( + DraftRoutedMoeParams const& params, DraftRoutedMoeTensorMaps const* maps, cudaStream_t stream); + +//! Rows of the FC1 weight box of the TMA-fed kernel (16 per chain). +inline int fc1TmaBoxRows(int flags, int fc1Warps) +{ + return 16 * fc1Warps * ((flags & kFlagFc1Pairs) ? 2 : 1); +} + +//! K chunks (slabs) of the FC1 weight box of the TMA-fed kernel. +inline int fc1TmaBoxSlabs(int flags, int fc1Warps) +{ + return kTmaStageChunks * 16 / fc1TmaBoxRows(flags, fc1Warps); +} + +} // namespace kernels::glm_moe_dsa::draft_routed_moe + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/draftTail.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/draftTail.cu new file mode 100644 index 000000000000..a22ae0d29395 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/draftTail.cu @@ -0,0 +1,1568 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/kernels/glmMoeDsa/draftTail.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa::draft_tail +{ + +namespace +{ + +// --------------------------------------------------------------------------- +// Candidate keys +// --------------------------------------------------------------------------- + +// Monotone map of BF16 bits to 16 bits: -0 and +0 compare equal, NaN of either +// sign ranks above +inf (as the native torch.max reduction does). +__device__ __forceinline__ uint32_t orderedBf16(uint32_t bits) +{ + uint32_t const mag = bits & 0x7FFFu; + if (mag > 0x7F80u) + { + return 0xFFFFu; + } + if (mag == 0u) + { + bits = 0u; + } + return (bits & 0x8000u) ? (~bits & 0xFFFFu) : (bits | 0x8000u); +} + +// Larger key == larger value, then smaller global id. Never zero. +__device__ __forceinline__ uint64_t makeKey(uint32_t bf16Bits, int64_t globalId) +{ + return (static_cast(orderedBf16(bf16Bits)) << 32) + | static_cast(0xFFFFFFFFu - static_cast(globalId)); +} + +__device__ __forceinline__ int32_t keyToToken(uint64_t key) +{ + return static_cast(0xFFFFFFFFu - static_cast(key & 0xFFFFFFFFull)); +} + +__device__ __forceinline__ uint64_t umax64(uint64_t a, uint64_t b) +{ + return a > b ? a : b; +} + +__device__ __forceinline__ uint64_t warpMax64(uint64_t v) +{ +#pragma unroll + for (int off = 16; off > 0; off >>= 1) + { + v = umax64(v, __shfl_xor_sync(0xFFFFFFFFu, v, off)); + } + return v; +} + +// --------------------------------------------------------------------------- +// Memory helpers +// --------------------------------------------------------------------------- + +__device__ __forceinline__ void griddepWait() +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif +} + +__device__ __forceinline__ uint64_t ldRelaxedSys(uint64_t const* addr) +{ + uint64_t v; + asm volatile("ld.relaxed.sys.global.u64 %0, [%1];" : "=l"(v) : "l"(addr) : "memory"); + return v; +} + +__device__ __forceinline__ void stRelaxedSys(uint64_t* addr, uint64_t v) +{ + asm volatile("st.relaxed.sys.global.u64 [%0], %1;" ::"l"(addr), "l"(v) : "memory"); +} + +__device__ __forceinline__ void mcStoreRelaxedSys(uint64_t* mcAddr, uint64_t v) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("multimem.st.relaxed.sys.global.b64 [%0], %1;" ::"l"(mcAddr), "l"(v) : "memory"); +#endif +} + +__device__ __forceinline__ uint64_t evictFirstPolicy() +{ + uint64_t policy; + asm volatile("createpolicy.fractional.L2::evict_first.b64 %0, 1.0;" : "=l"(policy)); + return policy; +} + +__device__ __forceinline__ uint64_t evictLastPolicy() +{ + uint64_t policy; + asm volatile("createpolicy.fractional.L2::evict_last.b64 %0, 1.0;" : "=l"(policy)); + return policy; +} + +// One 32-byte chunk = 16 BF16 values. +struct alignas(16) Chunk +{ + uint64_t v[4]; +}; + +enum LoadKind : int +{ + kLoad32Hint = 0, // 32-byte load, L2 evict_first policy + kLoad32 = 1, // 32-byte load + kLoad16x2Pf = 2, // two 16-byte loads with 256-byte L2 prefetch +}; + +template +__device__ __forceinline__ void loadChunk(Chunk& c, void const* ptr, uint64_t policy) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 1000) + // 32-byte global loads need SM100+. + __trap(); +#else + if constexpr (kKind == kLoad32Hint) + { + asm volatile("ld.global.nc.L1::no_allocate.L2::cache_hint.v4.b64 {%0, %1, %2, %3}, [%4], %5;" + : "=l"(c.v[0]), "=l"(c.v[1]), "=l"(c.v[2]), "=l"(c.v[3]) + : "l"(ptr), "l"(policy)); + } + else if constexpr (kKind == kLoad32) + { + asm volatile("ld.global.nc.L1::no_allocate.v4.b64 {%0, %1, %2, %3}, [%4];" + : "=l"(c.v[0]), "=l"(c.v[1]), "=l"(c.v[2]), "=l"(c.v[3]) + : "l"(ptr)); + } + else + { + int4 a; + int4 b; + asm volatile("ld.global.nc.L1::no_allocate.L2::256B.v4.s32 {%0, %1, %2, %3}, [%4];" + : "=r"(a.x), "=r"(a.y), "=r"(a.z), "=r"(a.w) + : "l"(ptr)); + asm volatile("ld.global.nc.L1::no_allocate.L2::256B.v4.s32 {%0, %1, %2, %3}, [%4];" + : "=r"(b.x), "=r"(b.y), "=r"(b.z), "=r"(b.w) + : "l"(static_cast(ptr) + 16)); + c.v[0] = (static_cast(static_cast(a.y)) << 32) | static_cast(a.x); + c.v[1] = (static_cast(static_cast(a.w)) << 32) | static_cast(a.z); + c.v[2] = (static_cast(static_cast(b.y)) << 32) | static_cast(b.x); + c.v[3] = (static_cast(static_cast(b.w)) << 32) | static_cast(b.z); + } +#endif +} + +// acc += dot(w, h) over 16 BF16 pairs, in a fixed order. +__device__ __forceinline__ float dotChunk(Chunk const& w, Chunk const& h, float acc) +{ +#pragma unroll + for (int i = 0; i < 4; ++i) + { + uint32_t const wl = static_cast(w.v[i]); + uint32_t const wh = static_cast(w.v[i] >> 32); + uint32_t const hl = static_cast(h.v[i]); + uint32_t const hh = static_cast(h.v[i] >> 32); + acc = fmaf(__uint_as_float(wl << 16), __uint_as_float(hl << 16), acc); + acc = fmaf(__uint_as_float(wl & 0xFFFF0000u), __uint_as_float(hl & 0xFFFF0000u), acc); + acc = fmaf(__uint_as_float(wh << 16), __uint_as_float(hh << 16), acc); + acc = fmaf(__uint_as_float(wh & 0xFFFF0000u), __uint_as_float(hh & 0xFFFF0000u), acc); + } + return acc; +} + +constexpr int kHidden6144ChunksDev = 6144 / 512; // 32-byte chunks per lane per row + +__device__ __forceinline__ uint32_t* arrivalCounter(uint64_t* scratch) +{ + return reinterpret_cast(scratch + kMaxRows); +} + +__device__ __forceinline__ uint32_t* epochCounter(uint64_t* scratch) +{ + return reinterpret_cast(scratch + kMaxRows) + 1; +} + +// --------------------------------------------------------------------------- +// Cross-rank exchange (one warp) +// +// Slot s of the exchange buffer holds kMaxRanks x kMaxRows keys; rank q writes +// only its own row block, every rank polls its local copy until all tpSize +// entries of each row are non-zero. Epoch e uses slot e % 3. Before publishing +// epoch e a rank clears slot (e + 1) % 3, whose last use was epoch e - 2: every +// key of that epoch was already observed here, and no peer can publish epoch +// e + 1 before it has observed this rank's epoch-e key, which the release fence +// orders after the clear (the peer's own fence before its next publish is the +// matching acquire). All ranks run the same sequence of epochs. +// --------------------------------------------------------------------------- +__device__ void exchangeAndWrite( + DraftTailParams const& p, uint64_t const* localKeys /* smem [rows] */, uint64_t (*gathered)[kMaxRows] /* smem */) +{ + int const lane = threadIdx.x & 31; + int const rows = p.rows; + bool const exchange = p.mode != ExchangeMode::kNone && p.tpSize > 1; + uint64_t result = 0; + if (!exchange) + { + if (lane < rows) + { + result = localKeys[lane]; + } + } + else + { + uint32_t* epochPtr = epochCounter(p.scratch); + uint32_t const epoch = *reinterpret_cast(epochPtr); + int const cur = static_cast(epoch % kNumSlots); + int const nxt = static_cast((epoch + 1) % kNumSlots); + for (int i = lane; i < kSlotWords; i += 32) + { + p.ucSlotsLocal[nxt * kSlotWords + i] = 0ull; + } + // Release: the clear above is ordered before this rank's publish below; + // a peer's next write to the cleared slot is causally after that publish. + asm volatile("fence.acq_rel.sys;" ::: "memory"); + __syncwarp(); + int const mine = cur * kSlotWords + p.tpRank * kMaxRows; + if (p.mode == ExchangeMode::kMulticast) + { + if (lane < rows) + { + mcStoreRelaxedSys(p.mcSlots + mine + lane, localKeys[lane]); + } + } + else + { + for (int i = lane; i < p.tpSize * rows; i += 32) + { + int const q = i / rows; + int const r = i - q * rows; + stRelaxedSys(p.ucSlotsPeers[q] + mine + r, localKeys[r]); + } + } + long long const start = clock64(); + for (int i = lane; i < p.tpSize * rows; i += 32) + { + int const q = i / rows; + int const r = i - q * rows; + uint64_t const* addr = p.ucSlotsLocal + cur * kSlotWords + q * kMaxRows + r; + uint64_t v = ldRelaxedSys(addr); + while (v == 0ull) + { + if (p.timeoutCycles > 0 && clock64() - start > p.timeoutCycles) + { + // A peer never published: the ranks disagree on the + // collective sequence. Record it and abort loudly. + atomicExch(reinterpret_cast(p.scratch + kMaxRows + 1), + 0x5733000000000000ull | (static_cast(epoch) << 8) + | static_cast(q)); + __threadfence_system(); + __trap(); + } + v = ldRelaxedSys(addr); + } + gathered[q][r] = v; + } + __syncwarp(); + if (lane < rows) + { + for (int q = 0; q < p.tpSize; ++q) + { + result = umax64(result, gathered[q][lane]); + } + } + if (lane == 0) + { + *epochPtr = epoch + 1; + } + } + if (lane < rows) + { + p.tokens[lane] = keyToToken(result); + if (p.keysOut != nullptr) + { + p.keysOut[lane] = result; + } + } +} + +// --------------------------------------------------------------------------- +// Fused kernel: GEMV + rank-local max + exchange +// +// CTA c owns a contiguous block of vocab rows; warp w of it takes rows +// begin + w, begin + w + kWarps, ... Each row is split into kSplit units of +// kUnit 32-byte chunks per lane, double-buffered in registers so a warp always +// has the next unit in flight while it computes the current one. The first +// unit is requested before griddepcontrol.wait, i.e. while the producer of +// the hidden rows may still be running under programmatic dependent launch. +// --------------------------------------------------------------------------- +// --------------------------------------------------------------------------- +// Direct exchange: every CTA publishes its partial key for each row into +// entry [slot][row][rank][cta] of every rank's buffer; CTA 0 of each rank polls +// all tpSize x numCtas entries of each row, writes the tokens, then clears +// exactly the entries it polled and advances the epoch. A slot's entries are +// written again only three epochs later, i.e. after this rank's next two +// publishes, which follow this launch; so no peer can write a cleared entry +// before the clear has completed. CTAs without vocab rows publish key 1, below +// every real key. +// --------------------------------------------------------------------------- +__device__ __forceinline__ bool isDirectMode(ExchangeMode m) +{ + return m == ExchangeMode::kMulticastDirect || m == ExchangeMode::kUnicastDirect; +} + +__device__ __forceinline__ size_t directEntry(int slot, int row, int rank, int cta) +{ + return kDirectOffsetWords + static_cast(slot) * kDirectSlotWords + + (static_cast(row) * kMaxRanks + rank) * kMaxDirectCtas + cta; +} + +// All threads of the CTA call this once ctaKeys[0..rows) (shared) is final. +__device__ void directExchange(DraftTailParams const& p, uint64_t const* ctaKeys, int ctaIdx, int numCtas) +{ + __shared__ uint32_t sEpoch; + __shared__ uint64_t red[32][kMaxRows]; + int const tid = threadIdx.x; + int const lane = tid & 31; + int const warp = tid >> 5; + int const rows = p.rows; + if (tid == 0) + { + sEpoch = *reinterpret_cast(epochCounter(p.scratch)); + } + __syncthreads(); + uint32_t const epoch = sEpoch; + int const cur = static_cast(epoch % kNumSlots); + if (p.mode == ExchangeMode::kMulticastDirect) + { + if (tid < rows) + { + uint64_t const key = ctaKeys[tid] != 0ull ? ctaKeys[tid] : 1ull; + mcStoreRelaxedSys(p.mcSlots + directEntry(cur, tid, p.tpRank, ctaIdx), key); + } + } + else + { + for (int i = tid; i < p.tpSize * rows; i += blockDim.x) + { + int const q = i / rows; + int const r = i - q * rows; + uint64_t const key = ctaKeys[r] != 0ull ? ctaKeys[r] : 1ull; + stRelaxedSys(p.ucSlotsPeers[q] + directEntry(cur, r, p.tpRank, ctaIdx), key); + } + } + if (ctaIdx != 0) + { + return; + } + long long const start = clock64(); + int const perRow = p.tpSize * numCtas; + // Each thread issues up to kPoll independent loads per round and re-polls + // only the entries still empty, so a row costs about one memory round + // trip after the last entry lands, however many CTAs published. + constexpr int kPoll = 32; + for (int r = 0; r < rows; ++r) + { + uint64_t m = 0ull; + for (int base = tid; base < perRow; base += kPoll * blockDim.x) + { + uint64_t v[kPoll]; +#pragma unroll + for (int b = 0; b < kPoll; ++b) + { + int const i = base + b * blockDim.x; + v[b] = i < perRow ? ldRelaxedSys(p.ucSlotsLocal + directEntry(cur, r, i / numCtas, i % numCtas)) : 1ull; + } + bool done = false; + while (!done) + { + done = true; +#pragma unroll + for (int b = 0; b < kPoll; ++b) + { + if (v[b] == 0ull) + { + int const i = base + b * blockDim.x; + done = false; + v[b] = ldRelaxedSys(p.ucSlotsLocal + directEntry(cur, r, i / numCtas, i % numCtas)); + } + } + if (!done && p.timeoutCycles > 0 && clock64() - start > p.timeoutCycles) + { + atomicExch(reinterpret_cast(p.scratch + kMaxRows + 1), + 0x5734000000000000ull | (static_cast(epoch) << 16) + | static_cast(base)); + __threadfence_system(); + __trap(); + } + } +#pragma unroll + for (int b = 0; b < kPoll; ++b) + { + m = umax64(m, v[b]); + } + } + m = warpMax64(m); + if (lane == 0) + { + red[warp][r] = m; + } + } + __syncthreads(); + if (tid < rows) + { + uint64_t result = 0ull; + int const numWarps = (blockDim.x + 31) >> 5; + for (int w = 0; w < numWarps; ++w) + { + result = umax64(result, red[w][tid]); + } + p.tokens[tid] = keyToToken(result); + if (p.keysOut != nullptr) + { + p.keysOut[tid] = result; + } + } + for (int i = tid; i < rows * perRow; i += blockDim.x) + { + int const r = i / perRow; + int const j = i - r * perRow; + int const q = j / numCtas; + int const c = j - q * numCtas; + p.ucSlotsLocal[directEntry(cur, r, q, c)] = 0ull; + } + if (tid == 0) + { + *epochCounter(p.scratch) = epoch + 1; + } +} + +// Per-CTA epilogue of the fused kernels: CTA max of the warps' keys, one +// atomicMax per row into rank-local scratch, then the last CTA to arrive takes +// (and resets) the rank-local keys and runs the exchange. +template +__device__ __forceinline__ void ctaFinish(DraftTailParams const& p, uint64_t const (&best)[R]) +{ + __shared__ uint64_t warpBest[kWarps][R]; + __shared__ uint64_t ctaKeys[kMaxRows]; + __shared__ uint64_t gathered[kMaxRanks][kMaxRows]; + __shared__ int isLast; + int const lane = threadIdx.x & 31; + int const warp = threadIdx.x >> 5; + if (lane == 0) + { +#pragma unroll + for (int r = 0; r < R; ++r) + { + warpBest[warp][r] = best[r]; + } + } + __syncthreads(); + if (isDirectMode(p.mode)) + { + if (threadIdx.x < R) + { + uint64_t m = 0ull; +#pragma unroll + for (int w = 0; w < kWarps; ++w) + { + m = umax64(m, warpBest[w][threadIdx.x]); + } + ctaKeys[threadIdx.x] = m; + } + __syncthreads(); + directExchange(p, ctaKeys, blockIdx.x, gridDim.x); + return; + } + if (threadIdx.x < R) + { + uint64_t m = 0ull; +#pragma unroll + for (int w = 0; w < kWarps; ++w) + { + m = umax64(m, warpBest[w][threadIdx.x]); + } + if (m != 0ull) + { + atomicMax(reinterpret_cast(p.scratch + threadIdx.x), m); + } + } + __syncthreads(); + if (threadIdx.x == 0) + { + __threadfence(); + uint32_t const prev = atomicAdd(arrivalCounter(p.scratch), 1u); + isLast = prev == gridDim.x - 1; + } + __syncthreads(); + if (!isLast || warp != 0) + { + return; + } + __threadfence(); + if (lane < R) + { + ctaKeys[lane] = atomicExch(reinterpret_cast(p.scratch + lane), 0ull); + } + if (lane == 0) + { + atomicExch(arrivalCounter(p.scratch), 0u); + } + __syncwarp(); + exchangeAndWrite(p, ctaKeys, gathered); +} + +template +__device__ __forceinline__ void loadUnit(Chunk (&buf)[kUnit], char const* base, uint64_t policy) +{ +#pragma unroll + for (int j = 0; j < kUnit; ++j) + { + loadChunk(buf[j], base + j * 32 * 32, policy); + } +} + +// Byte address of unit u of this lane: row firstRow + (u / kSplit) * kWarps, +// part u % kSplit; chunk c of a row lives at bytes [32 c, 32 c + 32) and lane l +// owns chunks part * kUnit * 32 + j * 32 + l. +template +__device__ __forceinline__ char const* unitAddress(__nv_bfloat16 const* weight, int firstRow, int u, int lane) +{ + constexpr int kUnit = kChunks / kSplit; + int const rowIdx = u / kSplit; + int const part = u - rowIdx * kSplit; + int64_t const row = firstRow + static_cast(rowIdx) * kWarps; + return reinterpret_cast(weight + row * (kChunks * 32 * 16)) + (part * kUnit * 32 + lane) * 32; +} + +template +__device__ __forceinline__ void computeUnit(Chunk const (&buf)[kChunks / kSplit], Chunk const* smemH, int u, + int firstRow, int lane, int64_t vocabOffset, float (&acc)[R], uint64_t (&best)[R]) +{ + constexpr int kUnit = kChunks / kSplit; + int const rowIdx = u / kSplit; + int const part = u - rowIdx * kSplit; +#pragma unroll + for (int j = 0; j < kUnit; ++j) + { + int const c = (part * kUnit + j) * 32 + lane; +#pragma unroll + for (int r = 0; r < R; ++r) + { + acc[r] = dotChunk(buf[j], smemH[r * kChunks * 32 + c], acc[r]); + } + } + if (part == kSplit - 1) + { + int64_t const row = firstRow + static_cast(rowIdx) * kWarps; +#pragma unroll + for (int r = 0; r < R; ++r) + { + float s = acc[r]; +#pragma unroll + for (int off = 16; off > 0; off >>= 1) + { + s += __shfl_xor_sync(0xFFFFFFFFu, s, off); + } + uint32_t const bits = static_cast(__bfloat16_as_ushort(__float2bfloat16_rn(s))); + best[r] = umax64(best[r], makeKey(bits, vocabOffset + row)); + acc[r] = 0.f; + } + } +} + +template +__global__ void __launch_bounds__(kWarps * 32, kMinBlocks) fusedDraftTailKernel(DraftTailParams p) +{ + constexpr int kUnit = kChunks / kSplit; + static_assert(kChunks % kSplit == 0, "row must split evenly"); + + extern __shared__ uint4 smemDyn[]; + Chunk* smemH = reinterpret_cast(smemDyn); // [R][kChunks * 32] + + int const lane = threadIdx.x & 31; + int const warp = threadIdx.x >> 5; + int const numRows = p.vocabLocal; + int const rowsPerCta = (numRows + gridDim.x - 1) / gridDim.x; + int const rowBegin = blockIdx.x * rowsPerCta; + int const rowEnd = min(numRows, rowBegin + rowsPerCta); + int const firstRow = rowBegin + warp; + int const myRows = firstRow < rowEnd ? (rowEnd - firstRow + kWarps - 1) / kWarps : 0; + int const numUnits = myRows * kSplit; + + uint64_t const policyFirst = kLoad == kLoad32Hint ? evictFirstPolicy() : 0ull; + uint64_t const policyLast = kLoad == kLoad32Hint ? evictLastPolicy() : 0ull; + int const keepRows = static_cast((static_cast(myRows) * p.keepPermille) / 1000); + __nv_bfloat16 const* weight = p.weight; + + Chunk bufA[kUnit]; + Chunk bufB[kUnit]; + if (numUnits > 0) + { + loadUnit(bufA, unitAddress(weight, firstRow, 0, lane), + keepRows > 0 ? policyLast : policyFirst); + } + + griddepWait(); + + { + int const total = R * kChunks * 32; + for (int i = threadIdx.x; i < total; i += blockDim.x) + { + int const r = i / (kChunks * 32); + int const c = i - r * (kChunks * 32); + uint4 const* src = reinterpret_cast(p.hidden + r * p.hiddenStride) + 2 * c; + uint4 const lo = src[0]; + uint4 const hi = src[1]; + Chunk ch; + ch.v[0] = (static_cast(lo.y) << 32) | lo.x; + ch.v[1] = (static_cast(lo.w) << 32) | lo.z; + ch.v[2] = (static_cast(hi.y) << 32) | hi.x; + ch.v[3] = (static_cast(hi.w) << 32) | hi.z; + smemH[i] = ch; + } + } + __syncthreads(); + + float acc[R]; + uint64_t best[R]; +#pragma unroll + for (int r = 0; r < R; ++r) + { + acc[r] = 0.f; + best[r] = 0ull; + } + + for (int u = 0; u < numUnits; u += 2) + { + if (u + 1 < numUnits) + { + loadUnit(bufB, unitAddress(weight, firstRow, u + 1, lane), + (u + 1) / kSplit < keepRows ? policyLast : policyFirst); + } + computeUnit(bufA, smemH, u, firstRow, lane, p.vocabOffset, acc, best); + if (u + 1 < numUnits) + { + if (u + 2 < numUnits) + { + loadUnit(bufA, unitAddress(weight, firstRow, u + 2, lane), + (u + 2) / kSplit < keepRows ? policyLast : policyFirst); + } + computeUnit(bufB, smemH, u + 1, firstRow, lane, p.vocabOffset, acc, best); + } + } + + ctaFinish(p, best); +} + +// --------------------------------------------------------------------------- +// Fused kernel, bulk-copy flavour: each warp streams its rows through its own +// ring of kSlots shared-memory row buffers filled by cp.async.bulk (one 12 KB +// copy per row, completion on a per-slot mbarrier), so the bytes in flight do +// not depend on registers. Lane l reads chunks l + 32 j of a row in the same +// order as the register kernel, so both produce bitwise-identical keys. +// --------------------------------------------------------------------------- +__device__ __forceinline__ uint32_t smemAddr(void const* ptr) +{ + return static_cast(__cvta_generic_to_shared(ptr)); +} + +__device__ __forceinline__ void mbarInit(uint64_t* bar, uint32_t count) +{ + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" ::"r"(smemAddr(bar)), "r"(count) : "memory"); +} + +__device__ __forceinline__ void mbarExpectTx(uint64_t* bar, uint32_t bytes) +{ + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(smemAddr(bar)), "r"(bytes) : "memory"); +} + +__device__ __forceinline__ void mbarWait(uint64_t* bar, uint32_t phase) +{ + uint32_t const addr = smemAddr(bar); + asm volatile( + "{\n" + ".reg .pred done;\n" + "DRAFT_TAIL_WAIT_%=:\n" + "mbarrier.try_wait.parity.shared::cta.b64 done, [%0], %1;\n" + "@!done bra DRAFT_TAIL_WAIT_%=;\n" + "}\n" ::"r"(addr), + "r"(phase) + : "memory"); +} + +__device__ __forceinline__ void bulkLoad(void* dst, void const* src, uint32_t bytes, uint64_t* bar, uint64_t policy) +{ + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint [%0], [%1], %2, [%3], %4;" :: + "r"(smemAddr(dst)), + "l"(src), "r"(bytes), "r"(smemAddr(bar)), "l"(policy) + : "memory"); +} + +template +__global__ void __launch_bounds__(kWarps * 32, 1) fusedBulkDraftTailKernel(DraftTailParams p) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 900) + // Bulk copies and transaction mbarriers need SM90+. + __trap(); +#else + constexpr int kChunks = kHidden6144ChunksDev; + constexpr int kRowBytes = kChunks * 32 * 32; + extern __shared__ uint4 smemDyn[]; + Chunk* smemH = reinterpret_cast(smemDyn); // [R][kChunks * 32] + Chunk* ring = smemH + R * kChunks * 32; // [kWarps][kSlots][kChunks * 32] + __shared__ __align__(8) uint64_t bars[kWarps][kSlots]; + + int const lane = threadIdx.x & 31; + int const warp = threadIdx.x >> 5; + int const numRows = p.vocabLocal; + int const rowsPerCta = (numRows + gridDim.x - 1) / gridDim.x; + int const rowBegin = blockIdx.x * rowsPerCta; + int const rowEnd = min(numRows, rowBegin + rowsPerCta); + int const firstRow = rowBegin + warp; + int const myRows = firstRow < rowEnd ? (rowEnd - firstRow + kWarps - 1) / kWarps : 0; + int const keepRows = static_cast((static_cast(myRows) * p.keepPermille) / 1000); + Chunk* myRing = ring + warp * kSlots * kChunks * 32; + + uint64_t const policyFirst = evictFirstPolicy(); + uint64_t const policyLast = evictLastPolicy(); + auto rowPtr = [&](int i) -> void const* + { return p.weight + (static_cast(firstRow) + static_cast(i) * kWarps) * (kChunks * 32 * 16); }; + + if (lane == 0) + { +#pragma unroll + for (int s = 0; s < kSlots; ++s) + { + mbarInit(&bars[warp][s], 1); + } + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); +#pragma unroll + for (int s = 0; s < kSlots; ++s) + { + if (s < myRows) + { + mbarExpectTx(&bars[warp][s], kRowBytes); + bulkLoad(myRing + s * kChunks * 32, rowPtr(s), kRowBytes, &bars[warp][s], + s < keepRows ? policyLast : policyFirst); + } + } + } + + griddepWait(); + { + int const total = R * kChunks * 32; + for (int i = threadIdx.x; i < total; i += blockDim.x) + { + int const r = i / (kChunks * 32); + int const c = i - r * (kChunks * 32); + uint4 const* src = reinterpret_cast(p.hidden + r * p.hiddenStride) + 2 * c; + uint4 const lo = src[0]; + uint4 const hi = src[1]; + Chunk ch; + ch.v[0] = (static_cast(lo.y) << 32) | lo.x; + ch.v[1] = (static_cast(lo.w) << 32) | lo.z; + ch.v[2] = (static_cast(hi.y) << 32) | hi.x; + ch.v[3] = (static_cast(hi.w) << 32) | hi.z; + smemH[i] = ch; + } + } + __syncthreads(); + + uint64_t best[R]; +#pragma unroll + for (int r = 0; r < R; ++r) + { + best[r] = 0ull; + } + for (int i = 0; i < myRows; ++i) + { + int const slot = i % kSlots; + mbarWait(&bars[warp][slot], static_cast((i / kSlots) & 1)); + Chunk const* row = myRing + slot * kChunks * 32; + float acc[R]; +#pragma unroll + for (int r = 0; r < R; ++r) + { + acc[r] = 0.f; + } +#pragma unroll + for (int j = 0; j < kChunks; ++j) + { + int const c = j * 32 + lane; + Chunk const w = row[c]; +#pragma unroll + for (int r = 0; r < R; ++r) + { + acc[r] = dotChunk(w, smemH[r * kChunks * 32 + c], acc[r]); + } + } + __syncwarp(); + if (lane == 0 && i + kSlots < myRows) + { + mbarExpectTx(&bars[warp][slot], kRowBytes); + bulkLoad(myRing + slot * kChunks * 32, rowPtr(i + kSlots), kRowBytes, &bars[warp][slot], + i + kSlots < keepRows ? policyLast : policyFirst); + } + int64_t const rowId = firstRow + static_cast(i) * kWarps; +#pragma unroll + for (int r = 0; r < R; ++r) + { + float sum = acc[r]; +#pragma unroll + for (int off = 16; off > 0; off >>= 1) + { + sum += __shfl_xor_sync(0xFFFFFFFFu, sum, off); + } + uint32_t const bits = static_cast(__bfloat16_as_ushort(__float2bfloat16_rn(sum))); + best[r] = umax64(best[r], makeKey(bits, p.vocabOffset + rowId)); + } + } + ctaFinish(p, best); +#endif +} + +// --------------------------------------------------------------------------- +// Fused kernel, tensor-core flavour: each warp owns 16-row tiles of the vocab +// shard and computes them with mma.sync m16n8k16 (BF16 in, FP32 accumulate), +// the R <= 8 hidden rows being the n columns. Lane (g = lane / 4, t = lane % 4) +// loads 32 contiguous bytes (16 values) of rows g and g + 8 per 64-value K +// block; inside a block the K order is permuted consistently for A and B +// (values 16t + 4s .. 16t + 4s + 3 feed k-step s), so a lane's two 16-byte +// halves map directly onto the A/B fragments without shuffles. This replaces +// the CUDA-core unpack + FMA loop (instruction bound) with 4 HMMA per 2 KB +// of weight per warp. +// +// Native-order flavours (MmaOrder below): k-step q of a 64-value block takes +// the 16 consecutive values 16q .. 16q + 15, so every MMA sums the same K +// values, in the same k-step order, as a GEMM that walks K in 16-value steps. +// They differ only in how the values reach the lanes. +// --------------------------------------------------------------------------- +__device__ __forceinline__ void mmaBf16( + float (&d)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t b0, uint32_t b1) +{ + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, " + "{%0, %1, %2, %3};\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)); +} + +constexpr int kMmaBlockElems = 64; // K values per lane-group block (4 lanes x 16) +constexpr int kMmaBlocks = 6144 / kMmaBlockElems; // 96 +constexpr int kMmaHStride = 6144 + 16; // padded smem row (elements) against bank conflicts + +// Which K values lane t (= lane % 4) of a quad feeds into the 4 mma k-steps +// q = 0..3 of each 64-value block, and how they are loaded. Every lane feeds +// 4 values (one 8-byte piece) per k-step and row. +enum MmaOrder : int +{ + // Lane t: values 16 t .. 16 t + 15 (one 32-byte load), piece q into k-step + // q; each k-step sums a K subset spread over the block (not native order). + kOrderPermuted = 0, + // Native orders: k-step q sums exactly values 16 q .. 16 q + 15. + // Lane t loads piece t of every k-step: four 8-byte loads. + kOrderNat8 = 1, + // Lane t loads k-step t whole (one 32-byte load); three in-quad shuffles + // hand piece (t - q) mod 4 of k-step q to lane t. + kOrderRot32 = 2, + // Lane t loads half (t >> 1) of k-steps 2 (t & 1) and 2 (t & 1) + 1 (two + // 16-byte loads), keeps the first piece of each half and swaps the second + // with lane t ^ 1. + kOrderPair16 = 3, +}; + +// Byte offset of lane t's first load inside a 128-byte block row. +template +__device__ __forceinline__ int mmaLaneOffset(int t) +{ + if constexpr (kOrder == kOrderNat8) + { + return 8 * t; + } + else if constexpr (kOrder == kOrderPair16) + { + return 64 * (t & 1) + 16 * (t >> 1); + } + else + { + return 32 * t; + } +} + +// Rot32 / Pair16: which 8-byte piece (0..3) of k-step q lane t feeds. +template +__device__ __forceinline__ int mmaNativePiece(int t, int q) +{ + if constexpr (kOrder == kOrderRot32) + { + return (t - q) & 3; + } + else + { + return q < 2 ? t : (t ^ 1); + } +} + +template +__device__ __forceinline__ void load16(uint64_t& a, uint64_t& b, void const* ptr, uint64_t policy) +{ + if constexpr (kLoad == kLoad32Hint) + { + asm volatile("ld.global.nc.L1::no_allocate.L2::cache_hint.v2.b64 {%0, %1}, [%2], %3;" + : "=l"(a), "=l"(b) + : "l"(ptr), "l"(policy)); + } + else + { + asm volatile("ld.global.nc.L1::no_allocate.v2.b64 {%0, %1}, [%2];" : "=l"(a), "=l"(b) : "l"(ptr)); + } +} + +template +__device__ __forceinline__ uint64_t load8(void const* ptr, uint64_t policy) +{ + uint64_t v; + if constexpr (kLoad == kLoad32Hint) + { + asm volatile("ld.global.nc.L1::no_allocate.L2::cache_hint.b64 %0, [%1], %2;" : "=l"(v) : "l"(ptr), "l"(policy)); + } + else + { + asm volatile("ld.global.nc.L1::no_allocate.b64 %0, [%1];" : "=l"(v) : "l"(ptr)); + } + return v; +} + +// row0 / row1 already include the lane offset (mmaLaneOffset). +template +__device__ __forceinline__ void mmaLoadStage( + Chunk (&x)[kStage], Chunk (&y)[kStage], char const* row0, char const* row1, int block, uint64_t policy) +{ +#pragma unroll + for (int j = 0; j < kStage; ++j) + { + int const off = (block + j) * kMmaBlockElems * 2; + if constexpr (kOrder == kOrderNat8) + { +#pragma unroll + for (int q = 0; q < 4; ++q) + { + x[j].v[q] = load8(row0 + off + 32 * q, policy); + y[j].v[q] = load8(row1 + off + 32 * q, policy); + } + } + else if constexpr (kOrder == kOrderPair16) + { + // v[0], v[1]: k-step 2 (t & 1); v[2], v[3]: k-step 2 (t & 1) + 1. + load16(x[j].v[0], x[j].v[1], row0 + off, policy); + load16(x[j].v[2], x[j].v[3], row0 + off + 32, policy); + load16(y[j].v[0], y[j].v[1], row1 + off, policy); + load16(y[j].v[2], y[j].v[3], row1 + off + 32, policy); + } + else + { + loadChunk(x[j], row0 + off, policy); + loadChunk(y[j], row1 + off, policy); + } + } +} + +// a[q]: the weight piece lane t feeds into k-step q, from what it loaded. +// Every lane of the warp must call this (in-quad shuffles). +template +__device__ __forceinline__ void mmaWeightPieces(uint64_t (&a)[4], Chunk const& c, int t) +{ + if constexpr (kOrder == kOrderRot32) + { + // c = k-step t. Round r takes piece r of k-step (t - r) mod 4 from the + // lane that loaded it, so k-step q arrives in round (t - q) mod 4. + uint64_t s[4]; + s[0] = c.v[0]; + int const quadBase = (threadIdx.x & 31) & ~3; +#pragma unroll + for (int r = 1; r < 4; ++r) + { + s[r] = __shfl_sync(0xFFFFFFFFu, c.v[r], quadBase | ((t - r) & 3)); + } +#pragma unroll + for (int q = 0; q < 4; ++q) + { + int const idx = (t - q) & 3; + uint64_t const lo = (idx & 1) ? s[1] : s[0]; + uint64_t const hi = (idx & 1) ? s[3] : s[2]; + a[q] = (idx & 2) ? hi : lo; + } + } + else if constexpr (kOrder == kOrderPair16) + { + // Even lanes hold k-steps 0, 1 and odd lanes k-steps 2, 3 (pieces + // 2 (t >> 1) and 2 (t >> 1) + 1 of each); swap the second pieces. + bool const odd = (t & 1) != 0; + uint64_t const r0 = __shfl_xor_sync(0xFFFFFFFFu, c.v[1], 1); + uint64_t const r1 = __shfl_xor_sync(0xFFFFFFFFu, c.v[3], 1); + a[0] = odd ? r0 : c.v[0]; + a[1] = odd ? r1 : c.v[2]; + a[2] = odd ? c.v[0] : r0; + a[3] = odd ? c.v[2] : r1; + } + else + { +#pragma unroll + for (int q = 0; q < 4; ++q) + { + a[q] = c.v[q]; + } + } +} + +template +__device__ __forceinline__ void mmaComputeStage(float (&d)[4], Chunk const (&x)[kStage], Chunk const (&y)[kStage], + __nv_bfloat16 const* hRow, bool hLive, int block, int t) +{ +#pragma unroll + for (int j = 0; j < kStage; ++j) + { + uint4 h0 = make_uint4(0, 0, 0, 0); + uint4 h1 = make_uint4(0, 0, 0, 0); + if (hLive) + { + if constexpr (kOrder == kOrderNat8) + { + // Values 16 q + 4 t .. + 3 for q = 0..3: 8 bytes every 32. + uint2 const* src = reinterpret_cast(hRow + (block + j) * kMmaBlockElems + 4 * t); + uint2 const q0 = src[0]; + uint2 const q1 = src[4]; + uint2 const q2 = src[8]; + uint2 const q3 = src[12]; + h0 = make_uint4(q0.x, q0.y, q1.x, q1.y); + h1 = make_uint4(q2.x, q2.y, q3.x, q3.y); + } + else + { + // Permuted: values 16 t .. 16 t + 15 of the block. Rot32 / + // Pair16: the row is staged in lane order, so the same 32 + // bytes are lane t's pieces for k-steps 0..3. + uint4 const* src = reinterpret_cast(hRow + (block + j) * kMmaBlockElems + 16 * t); + h0 = src[0]; + h1 = src[1]; + } + } + uint32_t const hb[8] = {h0.x, h0.y, h0.z, h0.w, h1.x, h1.y, h1.z, h1.w}; + uint64_t xp[4]; + uint64_t yp[4]; + mmaWeightPieces(xp, x[j], t); + mmaWeightPieces(yp, y[j], t); +#pragma unroll + for (int sIdx = 0; sIdx < 4; ++sIdx) + { + // Piece s holds the pairs mapped to k positions (2t, 2t + 1) and + // (2t + 8, 2t + 9) of k-step s, for A and B alike. + uint64_t const xw = xp[sIdx]; + uint64_t const yw = yp[sIdx]; + mmaBf16(d, static_cast(xw), static_cast(yw), static_cast(xw >> 32), + static_cast(yw >> 32), hb[2 * sIdx], hb[2 * sIdx + 1]); + } + } +} + +template +__global__ void __launch_bounds__(kWarps * 32, kMinBlocks) fusedMmaDraftTailKernel(DraftTailParams p) +{ + static_assert(R >= 1 && R <= 8, "n columns of m16n8k16"); + static_assert(kMmaBlocks % (2 * kStage) == 0, "double-buffered stages must tile K"); + extern __shared__ uint4 smemDyn[]; + __nv_bfloat16* smemH = reinterpret_cast<__nv_bfloat16*>(smemDyn); // [R][kMmaHStride] + + int const lane = threadIdx.x & 31; + int const warp = threadIdx.x >> 5; + int const g = lane >> 2; + int const t = lane & 3; + int const numRows = p.vocabLocal; + int const numTiles = (numRows + 15) / 16; + int const warpStride = gridDim.x * kWarps; + int tile = blockIdx.x * kWarps + warp; + + uint64_t const policy = kLoad == kLoad32Hint ? evictFirstPolicy() : 0ull; + char const* weight = reinterpret_cast(p.weight); + // Rows beyond the shard read row 0 (valid memory); their results are dropped. + auto rowPtr = [&](int row) -> char const* + { return weight + static_cast(row < numRows ? row : 0) * (6144 * 2) + mmaLaneOffset(t); }; + + Chunk xa[kStage], ya[kStage], xb[kStage], yb[kStage]; + if (tile < numTiles) + { + mmaLoadStage(xa, ya, rowPtr(tile * 16 + g), rowPtr(tile * 16 + g + 8), 0, policy); + } + + griddepWait(); + if constexpr (kOrder == kOrderRot32 || kOrder == kOrderPair16) + { + // Lane order: in each 64-value block, lane t's pieces for k-steps + // 0..3 are contiguous at values 16 t .. 16 t + 15. + for (int i = threadIdx.x; i < R * kMmaBlocks * 16; i += blockDim.x) + { + int const r = i / (kMmaBlocks * 16); + int const rem = i - r * (kMmaBlocks * 16); + int const block = rem >> 4; + int const lt = (rem >> 2) & 3; + int const q = rem & 3; + reinterpret_cast(smemH + r * kMmaHStride)[rem] = reinterpret_cast( + p.hidden + r * p.hiddenStride)[block * 16 + q * 4 + mmaNativePiece(lt, q)]; + } + } + else + { + for (int i = threadIdx.x; i < R * 6144 / 8; i += blockDim.x) + { + int const r = i / (6144 / 8); + int const c = i - r * (6144 / 8); + reinterpret_cast(smemH + r * kMmaHStride)[c] + = reinterpret_cast(p.hidden + r * p.hiddenStride)[c]; + } + } + __syncthreads(); + + bool const hLive = g < R; + __nv_bfloat16 const* hRow = smemH + (hLive ? g : 0) * kMmaHStride; + uint64_t bestJ[2] = {0ull, 0ull}; + + for (; tile < numTiles; tile += warpStride) + { + int const row0 = tile * 16 + g; + int const row1 = row0 + 8; + char const* p0 = rowPtr(row0); + char const* p1 = rowPtr(row1); + float d[4] = {0.f, 0.f, 0.f, 0.f}; + for (int b = 0; b < kMmaBlocks; b += 2 * kStage) + { + mmaLoadStage(xb, yb, p0, p1, b + kStage, policy); + mmaComputeStage(d, xa, ya, hRow, hLive, b, t); + if (b + 2 * kStage < kMmaBlocks) + { + mmaLoadStage(xa, ya, p0, p1, b + 2 * kStage, policy); + } + else if (tile + warpStride < numTiles) + { + int const next = tile + warpStride; + mmaLoadStage( + xa, ya, rowPtr(next * 16 + g), rowPtr(next * 16 + g + 8), 0, policy); + } + mmaComputeStage(d, xb, yb, hRow, hLive, b + kStage, t); + } + // d[j] = (row0, column 2t + j), d[2 + j] = (row1, column 2t + j). +#pragma unroll + for (int j = 0; j < 2; ++j) + { + int const n = 2 * t + j; + if (n < R) + { + if (row0 < numRows) + { + __nv_bfloat16 const v = __float2bfloat16_rn(d[j]); + bestJ[j] = umax64(bestJ[j], makeKey(__bfloat16_as_ushort(v), p.vocabOffset + row0)); + if (p.logitsOut != nullptr) + { + p.logitsOut[n * p.logitsOutStride + row0] = v; + } + } + if (row1 < numRows) + { + __nv_bfloat16 const v = __float2bfloat16_rn(d[2 + j]); + bestJ[j] = umax64(bestJ[j], makeKey(__bfloat16_as_ushort(v), p.vocabOffset + row1)); + if (p.logitsOut != nullptr) + { + p.logitsOut[n * p.logitsOutStride + row1] = v; + } + } + } + } + } + // Reduce over g (lane bits 2..4), then lane 0 collects column n from lane n / 2. +#pragma unroll + for (int j = 0; j < 2; ++j) + { +#pragma unroll + for (int off = 4; off < 32; off <<= 1) + { + bestJ[j] = umax64(bestJ[j], __shfl_xor_sync(0xFFFFFFFFu, bestJ[j], off)); + } + } + uint64_t best[R]; +#pragma unroll + for (int n = 0; n < R; ++n) + { + best[n] = __shfl_sync(0xFFFFFFFFu, bestJ[n & 1], n >> 1); + } + ctaFinish(p, best); +} + +// --------------------------------------------------------------------------- +// Logits kernel: rank-local max over precomputed BF16 logits + exchange. +// --------------------------------------------------------------------------- +template +__global__ void __launch_bounds__(512, 1) fromLogitsDraftTailKernel(DraftTailParams p) +{ + __shared__ uint64_t warpBest[32][R]; + __shared__ uint64_t ctaKeys[kMaxRows]; + __shared__ uint64_t gathered[kMaxRanks][kMaxRows]; + + griddepWait(); + int const lane = threadIdx.x & 31; + int const warp = threadIdx.x >> 5; + int const numWarps = blockDim.x >> 5; + uint64_t best[R]; +#pragma unroll + for (int r = 0; r < R; ++r) + { + best[r] = 0ull; + } + // Flat (row, 8-element vector) index space; each thread issues kBatch + // independent 16-byte loads before reducing them, so a row costs about one + // memory round trip instead of one per loop iteration. + constexpr int kBatch = 8; + bool vec = (p.logitsStride % 8) == 0 && (reinterpret_cast(p.logits) & 15) == 0; + int const nVec = vec ? p.vocabLocal / 8 : 0; + int const total = R * nVec; + for (int base = threadIdx.x; base < total; base += kBatch * blockDim.x) + { + uint4 x[kBatch]; +#pragma unroll + for (int b = 0; b < kBatch; ++b) + { + int const i = base + b * blockDim.x; + if (i < total) + { + int const r = i / nVec; + int const v = i - r * nVec; + x[b] = __ldg(reinterpret_cast(p.logits + r * p.logitsStride) + v); + } + } +#pragma unroll + for (int b = 0; b < kBatch; ++b) + { + int const i = base + b * blockDim.x; + if (i < total) + { + int const r = i / nVec; + int const v = i - r * nVec; + uint32_t const w[4] = {x[b].x, x[b].y, x[b].z, x[b].w}; + uint64_t m = 0ull; +#pragma unroll + for (int k = 0; k < 4; ++k) + { + int64_t const id = p.vocabOffset + 8 * v + 2 * k; + m = umax64(m, makeKey(w[k] & 0xFFFFu, id)); + m = umax64(m, makeKey(w[k] >> 16, id + 1)); + } +#pragma unroll + for (int rr = 0; rr < R; ++rr) + { + if (rr == r) + { + best[rr] = umax64(best[rr], m); + } + } + } + } + } +#pragma unroll + for (int r = 0; r < R; ++r) + { + __nv_bfloat16 const* row = p.logits + r * p.logitsStride; + for (int i = nVec * 8 + threadIdx.x; i < p.vocabLocal; i += blockDim.x) + { + uint32_t const bits = __bfloat16_as_ushort(row[i]); + best[r] = umax64(best[r], makeKey(bits, p.vocabOffset + i)); + } + best[r] = warpMax64(best[r]); + } + if (lane == 0) + { +#pragma unroll + for (int r = 0; r < R; ++r) + { + warpBest[warp][r] = best[r]; + } + } + __syncthreads(); + if (warp == 0) + { +#pragma unroll + for (int r = 0; r < R; ++r) + { + uint64_t m = lane < numWarps ? warpBest[lane][r] : 0ull; + m = warpMax64(m); + if (lane == 0) + { + ctaKeys[r] = m; + } + } + } + __syncthreads(); + if (isDirectMode(p.mode)) + { + directExchange(p, ctaKeys, 0, 1); + return; + } + if (warp != 0) + { + return; + } + exchangeAndWrite(p, ctaKeys, gathered); +} + +// --------------------------------------------------------------------------- +// Launch +// --------------------------------------------------------------------------- + +constexpr int kHidden6144Chunks = 6144 / 512; // 32-byte chunks per lane per row + +struct VariantDesc +{ + int warps; + int split; + int load; + int minBlocks; +}; + +// Kernel shape and load flavour, selectable at run time for tuning. +constexpr VariantDesc kVariants[] = { + {8, 2, kLoad32Hint, 1}, // 0: default + {8, 2, kLoad32, 1}, // 1: no L2 eviction hint + {8, 2, kLoad16x2Pf, 1}, // 2: 16-byte loads with L2 prefetch + {16, 3, kLoad32Hint, 1}, // 3: more warps, shorter units + {4, 2, kLoad32Hint, 2}, // 4: two CTAs per SM + {8, 1, kLoad32Hint, 1}, // 5: whole-row units + {8, 0, kLoad32Hint, 1}, // 6: per-warp cp.async.bulk row ring (R <= 2; else variant 0) + {8, 0, kLoad32Hint, 1}, // 7: tensor-core tiles, 8 warps, 2-block stages + {4, 0, kLoad32Hint, 2}, // 8: tensor-core tiles, 4 warps, 2 CTAs/SM + {8, 0, kLoad32, 1}, // 9: tensor-core tiles, no L2 hint + {8, 0, kLoad32Hint, 1}, // 10: tensor-core tiles, 8 warps, 4-block stages + {8, 0, kLoad32Hint, 1}, // 11: as 10, native K order (kOrderNat8: 8-byte lane loads) + {4, 0, kLoad32Hint, 2}, // 12: native K order (kOrderNat8), 4 warps, 2 CTAs/SM + {8, 0, kLoad32Hint, 1}, // 13: as 10, native K order (kOrderRot32: 32-byte loads, quad shuffles) + {8, 0, kLoad32Hint, 1}, // 14: as 10, native K order (kOrderPair16: 16-byte loads, pair shuffles) +}; +constexpr int kNumVariants = sizeof(kVariants) / sizeof(kVariants[0]); +// Variants from here on are tensor-core kernels; their default grid covers +// one 16-row tile per warp. +constexpr int kFirstMmaVariant = 7; + +template +void launchFusedShape(DraftTailParams const& p, int numCtas, bool pdl, cudaStream_t stream) +{ + auto kernel = fusedDraftTailKernel; + size_t const smem = static_cast(R) * 6144 * sizeof(__nv_bfloat16); + static bool const attrSet = [&]() + { + TLLM_CUDA_CHECK( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(smem))); + return true; + }(); + (void) attrSet; + cudaLaunchConfig_t config{}; + config.gridDim = dim3(numCtas); + config.blockDim = dim3(kWarps * 32); + config.dynamicSmemBytes = smem; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + config.attrs = attrs; + config.numAttrs = 1; + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, kernel, p)); +} + +template +void launchBulkShape(DraftTailParams const& p, int numCtas, bool pdl, cudaStream_t stream) +{ + auto kernel = fusedBulkDraftTailKernel; + size_t const smem = static_cast(R + kWarps * kSlots) * 6144 * sizeof(__nv_bfloat16); + static bool const attrSet = [&]() + { + TLLM_CUDA_CHECK( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(smem))); + return true; + }(); + (void) attrSet; + cudaLaunchConfig_t config{}; + config.gridDim = dim3(numCtas); + config.blockDim = dim3(kWarps * 32); + config.dynamicSmemBytes = smem; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + config.attrs = attrs; + config.numAttrs = 1; + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, kernel, p)); +} + +template +void launchMmaShape(DraftTailParams const& p, int numCtas, bool pdl, cudaStream_t stream) +{ + auto kernel = fusedMmaDraftTailKernel; + size_t const smem = static_cast(R) * kMmaHStride * sizeof(__nv_bfloat16); + static bool const attrSet = [&]() + { + TLLM_CUDA_CHECK( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(smem))); + return true; + }(); + (void) attrSet; + if (numCtas <= 0) + { + int const numTiles = (p.vocabLocal + 15) / 16; + numCtas = (numTiles + kWarps - 1) / kWarps; + } + TLLM_CHECK_WITH_INFO((p.mode != ExchangeMode::kMulticastDirect && p.mode != ExchangeMode::kUnicastDirect) + || numCtas <= kMaxDirectCtas, + "glm_moe_dsa_draft_tail: direct exchange supports at most %d CTAs", kMaxDirectCtas); + cudaLaunchConfig_t config{}; + config.gridDim = dim3(numCtas); + config.blockDim = dim3(kWarps * 32); + config.dynamicSmemBytes = smem; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + config.attrs = attrs; + config.numAttrs = 1; + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, kernel, p)); +} + +template +void launchFusedRows(DraftTailParams const& p, int variant, int numCtas, bool pdl, cudaStream_t stream) +{ + switch (variant) + { + case 0: launchFusedShape(p, numCtas, pdl, stream); break; + case 1: launchFusedShape(p, numCtas, pdl, stream); break; + case 2: launchFusedShape(p, numCtas, pdl, stream); break; + case 3: launchFusedShape(p, numCtas, pdl, stream); break; + case 4: launchFusedShape(p, numCtas, pdl, stream); break; + case 5: launchFusedShape(p, numCtas, pdl, stream); break; + case 6: + if constexpr (R <= 2) + { + launchBulkShape(p, numCtas, pdl, stream); + } + else + { + launchFusedShape(p, numCtas, pdl, stream); + } + break; + case 7: launchMmaShape(p, numCtas, pdl, stream); break; + case 8: launchMmaShape(p, numCtas, pdl, stream); break; + case 9: launchMmaShape(p, numCtas, pdl, stream); break; + case 10: launchMmaShape(p, numCtas, pdl, stream); break; + case 11: launchMmaShape(p, numCtas, pdl, stream); break; + case 12: launchMmaShape(p, numCtas, pdl, stream); break; + case 13: launchMmaShape(p, numCtas, pdl, stream); break; + case 14: launchMmaShape(p, numCtas, pdl, stream); break; + default: TLLM_THROW("glm_moe_dsa_draft_tail: unknown fused variant %d", variant); + } +} + +int smCount() +{ + static int const count = []() + { + int dev = 0; + TLLM_CUDA_CHECK(cudaGetDevice(&dev)); + int n = 0; + TLLM_CUDA_CHECK(cudaDeviceGetAttribute(&n, cudaDevAttrMultiProcessorCount, dev)); + return n; + }(); + return count; +} + +} // namespace + +int numFusedVariants() +{ + return kNumVariants; +} + +void invokeDraftTailFused(DraftTailParams const& p, FusedLaunchConfig const& config, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(p.hiddenDim == 6144, "glm_moe_dsa_draft_tail: fused kernel supports hidden size 6144 only"); + TLLM_CHECK_WITH_INFO(p.rows >= 1 && p.rows <= kMaxRows, "glm_moe_dsa_draft_tail: rows out of range"); + TLLM_CHECK_WITH_INFO(config.variant >= 0 && config.variant < kNumVariants, "glm_moe_dsa_draft_tail: bad variant"); + TLLM_CHECK_WITH_INFO(p.logitsOut == nullptr || config.variant >= kFirstMmaVariant, + "glm_moe_dsa_draft_tail: logits output needs a tensor-core variant"); + int numCtas = config.numCtas; + if (numCtas <= 0 && config.variant < kFirstMmaVariant) + { + numCtas = smCount() * kVariants[config.variant].minBlocks; + } + TLLM_CHECK_WITH_INFO((p.mode != ExchangeMode::kMulticastDirect && p.mode != ExchangeMode::kUnicastDirect) + || numCtas <= kMaxDirectCtas, + "glm_moe_dsa_draft_tail: direct exchange supports at most %d CTAs", kMaxDirectCtas); + switch (p.rows) + { + case 1: launchFusedRows<1>(p, config.variant, numCtas, config.pdl, stream); break; + case 2: launchFusedRows<2>(p, config.variant, numCtas, config.pdl, stream); break; + case 3: launchFusedRows<3>(p, config.variant, numCtas, config.pdl, stream); break; + case 4: launchFusedRows<4>(p, config.variant, numCtas, config.pdl, stream); break; + case 5: launchFusedRows<5>(p, config.variant, numCtas, config.pdl, stream); break; + case 6: launchFusedRows<6>(p, config.variant, numCtas, config.pdl, stream); break; + case 7: launchFusedRows<7>(p, config.variant, numCtas, config.pdl, stream); break; + case 8: launchFusedRows<8>(p, config.variant, numCtas, config.pdl, stream); break; + default: break; + } +} + +template +void launchFromLogitsRows(DraftTailParams const& p, bool pdl, cudaStream_t stream) +{ + cudaLaunchConfig_t config{}; + config.gridDim = dim3(1); + config.blockDim = dim3(512); + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + config.attrs = attrs; + config.numAttrs = 1; + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, fromLogitsDraftTailKernel, p)); +} + +void invokeDraftTailFromLogits(DraftTailParams const& p, bool pdl, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(p.rows >= 1 && p.rows <= kMaxRows, "glm_moe_dsa_draft_tail: rows out of range"); + switch (p.rows) + { + case 1: launchFromLogitsRows<1>(p, pdl, stream); break; + case 2: launchFromLogitsRows<2>(p, pdl, stream); break; + case 3: launchFromLogitsRows<3>(p, pdl, stream); break; + case 4: launchFromLogitsRows<4>(p, pdl, stream); break; + case 5: launchFromLogitsRows<5>(p, pdl, stream); break; + case 6: launchFromLogitsRows<6>(p, pdl, stream); break; + case 7: launchFromLogitsRows<7>(p, pdl, stream); break; + case 8: launchFromLogitsRows<8>(p, pdl, stream); break; + default: break; + } +} + +} // namespace kernels::glm_moe_dsa::draft_tail + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/draftTail.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/draftTail.h new file mode 100644 index 000000000000..82e13a7f084a --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/draftTail.h @@ -0,0 +1,144 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Greedy draft-token selection over a vocab-sharded LM head. +// +// For R hidden rows and this rank's vocab shard W [V_local, H] the kernels +// produce, on every rank, the token id of the global maximum of the BF16 +// logits h @ W^T over all tensor-parallel shards; ties go to the lowest global +// token id and NaN ranks above every number. That is the selection rule of +// SpecWorkerBase.greedy_sample_draft_with_tp_gather. +// +// The logits kernel selects from BF16 logits computed by the caller (the +// native LM-head GEMM), so its tokens are the native ones. The fused kernels +// compute the logits themselves (FP32 accumulation, then BF16 rounding); +// where the largest logits are within one BF16 rounding step of each other, +// the selected token depends on the exact FP32 sums, i.e. on the accumulation +// order. Fused variants 11 to 14 accumulate like the native cuBLAS LM-head +// GEMM selected for this shape on SM103 (K in 16-value tensor-core steps from +// 0 to H into one FP32 accumulator), so their BF16 logits and tokens equal the +// native ones bit for bit (checked by the draft_tail unit tests against F.linear; +// re-check whenever cuBLAS or the GPU changes). The other fused variants +// accumulate in other orders and are decision-level variants at BF16 +// near-ties. +// +// Each candidate is packed into one 64-bit key, (ordered BF16 bits << 32) | +// (0xFFFFFFFF - global id), so the selection is a plain unsigned max that is +// associative, commutative and independent of reduction order. The keys of +// all ranks are exchanged through a multicast (or unicast peer) buffer with a +// 0-sentinel protocol over three rotating slots; see draftTail.cu. + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa::draft_tail +{ + +constexpr int kMaxRows = 8; +constexpr int kMaxRanks = 8; +constexpr int kNumSlots = 3; +// Exchange buffer: kNumSlots x kMaxRanks x kMaxRows 64-bit keys. +constexpr int kSlotWords = kMaxRanks * kMaxRows; +constexpr int kExchangeWords = kNumSlots * kSlotWords; +// Rank-local scratch, in 64-bit words: kMaxRows partial keys, then the CTA +// arrival counter and the exchange epoch (32 bits each in one word), then an +// error word. +constexpr int kScratchWords = kMaxRows + 2; + +enum class ExchangeMode : int +{ + kNone = 0, // tpSize == 1, or debug: tokens from the rank-local maximum + kMulticast = 1, // one multimem.st per row into the multicast slot + kUnicast = 2, // one store per (peer, row) into each peer's slot + // Direct variants: every CTA publishes its own partial key (no rank-local + // atomics or last-CTA election); CTA 0 of each rank polls all + // tpSize x numCtas entries, writes the tokens and recycles the slot. + kMulticastDirect = 3, + kUnicastDirect = 4, +}; + +constexpr int kMaxDirectCtas = 512; +constexpr int kDirectSlotWords = kMaxRows * kMaxRanks * kMaxDirectCtas; +// The direct region follows the last-CTA region in the same buffer. +constexpr int kDirectOffsetWords = kExchangeWords; +constexpr int kExchangeBufferWords = kExchangeWords + kNumSlots * kDirectSlotWords; + +struct DraftTailParams +{ + // Fused mode: hidden [rows, hiddenDim] (row stride hiddenStride elements) + // and the vocab shard weight [vocabLocal, hiddenDim], row-major. + __nv_bfloat16 const* hidden = nullptr; + int64_t hiddenStride = 0; + __nv_bfloat16 const* weight = nullptr; + // Logits mode: BF16 logits [rows, vocabLocal] (row stride logitsStride). + __nv_bfloat16 const* logits = nullptr; + int64_t logitsStride = 0; + + int rows = 0; + int vocabLocal = 0; + int hiddenDim = 0; + int64_t vocabOffset = 0; // global id of local column 0 + + int32_t* tokens = nullptr; // [rows] output + uint64_t* keysOut = nullptr; // optional [rows] output: the winning key + + uint64_t* scratch = nullptr; // [kScratchWords], zero-initialized once + + int tpRank = 0; + int tpSize = 1; + ExchangeMode mode = ExchangeMode::kNone; + uint64_t* mcSlots = nullptr; // multicast address of the exchange buffer + uint64_t* ucSlotsLocal = nullptr; // this rank's unicast address of the buffer + uint64_t* const* ucSlotsPeers = nullptr; // device array [tpSize] of unicast addresses + long long timeoutCycles = 0; // 0: wait forever + // Optional BF16 logits [rows, vocabLocal] (row stride logitsOutStride), + // written by the tensor-core fused kernels only (tests compare them with + // the native GEMM bit for bit); nullptr in production. + __nv_bfloat16* logitsOut = nullptr; + int64_t logitsOutStride = 0; + // Fused kernel, L2-hinted variants: each warp loads the first + // keepPermille / 1000 of its rows with an L2 evict_last policy (the rest + // evict_first), so that part of the shard tends to stay in L2 across the + // draft steps that re-read the same weight. + int keepPermille = 0; +}; + +struct FusedLaunchConfig +{ + int numCtas = 0; // 0: one per SM + int variant = 0; // see draftTail.cu: kernel shape and load flavour + bool pdl = true; +}; + +int numFusedVariants(); + +// GEMV + rank-local max + exchange in one launch. +void invokeDraftTailFused(DraftTailParams const& params, FusedLaunchConfig const& config, cudaStream_t stream); + +// Rank-local max over precomputed BF16 logits + exchange in one launch. +void invokeDraftTailFromLogits(DraftTailParams const& params, bool pdl, cudaStream_t stream); + +} // namespace kernels::glm_moe_dsa::draft_tail + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/indexerLayerRopeKvAppend.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/indexerLayerRopeKvAppend.cu new file mode 100644 index 000000000000..8d21b5092336 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/indexerLayerRopeKvAppend.cu @@ -0,0 +1,439 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/glmMoeDsa/indexerLayerRopeKvAppend.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/kernels/decoderMaskedMultiheadAttentionUtils.h" + +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa::indexer_layer_front +{ + +using glm_moe_dsa::rope_kv_append::kKvLoraRank; +using glm_moe_dsa::rope_kv_append::kLatentDim; +using glm_moe_dsa::rope_kv_append::kRopeDim; +using glm_moe_dsa::rope_kv_append::RopeAppendParams; + +namespace +{ + +constexpr int32_t kRopeThreads = 256; +constexpr int32_t kEltsPerVec = 8; // bf16 elements per 16-byte load +constexpr int32_t kNopeVecs = kKvLoraRank / kEltsPerVec; // 64 +constexpr int32_t kRopeVecs = kRopeDim / kEltsPerVec; // 8 +constexpr int32_t kTotalVecsPerHead = kNopeVecs + kRopeVecs; // 72 +constexpr int32_t kIdxBytesPerItem = 4; + +__device__ __forceinline__ void pdlWait() +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif +} + +__device__ __forceinline__ void pdlLaunchDependents() +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); +#endif +} + +//! Same arithmetic as quantCopy<__nv_bfloat16, 8> in mlaKernels.cu: bf16 -> fp32, scale, saturating E4M3. +__device__ __forceinline__ uint2 quantizeBf16x8(uint4 const& src, float scale) +{ + __nv_bfloat162 const* in = reinterpret_cast<__nv_bfloat162 const*>(&src); + uint2 out; + __nv_fp8x2_e4m3* dst = reinterpret_cast<__nv_fp8x2_e4m3*>(&out); +#pragma unroll + for (int32_t j = 0; j < 4; ++j) + { + float2 v = __bfloat1622float2(in[j]); + v.x *= scale; + v.y *= scale; + dst[j] = __nv_fp8x2_e4m3(v); + } + return out; +} + +//! GPT-J RoPE on 8 bf16 values (4 adjacent pairs), identical to the per-pair transform of the native kernel. +__device__ __forceinline__ void rotateBf16x8(uint4& data, float2 const (&coef)[4]) +{ + __nv_bfloat162* pairs = reinterpret_cast<__nv_bfloat162*>(&data); +#pragma unroll + for (int32_t j = 0; j < 4; ++j) + { + pairs[j] = mmha::rotary_embedding_transform(pairs[j], coef[j]); + } +} + +//! Native unravel of a flat byte index over the strided 4-D indexer cache view (indexerKCacheScatter.cu); +//! 32-bit arithmetic when the index fits, same result. +__device__ __forceinline__ int64_t idxFlatToOffset(int64_t flat, int32_t const (&dims)[4], int64_t const (&strides)[4]) +{ + int32_t i0, i1, i2, i3; + if (flat >= 0 && flat <= 0x7fffffff) + { + uint32_t f = static_cast(flat); + i3 = static_cast(f % static_cast(dims[3])); + f /= static_cast(dims[3]); + i2 = static_cast(f % static_cast(dims[2])); + f /= static_cast(dims[2]); + i1 = static_cast(f % static_cast(dims[1])); + f /= static_cast(dims[1]); + i0 = static_cast(f); + } + else + { + i3 = flat % dims[3]; + flat /= dims[3]; + i2 = flat % dims[2]; + flat /= dims[2]; + i1 = flat % dims[1]; + flat /= dims[1]; + i0 = flat; + } + return i0 * strides[0] + i1 * strides[1] + i2 * strides[2] + i3 * strides[3]; +} + +enum class RopeItem : int32_t +{ + kNone = 0, + kQNope, + kQRope, + kKRope, + kKNope, + kIdxData, + kIdxScale, +}; + +//! One thread per 16-byte work item of one token. Per token the item space is +//! [0, H*64) q_nope FP8 quantization (head = i / 64, vec = i % 64) +//! [H*64, H*72) q_pe RoPE + FP8 (head, vec) +//! [H*72, H*72+8) k_pe RoPE + FP8 cache write (vec) +//! [H*72+8, H*72+72) compressed-KV FP8 cache write (vec) +//! then idxHeadBytes/4 indexer K data items (4 bytes each) and one indexer scale item. +//! Per-step metadata (cache lengths, KV block offsets, indexer slot mappings) and the item's source data are read only +//! after griddepcontrol.wait (see the header). Identical to ropeAppendScatterKernel except for the trigger position. +template +__global__ void __launch_bounds__(kRopeThreads) ropeAppendScatterZ14Kernel(RopeAppendParams const p) +{ + int32_t const token = blockIdx.x; + int32_t const item = blockIdx.y * kRopeThreads + threadIdx.x; + int32_t const qNopeEnd = p.numHeads * kNopeVecs; + int32_t const qRopeEnd = qNopeEnd + p.numHeads * kRopeVecs; + int32_t const kRopeEnd = qRopeEnd + kRopeVecs; + int32_t const kNopeEnd = kRopeEnd + kNopeVecs; + int32_t const idxDataEnd = kNopeEnd + p.idxHeadBytes / kIdxBytesPerItem; + + bool const ropeToken = token < p.numTokens; + bool const idxToken = token < p.idxNumTokens; + + RopeItem kind = RopeItem::kNone; + int32_t head = 0; + int32_t vec = 0; + if (item < qNopeEnd) + { + kind = ropeToken ? RopeItem::kQNope : RopeItem::kNone; + head = item / kNopeVecs; + vec = item % kNopeVecs; + } + else if (item < qRopeEnd) + { + kind = ropeToken ? RopeItem::kQRope : RopeItem::kNone; + head = (item - qNopeEnd) / kRopeVecs; + vec = (item - qNopeEnd) % kRopeVecs; + } + else if (item < kRopeEnd) + { + kind = ropeToken ? RopeItem::kKRope : RopeItem::kNone; + vec = item - qRopeEnd; + } + else if (item < kNopeEnd) + { + kind = ropeToken ? RopeItem::kKNope : RopeItem::kNone; + vec = item - kRopeEnd; + } + else if (item < idxDataEnd) + { + kind = idxToken ? RopeItem::kIdxData : RopeItem::kNone; + vec = item - kNopeEnd; + } + else if (item == idxDataEnd) + { + kind = idxToken ? RopeItem::kIdxScale : RopeItem::kNone; + } + + int32_t const batchIdx = token / p.seqLen; + int32_t const localIdx = token % p.seqLen; + bool const needPos = kind == RopeItem::kQRope || kind == RopeItem::kKRope || kind == RopeItem::kKNope; + bool const needCoef = kind == RopeItem::kQRope || kind == RopeItem::kKRope; + bool const needBlock = kind == RopeItem::kKRope || kind == RopeItem::kKNope; + bool const isIdx = kind == RopeItem::kIdxData || kind == RopeItem::kIdxScale; + + // Values that never change while kernels of the step run (the quantization scales) may be read before the wait; + // per-step metadata and producer outputs may not. + float scale = 1.F; + if (needBlock) + { + scale = p.quantScaleKv ? p.quantScaleKv[0] : 1.F; + } + else if (kind == RopeItem::kQNope || kind == RopeItem::kQRope) + { + scale = p.quantScaleQ ? p.quantScaleQ[0] : 1.F; + } + auto loadCoef = [&](int32_t pos, float2(&c)[4]) + { + float2 const* src = p.cosSin + static_cast(kRopeDim) * pos + (vec * kEltsPerVec) / 2; +#pragma unroll + for (int32_t j = 0; j < 4; ++j) + { + c[j] = src[j]; + } + }; + auto idxOffset = [&](int64_t slotData, int64_t slotScale) -> int64_t + { + if (slotData < 0 || slotScale < 0) + { + return -1; + } + return kind == RopeItem::kIdxData ? idxFlatToOffset(slotData + vec * kIdxBytesPerItem, p.idxDims, p.idxStrides) + : idxFlatToOffset(slotScale, p.idxDims, p.idxStrides); + }; + // Source data of this item; written by producers, so it is read after the wait. + uint4 vdata = make_uint4(0, 0, 0, 0); + uint32_t wdata = 0; + auto loadSource = [&]() + { + switch (kind) + { + case RopeItem::kQNope: + vdata = *reinterpret_cast( + p.fusedQ + (static_cast(token) * p.numHeads + head) * kLatentDim + vec * kEltsPerVec); + break; + case RopeItem::kQRope: + vdata = *reinterpret_cast(p.qPe + static_cast(token) * p.qPeStride + + static_cast(p.qPeLd) * head + vec * kEltsPerVec); + break; + case RopeItem::kKRope: + vdata = *reinterpret_cast( + p.latent + static_cast(token) * kLatentDim + kKvLoraRank + vec * kEltsPerVec); + break; + case RopeItem::kKNope: + vdata = *reinterpret_cast( + p.latent + static_cast(token) * kLatentDim + vec * kEltsPerVec); + break; + case RopeItem::kIdxData: + wdata = *reinterpret_cast( + p.idxK + static_cast(token) * p.idxHeadBytes + vec * kIdxBytesPerItem); + break; + case RopeItem::kIdxScale: + wdata = *reinterpret_cast(p.idxScale + static_cast(token) * kIdxBytesPerItem); + break; + default: break; + } + }; + + int32_t kvPos = 0; + float2 coef[4]; + __nv_fp8_e4m3* kBlock = nullptr; + int64_t idxDst = -1; + pdlWait(); + // Issue the source-data load first so it overlaps the metadata loads and the address math below. + loadSource(); + if (needPos) + { + kvPos = p.kvLens[batchIdx] - p.seqLen + localIdx; + } + if (needCoef) + { + loadCoef(kvPos, coef); + } + if (needBlock) + { + kBlock = reinterpret_cast<__nv_fp8_e4m3*>(p.kvCache.getKBlockPtr(batchIdx, kvPos)); + } + if (isIdx) + { + idxDst = idxOffset(p.idxSlotData[token], p.idxSlotScale[token]); + } + if constexpr (kEarlyTrigger) + { + // Every global read of this thread's item has been issued. The block counts as triggered at the first of + // its threads to get here (later invocations in the block have no further effect), so a dependent may be + // resident before this block's stores below and before block (0, 0)'s scheduler-buffer stores and their + // cache-length reads. Dependents read this kernel's outputs only after their own griddepcontrol.wait, and + // write no global memory before it, so neither the stores nor the reads below can race with them. + pdlLaunchDependents(); + } + + switch (kind) + { + case RopeItem::kQNope: + { + size_t const off = (static_cast(token) * p.numHeads + head) * kLatentDim + vec * kEltsPerVec; + *reinterpret_cast(p.quantQ + off) = quantizeBf16x8(vdata, scale); + break; + } + case RopeItem::kQRope: + { + uint4 data = vdata; + rotateBf16x8(data, coef); + size_t const dst + = (static_cast(token) * p.numHeads + head) * kLatentDim + kKvLoraRank + vec * kEltsPerVec; + *reinterpret_cast(p.quantQ + dst) = quantizeBf16x8(data, scale); + break; + } + case RopeItem::kKRope: + { + uint4 data = vdata; + rotateBf16x8(data, coef); + int32_t const inBlockIdx = p.kvCache.getKVLocalIdx(kvPos, 0, kTotalVecsPerHead, kNopeVecs + vec); + *reinterpret_cast(kBlock + inBlockIdx * kEltsPerVec) = quantizeBf16x8(data, scale); + break; + } + case RopeItem::kKNope: + { + int32_t const inBlockIdx = p.kvCache.getKVLocalIdx(kvPos, 0, kTotalVecsPerHead, vec); + *reinterpret_cast(kBlock + inBlockIdx * kEltsPerVec) = quantizeBf16x8(vdata, scale); + if (vec == 0 && localIdx == 0) + { + p.cuQSeqlens[batchIdx + 1] = p.numHeads * p.seqLen * (batchIdx + 1); + } + break; + } + case RopeItem::kIdxData: + case RopeItem::kIdxScale: + { + if (idxDst >= 0) + { + *reinterpret_cast(p.idxCache + idxDst) = wdata; + } + break; + } + default: break; + } + + if (blockIdx.x == 0 && blockIdx.y == 0) + { + if (threadIdx.x == 0) + { + p.fmhaTileCounter[0] = 0; + p.cuQSeqlens[0] = 0; + float const dequantQ = p.dequantScaleQ ? p.dequantScaleQ[0] : 1.F; + float const dequantKv = p.dequantScaleKv ? p.dequantScaleKv[0] : 1.F; + float const quantO = p.quantScaleO ? p.quantScaleO[0] : 1.F; + if (p.bmm1Scale) + { + constexpr float kLog2e = 1.4426950408889634074F; + float const bmm1 = dequantQ * dequantKv * p.hostBmm1Scale; + p.bmm1Scale[0] = bmm1; + p.bmm1Scale[1] = bmm1 * kLog2e; + } + if (p.bmm2Scale) + { + p.bmm2Scale[0] = quantO * dequantKv; + } + } + if (threadIdx.x < 32) + { + // Exclusive prefix sum of the cache lengths -> cuKvSeqlens[0..batch]. + int32_t const lane = threadIdx.x; + int32_t const batch = p.numTokens / p.seqLen; + int32_t carry = 0; + for (int32_t base = 0; base <= batch; base += 32) + { + int32_t const idx = base + lane; + int32_t const len = idx < batch ? p.kvLens[idx] : 0; + int32_t incl = len; +#pragma unroll + for (int32_t off = 1; off < 32; off <<= 1) + { + int32_t const n = __shfl_up_sync(0xffffffffU, incl, off); + if (lane >= off) + { + incl += n; + } + } + if (idx <= batch) + { + p.cuKvSeqlens[idx] = incl - len + carry; + } + carry += __shfl_sync(0xffffffffU, incl, 31); + } + } + } + + if constexpr (!kEarlyTrigger) + { + // ropeAppendScatterKernel's position. + pdlLaunchDependents(); + } +} + +template +void launchWithOptionalPdl(Kernel kernel, dim3 grid, dim3 block, bool enablePdl, cudaStream_t stream, Params const& p) +{ + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = enablePdl ? 1 : 0; + config.numAttrs = 1; + config.attrs = attrs; + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, kernel, p)); +} + +} // namespace + +void invokeIndexerLayerRopeKvAppend( + RopeAppendParams const& params, bool enablePdl, bool earlyTrigger, cudaStream_t stream) +{ + int32_t const gridTokens = std::max(params.numTokens, params.idxNumTokens); + if (gridTokens == 0) + { + return; + } + TLLM_CHECK_WITH_INFO(params.seqLen > 0, "seqLen must be positive"); + TLLM_CHECK_WITH_INFO(params.cosSinRows > 0, "the RoPE table must have at least one position"); + TLLM_CHECK_WITH_INFO(params.idxHeadBytes % kIdxBytesPerItem == 0, "indexer head bytes must be a multiple of 4"); + int32_t const itemsPerToken + = params.numHeads * kTotalVecsPerHead + kTotalVecsPerHead + params.idxHeadBytes / kIdxBytesPerItem + 1; + dim3 const grid(gridTokens, (itemsPerToken + kRopeThreads - 1) / kRopeThreads); + dim3 const block(kRopeThreads); + if (earlyTrigger) + { + launchWithOptionalPdl(ropeAppendScatterZ14Kernel, grid, block, enablePdl, stream, params); + } + else + { + launchWithOptionalPdl(ropeAppendScatterZ14Kernel, grid, block, enablePdl, stream, params); + } +} + +} // namespace kernels::glm_moe_dsa::indexer_layer_front + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/indexerLayerRopeKvAppend.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/indexerLayerRopeKvAppend.h new file mode 100644 index 000000000000..20a8e8b530de --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/indexerLayerRopeKvAppend.h @@ -0,0 +1,47 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" +#include "tensorrt_llm/kernels/glmMoeDsa/ropeKvAppend.h" + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa::indexer_layer_front +{ + +//! Fork of the generation rope of ropeKvAppend.h (ropeAppendScatterKernel): the same bytes for the same parameters +//! (glm_moe_dsa::rope_kv_append::RopeAppendParams), with the programmatic-launch trigger moved. +//! +//! Memory order: like ropeAppendScatterKernel, only the quantization scales are read before griddepcontrol.wait; every +//! producer output and every per-step metadata word is read after it. With earlyTrigger every thread executes +//! griddepcontrol.launch_dependents right after the global reads of its work item (source item, cache length, rotary +//! coefficients, KV block offset and indexer slot mapping), before its stores. A block counts as triggered at its first +//! such invocation, so dependents may become resident before any store of the kernel, including block (0, 0)'s +//! scheduler buffers (cuQSeqlens, cuKvSeqlens, the tile counter, the bmm scales) and the cache-length reads of its +//! prefix sum. Without earlyTrigger the trigger is at the end of the kernel, as there. A dependent must not read +//! anything this kernel writes, nor write anything it reads, before its own griddepcontrol.wait; the rope's dependents +//! already require this, since a trigger at the end of the code does not make the stores visible either. +void invokeIndexerLayerRopeKvAppend(glm_moe_dsa::rope_kv_append::RopeAppendParams const& params, bool enablePdl, + bool earlyTrigger, cudaStream_t stream); + +} // namespace kernels::glm_moe_dsa::indexer_layer_front + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHop.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHop.cu new file mode 100644 index 000000000000..5f85a965e1d7 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHop.cu @@ -0,0 +1,680 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/glmMoeDsa/postAttentionHop.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa::post_attention_hop +{ +namespace +{ +namespace cg = cooperative_groups; + +// Work decomposition. +// +// Every cluster has kSlices CTAs; CTA c owns slice [c * kSlice, (c + 1) * kSlice), the slice a CTA of the native +// one-shot kernel owns for numRows <= 8 (96 threads x 8 elements, 8 CTAs per row), so the RMSNorm reduction below +// reproduces the native tree exactly. +// +// Without router logits (dense layers) the layout is the native one: cluster y handles row y, and every CTA collects +// one row slice (8 ranks x 768 elements). +// +// With router logits there are kG x kRowSplit clusters (kG = 4, kRowSplit = min(R, 3): at most 12 clusters, which stay +// co-resident). Cluster (g, h) collects and normalizes the rows r with r % kRowSplit == h, so a CTA reads at most +// ceil(R / kRowSplit) row slices, and computes the logits of those rows for experts [64 g, 64 g + 64). Its router +// weights live in registers (loaded before the PDL wait). Each CTA multiplies its rows of slice c by the experts' +// weights of slice c with BF16 MMAs (experts on m16, rows on n8) and pushes the partial logits to the CTA that owns the +// expert; owners add the 8 slices in a fixed order, so the logits are identical on every rank. Clusters never +// synchronize with each other; inside a cluster, slice sums and partial logits travel as st.async stores that complete +// transactions on the consumer's mbarrier. +constexpr int kSlices = 8; +constexpr int kSlice = kHiddenDim / kSlices; +constexpr int kEltsPerVec = 8; +constexpr int kVecPerSlice = kSlice / kEltsPerVec; +constexpr int kNativeWarpsPerSlice = kVecPerSlice / 32; +// Row stride of the normalized-slice tile in 16-byte vectors; 100 = 96 + 4 keeps the 16-byte B-fragment loads of two +// rows in one quarter-warp phase on disjoint banks. +constexpr int kNormedStride = kVecPerSlice + 4; +constexpr int kMmaChains = 4; +constexpr uint32_t kLamportSentinel = 0x80000000U; + +static_assert(kHiddenDim % (kSlices * kEltsPerVec * 32) == 0); +static_assert(kVecPerSlice == kNativeWarpsPerSlice * 32); +static_assert(kMaxRows == 8, "the router MMA maps rows onto the n8 dimension"); + +// kG = number of expert groups (0: hop without router logits); kRowSplit = number of row groups. +template +struct HopConfig +{ + static constexpr int kThreads = 384; + static constexpr int kWarps = kThreads / 32; + static constexpr int kExpertsPerCluster = kNumExperts / kG; + static constexpr int kExpertTiles = kExpertsPerCluster / 16; + static constexpr int kKGroups = kWarps / kExpertTiles; + static constexpr int kKPerGroup = kSlice / kKGroups; + static constexpr int kPairsPerGroup = kKPerGroup / 32; + static constexpr int kExpertsPerOwner = kExpertsPerCluster / kSlices; + static_assert(kExpertsPerCluster % 16 == 0 && kWarps % kExpertTiles == 0 && kSlice % (kKGroups * 32) == 0); + static_assert(kPairsPerGroup % (kMmaChains / 2) == 0); +}; + +template +struct HopConfig<0, kRowSplit> +{ + static constexpr int kThreads = 128; + static constexpr int kWarps = kThreads / 32; + static constexpr int kMaxRowsPerCta = 1; + static_assert(kMaxRowsPerCta * kVecPerSlice <= kThreads && kMaxRowsPerCta * kNativeWarpsPerSlice <= kWarps); +}; + +constexpr int kMaxRowsPerCtaAll = kMaxRows / 2; + +struct alignas(16) HopStorage +{ + uint4 gamma[kVecPerSlice]; + float vecSq[kMaxRowsPerCtaAll][kVecPerSlice]; + float warpSq[kMaxRowsPerCtaAll][kNativeWarpsPerSlice]; + float blockSq[kMaxRowsPerCtaAll][kSlices]; + // Completes when all slice sums of this CTA's rows have arrived (st.async complete_tx). + alignas(8) uint64_t normBar; + // Completes when all partial logits of the experts this CTA owns have arrived. + alignas(8) uint64_t routerBar; +}; + +template +struct alignas(16) RouterStorage +{ + using Cfg = HopConfig; + uint4 normed[kMaxRows][kNormedStride]; // local row on the MMA n dimension; unused rows stay zero + float partial[kSlices][Cfg::kKGroups][Cfg::kExpertsPerOwner][kMaxRows]; +}; + +struct EmptyStorage +{ +}; + +#ifdef TLLM_GLM_LL_W2_TIMING +// Development-only phase timestamps (globaltimer, ns) of thread 0 of every CTA for the most recent launch. +__device__ unsigned long long gW2Timing[kTimingMaxCtas][kTimingStamps]; +#define W2_STAMP(i) \ + do \ + { \ + if (threadIdx.x == 0) \ + { \ + unsigned long long t; \ + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t)); \ + gW2Timing[blockIdx.y * gridDim.x + blockIdx.x][i] = t; \ + } \ + } while (0) +#else +#define W2_STAMP(i) \ + do \ + { \ + } while (0) +#endif + +union BfVec +{ + uint4 packed; + __nv_bfloat16 elts[kEltsPerVec]; +}; + +__device__ __forceinline__ uint4 loadVolatile(void const* ptr) +{ + uint4 v; + asm volatile("ld.volatile.global.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(ptr) + : "memory"); + return v; +} + +__device__ __forceinline__ uint4 loadConstant(void const* ptr) +{ + uint4 v; + asm volatile("ld.global.nc.v4.u32 {%0, %1, %2, %3}, [%4];" : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) : "l"(ptr)); + return v; +} + +__device__ __forceinline__ bool isDirty(uint4 const& v) +{ + return v.x == kLamportSentinel || v.y == kLamportSentinel || v.z == kLamportSentinel || v.w == kLamportSentinel; +} + +// Lamport payloads must never contain the sentinel word: map bf16 -0.0 to +0.0 (the reduction starts from +0.0, so +// the sum is unchanged). +__device__ __forceinline__ uint4 sanitize(uint4 v) +{ + BfVec b; + b.packed = v; +#pragma unroll + for (int i = 0; i < kEltsPerVec; ++i) + { + if (__bfloat16_as_ushort(b.elts[i]) == 0x8000U) + { + b.elts[i] = __ushort_as_bfloat16(0); + } + } + return b.packed; +} + +__device__ __forceinline__ uint32_t sharedAddr(void const* ptr) +{ + return static_cast(__cvta_generic_to_shared(ptr)); +} + +// Address of the same shared-memory object in CTA `rank` of the cluster. +__device__ __forceinline__ uint32_t clusterAddr(uint32_t addr, uint32_t rank) +{ + uint32_t out; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(out) : "r"(addr), "r"(rank)); + return out; +} + +__device__ __forceinline__ void mbarrierInitExpectTx(uint64_t* bar, uint32_t bytes) +{ + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(sharedAddr(bar)) : "memory"); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(sharedAddr(bar)), "r"(bytes) + : "memory"); +} + +__device__ __forceinline__ void mbarrierWaitPhase0(uint64_t* bar) +{ + asm volatile( + "{\n" + ".reg .pred done;\n" + "HOP_MBAR_WAIT_%=:\n" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 done, [%0], 0;\n" + "@!done bra HOP_MBAR_WAIT_%=;\n" + "}\n" ::"r"(sharedAddr(bar)) + : "memory"); +} + +// Remote shared-memory stores that complete transactions on the destination CTA's mbarrier. +__device__ __forceinline__ void storeAsync(uint32_t remoteAddr, float v, uint32_t remoteBar) +{ + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.f32 [%0], %1, [%2];" ::"r"(remoteAddr), "f"(v), + "r"(remoteBar) + : "memory"); +} + +__device__ __forceinline__ void storeAsync(uint32_t remoteAddr, float a, float b, uint32_t remoteBar) +{ + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.v2.f32 [%0], {%1, %2}, [%3];" ::"r"(remoteAddr), + "f"(a), "f"(b), "r"(remoteBar) + : "memory"); +} + +__device__ __forceinline__ void mmaBf16( + float (&d)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t b0, uint32_t b1) +{ + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, " + "{%0, %1, %2, %3};" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)); +} + +template +__global__ void __launch_bounds__(HopConfig::kThreads, 1) + postAttnHopKernel(PostAttnHopParams const params) +{ + static_assert(R >= 1 && R <= kMaxRows); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + constexpr bool kRouter = kG > 0; + using Cfg = HopConfig; + constexpr int kClusters = kRouter ? kG * kRowSplit : R; + constexpr int kGridCtas = kClusters * kSlices; + // Rows one CTA collects, normalizes and multiplies. + constexpr int kRowsPerCta = kRouter ? (R + kRowSplit - 1) / kRowSplit : 1; + static_assert(kRowsPerCta <= kMaxRowsPerCtaAll); + static_assert(kRowsPerCta * kVecPerSlice <= Cfg::kThreads && kRowsPerCta * kNativeWarpsPerSlice <= Cfg::kWarps); + if constexpr (kRouter) + { + static_assert( + HopConfig<(kG > 0 ? kG : 4), kRowSplit>::kExpertsPerOwner * kRowsPerCta * kSlices <= Cfg::kThreads); + } + + __shared__ HopStorage hop; + __shared__ std::conditional_t 0 ? kG : 4), (kG > 0 ? kRowSplit : 2)>, EmptyStorage> rs; + + cg::cluster_group cluster = cg::this_cluster(); + int const slice = static_cast(cluster.block_rank()); + int const clusterId = blockIdx.y; + int const expertGroup = kRouter ? clusterId % (kG > 0 ? kG : 1) : 0; + int const rowGroup = kRouter ? clusterId / (kG > 0 ? kG : 1) : clusterId; + int const tid = threadIdx.x; + int const lane = tid & 31; + int const warp = tid >> 5; + int const sliceBase = slice * kSlice; + + // Rows this CTA collects and normalizes: local row j is global row rowGroup + j * rowStride. + constexpr int kRowStride = kRouter ? kRowSplit : 1; + int const numRowsCta = kRouter ? (R - rowGroup + kRowSplit - 1) / kRowSplit : 1; + // This thread's item: one 16-byte vector of one row slice. + int const itemRow = tid / kVecPerSlice; + int const itemVec = tid % kVecPerSlice; + bool const hasItem = itemRow < numRowsCta; + int const row = rowGroup + itemRow * kRowStride; + // The rank's own partial of a row slice is published by exactly one CTA of the grid. + bool const sends = hasItem && (!kRouter || itemRow % (kG > 0 ? kG : 1) == expertGroup); + size_t const rowOff = static_cast(row) * kHiddenDim + sliceBase + itemVec * kEltsPerVec; + // The cluster that writes the normalized rows and the residual. + bool const writesOutputs = !kRouter || expertGroup == 0; + + // ---- Prologue: constant model weights only; may overlap the producer of `input` under PDL. ---- + if (tid < kVecPerSlice) + { + hop.gamma[tid] = loadConstant(reinterpret_cast(params.gamma) + slice * kVecPerSlice + tid); + } + [[maybe_unused]] int const gid = lane >> 2; + [[maybe_unused]] int const tig = lane & 3; + [[maybe_unused]] uint4 wLo[kRouter ? HopConfig<(kG > 0 ? kG : 4), (kG > 0 ? kRowSplit : 2)>::kPairsPerGroup : 1]; + [[maybe_unused]] uint4 wHi[kRouter ? HopConfig<(kG > 0 ? kG : 4), (kG > 0 ? kRowSplit : 2)>::kPairsPerGroup : 1]; + if constexpr (kRouter) + { + int const tile = warp % Cfg::kExpertTiles; + int const kGroup = warp / Cfg::kExpertTiles; + auto const* w = reinterpret_cast<__nv_bfloat16 const*>(params.routerWeight); + int const expert = expertGroup * Cfg::kExpertsPerCluster + tile * 16 + gid; +#pragma unroll + for (int p = 0; p < Cfg::kPairsPerGroup; ++p) + { + int const k = sliceBase + kGroup * Cfg::kKPerGroup + p * 32 + tig * kEltsPerVec; + wLo[p] = loadConstant(w + static_cast(expert) * kHiddenDim + k); + wHi[p] = loadConstant(w + static_cast(expert + 8) * kHiddenDim + k); + } + // MMA columns of rows this CTA does not own stay zero. + for (int i = tid; i < kMaxRows * kNormedStride; i += Cfg::kThreads) + { + rs.normed[i / kNormedStride][i % kNormedStride] = make_uint4(0, 0, 0, 0); + } + } + + // Cluster-local completion barriers: every CTA of a cluster has the same rows, so each expects the 8 slice sums + // of its rows and, with the router, the 8 x kKGroups partial tiles of the experts it owns. A split cluster + // barrier publishes the initialization; its wait precedes the first remote store, long after all CTAs arrived. + if (tid == 0) + { + mbarrierInitExpectTx(&hop.normBar, static_cast(numRowsCta * kSlices * sizeof(float))); + if constexpr (kRouter) + { + mbarrierInitExpectTx(&hop.routerBar, numRowsCta > 0 ? static_cast(sizeof(rs.partial)) : 0U); + } + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + + // ---- Wait for the producer of `input` (and transitively of `residualIn` and the protocol state). ---- + W2_STAMP(0); + asm volatile("griddepcontrol.wait;" ::: "memory"); + W2_STAMP(1); + // Dependents only start their own prologue; they still wait for this grid to complete. + asm volatile("griddepcontrol.launch_dependents;"); + + uint4 const flags = loadVolatile(params.flags); + auto const* input = reinterpret_cast<__nv_bfloat16 const*>(params.input); + auto const* residualIn = reinterpret_cast<__nv_bfloat16 const*>(params.residualIn); + uint4 res = make_uint4(0, 0, 0, 0); + // This rank's own partial, as published (sanitized); it is used from registers, not from the multicast copy. + uint4 own = make_uint4(0, 0, 0, 0); + if (hasItem) + { + res = *reinterpret_cast(residualIn + rowOff); + own = sanitize(*reinterpret_cast(input + rowOff)); + } + uint32_t const cur = flags.x; + uint32_t const clearBytes = flags.y; + char* ucBase = reinterpret_cast(params.ucBuffer); + auto* curUc = reinterpret_cast<__nv_bfloat16*>(ucBase + cur * kLamportBufferBytes); + auto* curMc + = reinterpret_cast<__nv_bfloat16*>(reinterpret_cast(params.mcBuffer) + cur * kLamportBufferBytes); + // Lamport layout of one buffer: [row][source rank][kHiddenDim]. + size_t const slotBase = static_cast(row) * kNumRanks * kHiddenDim + sliceBase + itemVec * kEltsPerVec; + + // ---- Publish this rank's partial to every rank (one multicast store per vector). ---- + if (sends) + { + *reinterpret_cast(curMc + slotBase + static_cast(params.rank) * kHiddenDim) = own; + } + + W2_STAMP(2); + // ---- Collect the other ranks' partials. ---- + uint4 vals[kNumRanks]; + if (hasItem) + { +#pragma unroll + for (int s = 0; s < kNumRanks; ++s) + { + vals[s] = s == params.rank ? own : loadVolatile(curUc + slotBase + static_cast(s) * kHiddenDim); + } + } + + // ---- Reset the previous call's buffer to the sentinel (reused two calls from now), off the poll's path. ---- + { + uint32_t const dirty = (cur + kNumLamportBuffers - 1) % kNumLamportBuffers; + auto* dirtyBuf = reinterpret_cast(ucBase + dirty * kLamportBufferBytes); + uint32_t const numVec = clearBytes / sizeof(uint4); + uint4 const sentinel = make_uint4(kLamportSentinel, kLamportSentinel, kLamportSentinel, kLamportSentinel); + uint32_t const ctaLinear = clusterId * kSlices + slice; + for (uint32_t j = ctaLinear * Cfg::kThreads + tid; j < numVec; j += kGridCtas * Cfg::kThreads) + { + dirtyBuf[j] = sentinel; + } + } + + if (hasItem) + { + bool pending = true; + while (pending) + { + pending = false; +#pragma unroll + for (int s = 0; s < kNumRanks; ++s) + { + if (isDirty(vals[s])) + { + pending = true; + vals[s] = loadVolatile(curUc + slotBase + static_cast(s) * kHiddenDim); + } + } + } + } + + W2_STAMP(3); + // ---- Reduce in rank order, add the residual, and form per-vector sums of squares (native arithmetic). ---- + uint4 resOut = make_uint4(0, 0, 0, 0); + if (hasItem) + { + float acc[kEltsPerVec]; +#pragma unroll + for (int e = 0; e < kEltsPerVec; ++e) + { + acc[e] = 0.F; + } +#pragma unroll + for (int s = 0; s < kNumRanks; ++s) + { + BfVec x; + x.packed = vals[s]; +#pragma unroll + for (int e = 0; e < kEltsPerVec; ++e) + { + acc[e] += __bfloat162float(x.elts[e]); + } + } + BfVec rv; + rv.packed = res; + BfVec y; + float sq = 0.F; +#pragma unroll + for (int e = 0; e < kEltsPerVec; ++e) + { + y.elts[e] = __hadd(__float2bfloat16(acc[e]), rv.elts[e]); + sq += __bfloat162float(__hmul(y.elts[e], y.elts[e])); + } + resOut = y.packed; + hop.vecSq[itemRow][itemVec] = sq; + } + __syncthreads(); + if (tid == 0) + { + // Every thread of this CTA has read the protocol state; CTA 0 advances it only after all CTAs report. + asm volatile("red.release.gpu.global.add.u32 [%0], %1;" ::"l"(params.flags + 2), "r"(1U) : "memory"); + } + + // ---- RMS: warp butterfly over 32 vectors, then (w0 + w2) + w1 per slice, then the 8 slices in order. ---- + if (warp < numRowsCta * kNativeWarpsPerSlice) + { + int const j = warp / kNativeWarpsPerSlice; + int const q = warp % kNativeWarpsPerSlice; + float val = hop.vecSq[j][q * 32 + lane]; +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) + { + val = val + __shfl_xor_sync(0xffffffffU, val, mask, 32); + } + if (lane == 0) + { + hop.warpSq[j][q] = val; + } + } + __syncthreads(); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + if (tid < numRowsCta * kSlices) + { + int const j = tid / kSlices; + int const dst = tid % kSlices; + float const blockSum = (hop.warpSq[j][0] + hop.warpSq[j][2]) + hop.warpSq[j][1]; + storeAsync( + clusterAddr(sharedAddr(&hop.blockSq[j][slice]), dst), blockSum, clusterAddr(sharedAddr(&hop.normBar), dst)); + } + W2_STAMP(4); + mbarrierWaitPhase0(&hop.normBar); + W2_STAMP(5); + + if (hasItem) + { + float fullSum = 0.F; +#pragma unroll + for (int c = 0; c < kSlices; ++c) + { + fullSum += hop.blockSq[itemRow][c]; + } + float const rcpRms = rsqrtf(fullSum / static_cast(kHiddenDim) + params.eps); + BfVec y; + y.packed = resOut; + BfVec g; + g.packed = hop.gamma[itemVec]; + BfVec n; +#pragma unroll + for (int e = 0; e < kEltsPerVec; ++e) + { + n.elts[e] = __float2bfloat16(__bfloat162float(y.elts[e]) * rcpRms * __bfloat162float(g.elts[e])); + } + if (writesOutputs) + { + *reinterpret_cast(reinterpret_cast<__nv_bfloat16*>(params.normOut) + rowOff) = n.packed; + *reinterpret_cast(reinterpret_cast<__nv_bfloat16*>(params.residualOut) + rowOff) = y.packed; + } + if constexpr (kRouter) + { + rs.normed[itemRow][itemVec] = n.packed; + } + } + + if constexpr (kRouter) + { + if (numRowsCta > 0) + { + __syncthreads(); + W2_STAMP(6); + int const tile = warp % Cfg::kExpertTiles; + int const kGroup = warp / Cfg::kExpertTiles; + // Independent accumulator chains keep the MMA latency off a single dependency chain. + float acc[kMmaChains][4] = {}; +#pragma unroll + for (int p = 0; p < Cfg::kPairsPerGroup; ++p) + { + // One 16-byte load covers two k16 steps; the k-permutation inside a step is the same for both operands. + uint4 const b = rs.normed[gid][kGroup * (Cfg::kKPerGroup / kEltsPerVec) + p * 4 + tig]; + int const chain = (2 * p) % kMmaChains; + mmaBf16(acc[chain], wLo[p].x, wHi[p].x, wLo[p].y, wHi[p].y, b.x, b.y); + mmaBf16(acc[chain + 1], wLo[p].z, wHi[p].z, wLo[p].w, wHi[p].w, b.z, b.w); + } + float sum4[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + sum4[i] = (acc[0][i] + acc[1][i]) + (acc[2][i] + acc[3][i]); + } + // Expert e (cluster-local) is owned by CTA e / kExpertsPerOwner. + int const e0 = tile * 16 + gid; + int const e1 = e0 + 8; + uint32_t const owner0 = e0 / Cfg::kExpertsPerOwner; + uint32_t const owner1 = e1 / Cfg::kExpertsPerOwner; + storeAsync(clusterAddr(sharedAddr(&rs.partial[slice][kGroup][e0 % Cfg::kExpertsPerOwner][tig * 2]), owner0), + sum4[0], sum4[1], clusterAddr(sharedAddr(&hop.routerBar), owner0)); + storeAsync(clusterAddr(sharedAddr(&rs.partial[slice][kGroup][e1 % Cfg::kExpertsPerOwner][tig * 2]), owner1), + sum4[2], sum4[3], clusterAddr(sharedAddr(&hop.routerBar), owner1)); + W2_STAMP(7); + // Also keeps this CTA resident until every partial addressed to it has landed. + mbarrierWaitPhase0(&hop.routerBar); + W2_STAMP(8); + + // This CTA owns experts [slice * kExpertsPerOwner, (slice + 1) * kExpertsPerOwner) of the cluster. Lane + // groups of 8 sum over the 8 slices. + if (tid < Cfg::kExpertsPerOwner * numRowsCta * kSlices) + { + int const c = tid % kSlices; + int const ij = tid / kSlices; + int const ei = ij / numRowsCta; + int const j = ij % numRowsCta; + float sum = 0.F; +#pragma unroll + for (int kg = 0; kg < Cfg::kKGroups; ++kg) + { + sum += rs.partial[c][kg][ei][j]; + } +#pragma unroll + for (int mask = kSlices / 2; mask > 0; mask >>= 1) + { + sum = sum + __shfl_xor_sync(0xffffffffU, sum, mask, kSlices); + } + if (c == 0) + { + int const r = rowGroup + j * kRowStride; + params.routerLogits[r * kNumExperts + expertGroup * Cfg::kExpertsPerCluster + + slice * Cfg::kExpertsPerOwner + ei] + = sum; + } + } + } + } + + W2_STAMP(9); + // ---- CTA 0 advances the protocol state once every CTA has read it. ---- + if (tid == 0 && slice == 0 && clusterId == 0) + { + uint32_t arrived = 0; + do + { + asm volatile("ld.acquire.gpu.global.u32 %0, [%1];" : "=r"(arrived) : "l"(params.flags + 2) : "memory"); + } while (arrived < static_cast(kGridCtas)); + uint4 const next = make_uint4((cur + 1) % kNumLamportBuffers, + static_cast(R * kNumRanks * kHiddenDim * sizeof(__nv_bfloat16)), 0U, 0U); + asm volatile("st.relaxed.gpu.global.v4.u32 [%0], {%1, %2, %3, %4};" ::"l"(params.flags), "r"(next.x), + "r"(next.y), "r"(next.z), "r"(next.w) + : "memory"); + } +#else + // Clusters, distributed shared memory and PDL need SM90+; the launcher is only reached on supported devices. + __trap(); +#endif +} + +template +cudaError_t launchTyped(PostAttnHopParams const& params) +{ + using Cfg = HopConfig; + cudaLaunchAttribute attrs[2]; + attrs[0].id = cudaLaunchAttributeClusterDimension; + attrs[0].val.clusterDim.x = kSlices; + attrs[0].val.clusterDim.y = 1; + attrs[0].val.clusterDim.z = 1; + attrs[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[1].val.programmaticStreamSerializationAllowed = params.enablePdl ? 1 : 0; + + cudaLaunchConfig_t config{}; + config.gridDim = dim3(kSlices, kG > 0 ? kG * kRowSplit : R, 1); + config.blockDim = dim3(Cfg::kThreads, 1, 1); + config.dynamicSmemBytes = 0; + config.stream = params.stream; + config.attrs = attrs; + config.numAttrs = 2; + return cudaLaunchKernelEx(&config, postAttnHopKernel, params); +} + +template +cudaError_t launchRows(PostAttnHopParams const& params) +{ + switch (params.numRows) + { + case 1: return launchTyped<1, kG, kRowSplit>(params); + case 2: return launchTyped<2, kG, kRowSplit>(params); + case 3: return launchTyped<3, kG, kRowSplit>(params); + case 4: return launchTyped<4, kG, kRowSplit>(params); + case 5: return launchTyped<5, kG, kRowSplit>(params); + case 6: return launchTyped<6, kG, kRowSplit>(params); + case 7: return launchTyped<7, kG, kRowSplit>(params); + case 8: return launchTyped<8, kG, kRowSplit>(params); + default: return cudaErrorInvalidValue; + } +} + +// Router hop: 4 expert groups x min(R, 3) row groups (at most 12 clusters of 8 CTAs, each CTA <= 3 row slices). +cudaError_t launchRouter(PostAttnHopParams const& params) +{ + switch (params.numRows) + { + case 1: return launchTyped<1, 4, 1>(params); + case 2: return launchTyped<2, 4, 2>(params); + case 3: return launchTyped<3, 4, 3>(params); + case 4: return launchTyped<4, 4, 3>(params); + case 5: return launchTyped<5, 4, 3>(params); + case 6: return launchTyped<6, 4, 3>(params); + case 7: return launchTyped<7, 4, 3>(params); + case 8: return launchTyped<8, 4, 3>(params); + default: return cudaErrorInvalidValue; + } +} + +} // namespace + +#ifdef TLLM_GLM_LL_W2_TIMING +cudaError_t copyPostAttnHopTiming(unsigned long long* dst, cudaStream_t stream) +{ + return cudaMemcpyFromSymbolAsync(dst, gW2Timing, sizeof(gW2Timing), 0, cudaMemcpyDeviceToHost, stream); +} +#endif + +cudaError_t invokePostAttentionHop(PostAttnHopParams const& params) +{ + if (params.rank < 0 || params.rank >= kNumRanks || params.flags == nullptr || params.ucBuffer == nullptr + || params.mcBuffer == nullptr) + { + return cudaErrorInvalidValue; + } + if (params.routerWeight != nullptr) + { + if (params.routerLogits == nullptr) + { + return cudaErrorInvalidValue; + } + return launchRouter(params); + } + return launchRows<0, 1>(params); +} + +} // namespace kernels::glm_moe_dsa::post_attention_hop + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHop.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHop.h new file mode 100644 index 000000000000..ee7c62202f5b --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHop.h @@ -0,0 +1,88 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa::post_attention_hop +{ + +//! Fused post-attention hop. +//! +//! For ``numRows`` rows of hidden size ``kHiddenDim`` on a TP group of ``kNumRanks`` ranks it computes, on every +//! rank and with identical bits on every rank: +//! sum = bf16(sum over ranks 0..7 of float(partial_rank)) (one-shot Lamport exchange, fixed order) +//! residual = sum + residualIn (bf16 add) +//! normed = bf16(float(residual) * rsqrt(mean(bf16(residual^2)) + eps) * float(gamma)) +//! logits = normed @ routerWeight^T in FP32 (optional, [numRows, kNumExperts]) +//! ``sum``, ``residual`` and ``normed`` follow the arithmetic and reduction tree of the native MNNVL one-shot +//! AR + residual + RMSNorm kernel, so they are bit-identical to it. +//! +//! The exchange uses a private multicast workspace of ``kNumLamportBuffers`` rotating buffers. ``flags`` holds the +//! device-side protocol state ``{current buffer, bytes to clear in the previous buffer, arrival counter, unused}`` +//! and must be zero-initialized together with the buffers filled with the Lamport sentinel (all 32-bit words equal +//! to 0x80000000) on every rank before the first call. +struct PostAttnHopParams +{ + void const* input = nullptr; //!< [numRows, kHiddenDim] bf16, rank-local partial (o_proj output) + void const* residualIn = nullptr; //!< [numRows, kHiddenDim] bf16 + void const* gamma = nullptr; //!< [kHiddenDim] bf16 RMSNorm weight + void const* routerWeight = nullptr; //!< [kNumExperts, kHiddenDim] bf16 or nullptr (no router logits) + void* normOut = nullptr; //!< [numRows, kHiddenDim] bf16 + void* residualOut = nullptr; //!< [numRows, kHiddenDim] bf16 + float* routerLogits = nullptr; //!< [numRows, kNumExperts] fp32, required iff routerWeight != nullptr + void* ucBuffer = nullptr; //!< local unicast view of the multicast workspace + void* mcBuffer = nullptr; //!< multicast address of the same workspace + uint32_t* flags = nullptr; //!< [4] device-side protocol state of this workspace + int numRows = 0; + int rank = 0; + float eps = 1e-5F; + bool enablePdl = true; + cudaStream_t stream = nullptr; +}; + +constexpr int kHiddenDim = 6144; +constexpr int kNumRanks = 8; +constexpr int kNumExperts = 256; +constexpr int kMaxRows = 8; +constexpr int kNumLamportBuffers = 3; + +//! Bytes of one Lamport buffer: [kMaxRows][kNumRanks][kHiddenDim] bf16. +constexpr size_t kLamportBufferBytes = size_t(kMaxRows) * kNumRanks * kHiddenDim * 2; +//! Bytes of the whole multicast workspace. +constexpr size_t kWorkspaceBytes = kLamportBufferBytes * kNumLamportBuffers; + +//! Launch the fused hop on ``params.stream``. Returns the launch status. +cudaError_t invokePostAttentionHop(PostAttnHopParams const& params); + +#ifdef TLLM_GLM_LL_W2_TIMING +//! Development builds only: phase timestamps of the most recent launch, [kTimingMaxCtas][kTimingStamps] ns. +constexpr int kTimingMaxCtas = 128; +constexpr int kTimingStamps = 10; +cudaError_t copyPostAttnHopTiming(unsigned long long* dst, cudaStream_t stream); +#endif + +} // namespace kernels::glm_moe_dsa::post_attention_hop + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopFromPartials.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopFromPartials.cu new file mode 100644 index 000000000000..17c81e24877a --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopFromPartials.cu @@ -0,0 +1,707 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/glmMoeDsa/postAttentionHopFromPartials.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa::attention_output_into_hop +{ +namespace +{ +namespace cg = cooperative_groups; + +// The post-attention hop (glmMoeDsa/postAttentionHop.cu) with its input taken from the transport: the only differences +// are marked "Transport" below; the work decomposition, the arithmetic and the protocol are that hop's. +// +// Work decomposition. +// +// Every cluster has kSlices CTAs; CTA c owns slice [c * kSlice, (c + 1) * kSlice), the slice a CTA of the native +// one-shot kernel owns for numRows <= 8 (96 threads x 8 elements, 8 CTAs per row), so the RMSNorm reduction below +// reproduces the native tree exactly. +// +// Without router logits (dense layers) the layout is the native one: cluster y handles row y, and every CTA collects +// one row slice (8 ranks x 768 elements). +// +// With router logits there are kG x kRowSplit clusters (kG = 4, kRowSplit = min(R, 3): at most 12 clusters, which stay +// co-resident). Cluster (g, h) collects and normalizes the rows r with r % kRowSplit == h, so a CTA reads at most +// ceil(R / kRowSplit) row slices, and computes the logits of those rows for experts [64 g, 64 g + 64). Its router +// weights live in registers (loaded before the PDL wait). Each CTA multiplies its rows of slice c by the experts' +// weights of slice c with BF16 MMAs (experts on m16, rows on n8) and pushes the partial logits to the CTA that owns the +// expert; owners add the 8 slices in a fixed order, so the logits are identical on every rank. Clusters never +// synchronize with each other; inside a cluster, slice sums and partial logits travel as st.async stores that complete +// transactions on the consumer's mbarrier. +constexpr int kSlices = 8; +constexpr int kSlice = kHiddenDim / kSlices; +constexpr int kEltsPerVec = 8; +constexpr int kVecPerSlice = kSlice / kEltsPerVec; +constexpr int kNativeWarpsPerSlice = kVecPerSlice / 32; +// Row stride of the normalized-slice tile in 16-byte vectors; 100 = 96 + 4 keeps the 16-byte B-fragment loads of two +// rows in one quarter-warp phase on disjoint banks. +constexpr int kNormedStride = kVecPerSlice + 4; +constexpr int kMmaChains = 4; +constexpr uint32_t kLamportSentinel = 0x80000000U; + +static_assert(kHiddenDim % (kSlices * kEltsPerVec * 32) == 0); +static_assert(kVecPerSlice == kNativeWarpsPerSlice * 32); +static_assert(kMaxRows == 8, "the router MMA maps rows onto the n8 dimension"); + +// kG = number of expert groups (0: hop without router logits); kRowSplit = number of row groups. +template +struct HopConfig +{ + static constexpr int kThreads = 384; + static constexpr int kWarps = kThreads / 32; + static constexpr int kExpertsPerCluster = kNumExperts / kG; + static constexpr int kExpertTiles = kExpertsPerCluster / 16; + static constexpr int kKGroups = kWarps / kExpertTiles; + static constexpr int kKPerGroup = kSlice / kKGroups; + static constexpr int kPairsPerGroup = kKPerGroup / 32; + static constexpr int kExpertsPerOwner = kExpertsPerCluster / kSlices; + static_assert(kExpertsPerCluster % 16 == 0 && kWarps % kExpertTiles == 0 && kSlice % (kKGroups * 32) == 0); + static_assert(kPairsPerGroup % (kMmaChains / 2) == 0); +}; + +template +struct HopConfig<0, kRowSplit> +{ + static constexpr int kThreads = 128; + static constexpr int kWarps = kThreads / 32; + static constexpr int kMaxRowsPerCta = 1; + static_assert(kMaxRowsPerCta * kVecPerSlice <= kThreads && kMaxRowsPerCta * kNativeWarpsPerSlice <= kWarps); +}; + +constexpr int kMaxRowsPerCtaAll = kMaxRows / 2; + +struct alignas(16) HopStorage +{ + uint4 gamma[kVecPerSlice]; + float vecSq[kMaxRowsPerCtaAll][kVecPerSlice]; + float warpSq[kMaxRowsPerCtaAll][kNativeWarpsPerSlice]; + float blockSq[kMaxRowsPerCtaAll][kSlices]; + // Completes when all slice sums of this CTA's rows have arrived (st.async complete_tx). + alignas(8) uint64_t normBar; + // Completes when all partial logits of the experts this CTA owns have arrived. + alignas(8) uint64_t routerBar; +}; + +template +struct alignas(16) RouterStorage +{ + using Cfg = HopConfig; + uint4 normed[kMaxRows][kNormedStride]; // local row on the MMA n dimension; unused rows stay zero + float partial[kSlices][Cfg::kKGroups][Cfg::kExpertsPerOwner][kMaxRows]; +}; + +struct EmptyStorage +{ +}; + +union BfVec +{ + uint4 packed; + __nv_bfloat16 elts[kEltsPerVec]; +}; + +__device__ __forceinline__ uint4 loadVolatile(void const* ptr) +{ + uint4 v; + asm volatile("ld.volatile.global.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(ptr) + : "memory"); + return v; +} + +__device__ __forceinline__ uint4 loadConstant(void const* ptr) +{ + uint4 v; + asm volatile("ld.global.nc.v4.u32 {%0, %1, %2, %3}, [%4];" : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) : "l"(ptr)); + return v; +} + +__device__ __forceinline__ void prefetchL2(void const* ptr) +{ + asm volatile("prefetch.global.L2 [%0];" ::"l"(ptr)); +} + +__device__ __forceinline__ bool isDirty(uint4 const& v) +{ + return v.x == kLamportSentinel || v.y == kLamportSentinel || v.z == kLamportSentinel || v.w == kLamportSentinel; +} + +// Transport: the producer writes the Lamport payloads already sanitized (bf16 -0.0 mapped to +0.0), so they never +// contain the sentinel word. + +__device__ __forceinline__ uint32_t sharedAddr(void const* ptr) +{ + return static_cast(__cvta_generic_to_shared(ptr)); +} + +// Address of the same shared-memory object in CTA `rank` of the cluster. +__device__ __forceinline__ uint32_t clusterAddr(uint32_t addr, uint32_t rank) +{ + uint32_t out; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(out) : "r"(addr), "r"(rank)); + return out; +} + +__device__ __forceinline__ void mbarrierInitExpectTx(uint64_t* bar, uint32_t bytes) +{ + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(sharedAddr(bar)) : "memory"); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(sharedAddr(bar)), "r"(bytes) + : "memory"); +} + +__device__ __forceinline__ void mbarrierWaitPhase0(uint64_t* bar) +{ + asm volatile( + "{\n" + ".reg .pred done;\n" + "HOP_MBAR_WAIT_%=:\n" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 done, [%0], 0;\n" + "@!done bra HOP_MBAR_WAIT_%=;\n" + "}\n" ::"r"(sharedAddr(bar)) + : "memory"); +} + +// Remote shared-memory stores that complete transactions on the destination CTA's mbarrier. +__device__ __forceinline__ void storeAsync(uint32_t remoteAddr, float v, uint32_t remoteBar) +{ + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.f32 [%0], %1, [%2];" ::"r"(remoteAddr), "f"(v), + "r"(remoteBar) + : "memory"); +} + +__device__ __forceinline__ void storeAsync(uint32_t remoteAddr, float a, float b, uint32_t remoteBar) +{ + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.v2.f32 [%0], {%1, %2}, [%3];" ::"r"(remoteAddr), + "f"(a), "f"(b), "r"(remoteBar) + : "memory"); +} + +__device__ __forceinline__ void mmaBf16( + float (&d)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t b0, uint32_t b1) +{ + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, " + "{%0, %1, %2, %3};" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)); +} + +// Transport: one warp's router-weight MMA fragments (the post-attention hop's register layout): experts +// tile * 16 + gid (+ 8) of this cluster's expert group, k pairs of this warp's k group within the CTA's slice. +template +__device__ __forceinline__ void loadRouterFragments(uint4 (&wLo)[Cfg::kPairsPerGroup], + uint4 (&wHi)[Cfg::kPairsPerGroup], void const* routerWeight, int expertGroup, int warp, int gid, int tig, + int sliceBase) +{ + int const tile = warp % Cfg::kExpertTiles; + int const kGroup = warp / Cfg::kExpertTiles; + auto const* w = reinterpret_cast<__nv_bfloat16 const*>(routerWeight); + int const expert = expertGroup * Cfg::kExpertsPerCluster + tile * 16 + gid; +#pragma unroll + for (int p = 0; p < Cfg::kPairsPerGroup; ++p) + { + int const k = sliceBase + kGroup * Cfg::kKPerGroup + p * 32 + tig * kEltsPerVec; + wLo[p] = loadConstant(w + static_cast(expert) * kHiddenDim + k); + wHi[p] = loadConstant(w + static_cast(expert + 8) * kHiddenDim + k); + } +} + +// Transport: where the router-weight fragments are loaded (all orders give the same values and arithmetic): +// kRouterAfterBarrier: after the barrier that ends the rank sum, overlapping the RMS tree and cluster exchange; +// kRouterAfterSum: right after the rank sum, before that barrier (the default; the loads go through +// loadRouterFragments so the rank sum is not sunk below them and the kernel stays at +// few enough registers to share an SM with a producer CTA); +// kRouterBeforeMma: right before the MMAs (fewest registers, the load latency follows the RMSNorm). +constexpr int kRouterAfterBarrier = 0; +constexpr int kRouterAfterSum = 1; +constexpr int kRouterBeforeMma = 2; + +template +__global__ void __launch_bounds__(HopConfig::kThreads, 1) + postAttnHopX1Kernel(PostAttentionHopFromPartialsParams const params) +{ + static_assert(R >= 1 && R <= kMaxRows); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + constexpr bool kRouter = kG > 0; + using Cfg = HopConfig; + constexpr int kClusters = kRouter ? kG * kRowSplit : R; + constexpr int kGridCtas = kClusters * kSlices; + // Rows one CTA collects, normalizes and multiplies. + constexpr int kRowsPerCta = kRouter ? (R + kRowSplit - 1) / kRowSplit : 1; + static_assert(kRowsPerCta <= kMaxRowsPerCtaAll); + static_assert(kRowsPerCta * kVecPerSlice <= Cfg::kThreads && kRowsPerCta * kNativeWarpsPerSlice <= Cfg::kWarps); + if constexpr (kRouter) + { + static_assert( + HopConfig<(kG > 0 ? kG : 4), kRowSplit>::kExpertsPerOwner * kRowsPerCta * kSlices <= Cfg::kThreads); + } + + __shared__ HopStorage hop; + __shared__ std::conditional_t 0 ? kG : 4), (kG > 0 ? kRowSplit : 2)>, EmptyStorage> rs; + + cg::cluster_group cluster = cg::this_cluster(); + int const slice = static_cast(cluster.block_rank()); + int const clusterId = blockIdx.y; + int const expertGroup = kRouter ? clusterId % (kG > 0 ? kG : 1) : 0; + int const rowGroup = kRouter ? clusterId / (kG > 0 ? kG : 1) : clusterId; + int const tid = threadIdx.x; + int const lane = tid & 31; + int const warp = tid >> 5; + int const sliceBase = slice * kSlice; + + // Rows this CTA collects and normalizes: local row j is global row rowGroup + j * rowStride. + constexpr int kRowStride = kRouter ? kRowSplit : 1; + int const numRowsCta = kRouter ? (R - rowGroup + kRowSplit - 1) / kRowSplit : 1; + // This thread's item: one 16-byte vector of one row slice. + int const itemRow = tid / kVecPerSlice; + int const itemVec = tid % kVecPerSlice; + bool const hasItem = itemRow < numRowsCta; + int const row = rowGroup + itemRow * kRowStride; + size_t const rowOff = static_cast(row) * kHiddenDim + sliceBase + itemVec * kEltsPerVec; + // The cluster that writes the normalized rows and the residual. + bool const writesOutputs = !kRouter || expertGroup == 0; + + // ---- Prologue: constant model weights only; may overlap the producer of `input` under PDL. ---- + if (tid < kVecPerSlice) + { + hop.gamma[tid] = loadConstant(reinterpret_cast(params.gamma) + slice * kVecPerSlice + tid); + } + [[maybe_unused]] int const gid = lane >> 2; + [[maybe_unused]] int const tig = lane & 3; + // Transport: the router-weight fragments are loaded after the poll (not here), so they are not live across it; with + // fewer registers a CTA of this kernel can be resident next to a CTA of the producer. The prologue only asks L2 + // for the same lines (no registers held, no architectural effect). + if constexpr (kRouter) + { + { + int const tile = warp % Cfg::kExpertTiles; + int const kGroup = warp / Cfg::kExpertTiles; + auto const* w = reinterpret_cast<__nv_bfloat16 const*>(params.routerWeight); + int const expert = expertGroup * Cfg::kExpertsPerCluster + tile * 16 + gid; +#pragma unroll + for (int p = 0; p < Cfg::kPairsPerGroup; ++p) + { + int const k = sliceBase + kGroup * Cfg::kKPerGroup + p * 32 + tig * kEltsPerVec; + prefetchL2(w + static_cast(expert) * kHiddenDim + k); + prefetchL2(w + static_cast(expert + 8) * kHiddenDim + k); + } + } + // MMA columns of rows this CTA does not own stay zero. + for (int i = tid; i < kMaxRows * kNormedStride; i += Cfg::kThreads) + { + rs.normed[i / kNormedStride][i % kNormedStride] = make_uint4(0, 0, 0, 0); + } + } + + // Cluster-local completion barriers: every CTA of a cluster has the same rows, so each expects the 8 slice sums + // of its rows and, with the router, the 8 x kKGroups partial tiles of the experts it owns. A split cluster + // barrier publishes the initialization; its wait precedes the first remote store, long after all CTAs arrived. + if (tid == 0) + { + mbarrierInitExpectTx(&hop.normBar, static_cast(numRowsCta * kSlices * sizeof(float))); + if constexpr (kRouter) + { + mbarrierInitExpectTx(&hop.routerBar, numRowsCta > 0 ? static_cast(sizeof(rs.partial)) : 0U); + } + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + + // Transport: with earlyPoll no wait here. The producer read the protocol state and wrote the transport after its + // own griddepcontrol.wait and triggered this grid only after that wait, so every grid before the producer has + // completed and the protocol state written by the previous hop is final (see postAttentionHopFromPartials.h). + auto const* residualIn = reinterpret_cast<__nv_bfloat16 const*>(params.residualIn); + uint4 res = make_uint4(0, 0, 0, 0); + if (!params.earlyPoll) + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + // As in the post-attention hop: dependents may start their prologue; they still wait for this grid to complete. + asm volatile("griddepcontrol.launch_dependents;"); + // Transport: after the wait the residual is final; its load overlaps the poll instead of following it. + if (hasItem) + { + res = *reinterpret_cast(residualIn + rowOff); + } + } + uint4 const flags = loadVolatile(params.flags); + uint32_t const cur = flags.x; + uint32_t const clearBytes = flags.y; + char* ucBase = reinterpret_cast(params.ucBuffer); + auto* curUc = reinterpret_cast<__nv_bfloat16*>(ucBase + cur * kLamportBufferBytes); + // Lamport layout of one buffer: [row][source rank][kHiddenDim]. + size_t const slotBase = static_cast(row) * kNumRanks * kHiddenDim + sliceBase + itemVec * kEltsPerVec; + + // ---- Transport: collect every rank's partial, this rank's own included, from the transport. ---- + uint4 vals[kNumRanks]; + if (hasItem) + { +#pragma unroll + for (int s = 0; s < kNumRanks; ++s) + { + vals[s] = loadVolatile(curUc + slotBase + static_cast(s) * kHiddenDim); + } + } + + // ---- Reset the previous call's buffer to the sentinel (reused two calls from now), off the poll's path. ---- + { + uint32_t const dirty = (cur + kNumLamportBuffers - 1) % kNumLamportBuffers; + auto* dirtyBuf = reinterpret_cast(ucBase + dirty * kLamportBufferBytes); + uint32_t const numVec = clearBytes / sizeof(uint4); + uint4 const sentinel = make_uint4(kLamportSentinel, kLamportSentinel, kLamportSentinel, kLamportSentinel); + uint32_t const ctaLinear = clusterId * kSlices + slice; + for (uint32_t j = ctaLinear * Cfg::kThreads + tid; j < numVec; j += kGridCtas * Cfg::kThreads) + { + dirtyBuf[j] = sentinel; + } + } + + if (hasItem) + { + bool pending = true; + while (pending) + { + pending = false; +#pragma unroll + for (int s = 0; s < kNumRanks; ++s) + { + if (isDirty(vals[s])) + { + pending = true; + vals[s] = loadVolatile(curUc + slotBase + static_cast(s) * kHiddenDim); + } + } + } + } + + // ---- Transport (earlyPoll): wait for the producer grid before reading the residual or writing any output. ---- + if (params.earlyPoll) + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + // Dependents only start their own prologue; they still wait for this grid to complete. + asm volatile("griddepcontrol.launch_dependents;"); + if (hasItem) + { + res = *reinterpret_cast(residualIn + rowOff); + } + } + + // Transport: router-weight fragments (see kRouterLoad). + constexpr int kWPairs = kRouter ? HopConfig<(kG > 0 ? kG : 4), kRowSplit>::kPairsPerGroup : 1; + [[maybe_unused]] uint4 wLo[kWPairs]; + [[maybe_unused]] uint4 wHi[kWPairs]; + + // ---- Reduce in rank order, add the residual, and form per-vector sums of squares (native arithmetic). ---- + uint4 resOut = make_uint4(0, 0, 0, 0); + if (hasItem) + { + float acc[kEltsPerVec]; +#pragma unroll + for (int e = 0; e < kEltsPerVec; ++e) + { + acc[e] = 0.F; + } +#pragma unroll + for (int s = 0; s < kNumRanks; ++s) + { + BfVec x; + x.packed = vals[s]; +#pragma unroll + for (int e = 0; e < kEltsPerVec; ++e) + { + acc[e] += __bfloat162float(x.elts[e]); + } + } + BfVec rv; + rv.packed = res; + BfVec y; + float sq = 0.F; +#pragma unroll + for (int e = 0; e < kEltsPerVec; ++e) + { + y.elts[e] = __hadd(__float2bfloat16(acc[e]), rv.elts[e]); + sq += __bfloat162float(__hmul(y.elts[e], y.elts[e])); + } + resOut = y.packed; + hop.vecSq[itemRow][itemVec] = sq; + } + if constexpr (kRouter && kRouterLoad == kRouterAfterSum) + { + if (numRowsCta > 0) + { + loadRouterFragments(wLo, wHi, params.routerWeight, expertGroup, warp, gid, tig, sliceBase); + } + } + __syncthreads(); + if (tid == 0) + { + // Every thread of this CTA has read the protocol state; CTA 0 advances it only after all CTAs report. + asm volatile("red.release.gpu.global.add.u32 [%0], %1;" ::"l"(params.flags + 2), "r"(1U) : "memory"); + } + if constexpr (kRouter && kRouterLoad == kRouterAfterBarrier) + { + if (numRowsCta > 0) + { + loadRouterFragments(wLo, wHi, params.routerWeight, expertGroup, warp, gid, tig, sliceBase); + } + } + + // ---- RMS: warp butterfly over 32 vectors, then (w0 + w2) + w1 per slice, then the 8 slices in order. ---- + if (warp < numRowsCta * kNativeWarpsPerSlice) + { + int const j = warp / kNativeWarpsPerSlice; + int const q = warp % kNativeWarpsPerSlice; + float val = hop.vecSq[j][q * 32 + lane]; +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) + { + val = val + __shfl_xor_sync(0xffffffffU, val, mask, 32); + } + if (lane == 0) + { + hop.warpSq[j][q] = val; + } + } + __syncthreads(); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + if (tid < numRowsCta * kSlices) + { + int const j = tid / kSlices; + int const dst = tid % kSlices; + float const blockSum = (hop.warpSq[j][0] + hop.warpSq[j][2]) + hop.warpSq[j][1]; + storeAsync( + clusterAddr(sharedAddr(&hop.blockSq[j][slice]), dst), blockSum, clusterAddr(sharedAddr(&hop.normBar), dst)); + } + mbarrierWaitPhase0(&hop.normBar); + + if (hasItem) + { + float fullSum = 0.F; +#pragma unroll + for (int c = 0; c < kSlices; ++c) + { + fullSum += hop.blockSq[itemRow][c]; + } + float const rcpRms = rsqrtf(fullSum / static_cast(kHiddenDim) + params.eps); + BfVec y; + y.packed = resOut; + BfVec g; + g.packed = hop.gamma[itemVec]; + BfVec n; +#pragma unroll + for (int e = 0; e < kEltsPerVec; ++e) + { + n.elts[e] = __float2bfloat16(__bfloat162float(y.elts[e]) * rcpRms * __bfloat162float(g.elts[e])); + } + if (writesOutputs) + { + *reinterpret_cast(reinterpret_cast<__nv_bfloat16*>(params.normOut) + rowOff) = n.packed; + *reinterpret_cast(reinterpret_cast<__nv_bfloat16*>(params.residualOut) + rowOff) = y.packed; + } + if constexpr (kRouter) + { + rs.normed[itemRow][itemVec] = n.packed; + } + } + + if constexpr (kRouter) + { + if (numRowsCta > 0) + { + __syncthreads(); + int const tile = warp % Cfg::kExpertTiles; + int const kGroup = warp / Cfg::kExpertTiles; + if constexpr (kRouterLoad == kRouterBeforeMma) + { + loadRouterFragments(wLo, wHi, params.routerWeight, expertGroup, warp, gid, tig, sliceBase); + } + // Independent accumulator chains keep the MMA latency off a single dependency chain. + float acc[kMmaChains][4] = {}; +#pragma unroll + for (int p = 0; p < Cfg::kPairsPerGroup; ++p) + { + // One 16-byte load covers two k16 steps; the k-permutation inside a step is the same for both operands. + uint4 const b = rs.normed[gid][kGroup * (Cfg::kKPerGroup / kEltsPerVec) + p * 4 + tig]; + int const chain = (2 * p) % kMmaChains; + mmaBf16(acc[chain], wLo[p].x, wHi[p].x, wLo[p].y, wHi[p].y, b.x, b.y); + mmaBf16(acc[chain + 1], wLo[p].z, wHi[p].z, wLo[p].w, wHi[p].w, b.z, b.w); + } + float sum4[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + sum4[i] = (acc[0][i] + acc[1][i]) + (acc[2][i] + acc[3][i]); + } + // Expert e (cluster-local) is owned by CTA e / kExpertsPerOwner. + int const e0 = tile * 16 + gid; + int const e1 = e0 + 8; + uint32_t const owner0 = e0 / Cfg::kExpertsPerOwner; + uint32_t const owner1 = e1 / Cfg::kExpertsPerOwner; + storeAsync(clusterAddr(sharedAddr(&rs.partial[slice][kGroup][e0 % Cfg::kExpertsPerOwner][tig * 2]), owner0), + sum4[0], sum4[1], clusterAddr(sharedAddr(&hop.routerBar), owner0)); + storeAsync(clusterAddr(sharedAddr(&rs.partial[slice][kGroup][e1 % Cfg::kExpertsPerOwner][tig * 2]), owner1), + sum4[2], sum4[3], clusterAddr(sharedAddr(&hop.routerBar), owner1)); + // Also keeps this CTA resident until every partial addressed to it has landed. + mbarrierWaitPhase0(&hop.routerBar); + + // This CTA owns experts [slice * kExpertsPerOwner, (slice + 1) * kExpertsPerOwner) of the cluster. Lane + // groups of 8 sum over the 8 slices. + if (tid < Cfg::kExpertsPerOwner * numRowsCta * kSlices) + { + int const c = tid % kSlices; + int const ij = tid / kSlices; + int const ei = ij / numRowsCta; + int const j = ij % numRowsCta; + float sum = 0.F; +#pragma unroll + for (int kg = 0; kg < Cfg::kKGroups; ++kg) + { + sum += rs.partial[c][kg][ei][j]; + } +#pragma unroll + for (int mask = kSlices / 2; mask > 0; mask >>= 1) + { + sum = sum + __shfl_xor_sync(0xffffffffU, sum, mask, kSlices); + } + if (c == 0) + { + int const r = rowGroup + j * kRowStride; + params.routerLogits[r * kNumExperts + expertGroup * Cfg::kExpertsPerCluster + + slice * Cfg::kExpertsPerOwner + ei] + = sum; + } + } + } + } + + // ---- CTA 0 advances the protocol state once every CTA has read it. ---- + if (tid == 0 && slice == 0 && clusterId == 0) + { + uint32_t arrived = 0; + do + { + asm volatile("ld.acquire.gpu.global.u32 %0, [%1];" : "=r"(arrived) : "l"(params.flags + 2) : "memory"); + } while (arrived < static_cast(kGridCtas)); + uint4 const next = make_uint4((cur + 1) % kNumLamportBuffers, + static_cast(R * kNumRanks * kHiddenDim * sizeof(__nv_bfloat16)), 0U, 0U); + asm volatile("st.relaxed.gpu.global.v4.u32 [%0], {%1, %2, %3, %4};" ::"l"(params.flags), "r"(next.x), + "r"(next.y), "r"(next.z), "r"(next.w) + : "memory"); + } +#else + // Clusters, distributed shared memory and PDL need SM90+; the launcher is only reached on supported devices. + __trap(); +#endif +} + +template +cudaError_t launchTyped(PostAttentionHopFromPartialsParams const& params) +{ + using Cfg = HopConfig; + cudaLaunchAttribute attrs[2]; + attrs[0].id = cudaLaunchAttributeClusterDimension; + attrs[0].val.clusterDim.x = kSlices; + attrs[0].val.clusterDim.y = 1; + attrs[0].val.clusterDim.z = 1; + attrs[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[1].val.programmaticStreamSerializationAllowed = params.enablePdl ? 1 : 0; + + cudaLaunchConfig_t config{}; + config.gridDim = dim3(kSlices, kG > 0 ? kG * kRowSplit : R, 1); + config.blockDim = dim3(Cfg::kThreads, 1, 1); + config.dynamicSmemBytes = 0; + config.stream = params.stream; + config.attrs = attrs; + config.numAttrs = 2; + return cudaLaunchKernelEx(&config, postAttnHopX1Kernel, params); +} + +template +cudaError_t launchRows(PostAttentionHopFromPartialsParams const& params) +{ + switch (params.numRows) + { + case 1: return launchTyped<1, kG, kRowSplit>(params); + case 2: return launchTyped<2, kG, kRowSplit>(params); + case 3: return launchTyped<3, kG, kRowSplit>(params); + case 4: return launchTyped<4, kG, kRowSplit>(params); + case 5: return launchTyped<5, kG, kRowSplit>(params); + case 6: return launchTyped<6, kG, kRowSplit>(params); + case 7: return launchTyped<7, kG, kRowSplit>(params); + case 8: return launchTyped<8, kG, kRowSplit>(params); + default: return cudaErrorInvalidValue; + } +} + +// Router hop: 4 expert groups x min(R, 3) row groups (at most 12 clusters of 8 CTAs, each CTA <= 3 row slices). +template +cudaError_t launchRouterAt(PostAttentionHopFromPartialsParams const& params) +{ + switch (params.numRows) + { + case 1: return launchTyped<1, 4, 1, kRouterLoad>(params); + case 2: return launchTyped<2, 4, 2, kRouterLoad>(params); + case 3: return launchTyped<3, 4, 3, kRouterLoad>(params); + case 4: return launchTyped<4, 4, 3, kRouterLoad>(params); + case 5: return launchTyped<5, 4, 3, kRouterLoad>(params); + case 6: return launchTyped<6, 4, 3, kRouterLoad>(params); + case 7: return launchTyped<7, 4, 3, kRouterLoad>(params); + case 8: return launchTyped<8, 4, 3, kRouterLoad>(params); + default: return cudaErrorInvalidValue; + } +} + +cudaError_t launchRouter(PostAttentionHopFromPartialsParams const& params) +{ + switch (params.routerLoad) + { + case kRouterAfterBarrier: return launchRouterAt(params); + case kRouterAfterSum: return launchRouterAt(params); + case kRouterBeforeMma: return launchRouterAt(params); + default: return cudaErrorInvalidValue; + } +} + +} // namespace + +cudaError_t invokePostAttentionHopFromPartials(PostAttentionHopFromPartialsParams const& params) +{ + if (params.rank < 0 || params.rank >= kNumRanks || params.flags == nullptr || params.ucBuffer == nullptr) + { + return cudaErrorInvalidValue; + } + if (params.routerWeight != nullptr) + { + if (params.routerLogits == nullptr) + { + return cudaErrorInvalidValue; + } + return launchRouter(params); + } + return launchRows<0, 1>(params); +} + +} // namespace kernels::glm_moe_dsa::attention_output_into_hop + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopFromPartials.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopFromPartials.h new file mode 100644 index 000000000000..362fa1da20fa --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopFromPartials.h @@ -0,0 +1,92 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" +#include "tensorrt_llm/kernels/glmMoeDsa/postAttentionHop.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa::attention_output_into_hop +{ + +//! Post-attention hop that takes its input from the transport: the hop of postAttentionHop.h with its input taken +//! from the transport instead of from a partial tensor. +//! +//! The producer of this call (the attention output tail's o_proj kernel, publishing into the hop) has already +//! written every rank's sanitized rank-local partial (bf16 -0.0 mapped to +0.0) into the current Lamport buffer of +//! the post-attention hop's workspace, slot ``[row][rank][kHiddenDim]``, with multicast stores. This kernel polls +//! all ``kNumRanks`` slots of its rows (its own included) and then performs exactly the post-attention hop's +//! arithmetic: the rank sum in rank order 0..7, the bf16 residual add, RMSNorm with the native reduction tree and +//! the optional router logits with the post-attention hop's MMA decomposition and summation order. Given the same +//! partials, ``normed``, ``residualOut`` and ``routerLogits`` are bit-identical to ``postAttnHop`` consuming the +//! ordinary partial (whose own contribution is the same sanitized value, taken from registers there). +//! +//! The workspace, buffer rotation and protocol state are the post-attention hop's +//! (``kernels::glm_moe_dsa::post_attention_hop``), so calls of either kind can follow each other on one workspace. +//! Memory-order contract with the producer (programmatic dependent launch): +//! * the producer reads the protocol state and writes the transport only after its own ``griddepcontrol.wait`` and +//! triggers its dependents only after that wait, so when this kernel starts, every grid before the producer has +//! completed: the protocol state written by the previous hop and ``residualIn`` are final; +//! * with ``earlyPoll`` this kernel reads the protocol state and polls the transport before its +//! ``griddepcontrol.wait`` (the transport is private to the package and Lamport-validated); it reads +//! ``residualIn`` and writes every output only after the wait, i.e. after the producer grid has completed. +//! Without ``earlyPoll`` it waits first. The early read relies on every kernel between two hops waiting on its +//! predecessor before it completes, which holds for the served graph's kernels (each consumes its predecessor's +//! output). +struct PostAttentionHopFromPartialsParams +{ + void const* residualIn = nullptr; //!< [numRows, kHiddenDim] bf16 + void const* gamma = nullptr; //!< [kHiddenDim] bf16 RMSNorm weight + void const* routerWeight = nullptr; //!< [kNumExperts, kHiddenDim] bf16 or nullptr (no router logits) + void* normOut = nullptr; //!< [numRows, kHiddenDim] bf16 + void* residualOut = nullptr; //!< [numRows, kHiddenDim] bf16 + float* routerLogits = nullptr; //!< [numRows, kNumExperts] fp32, required iff routerWeight != nullptr + void* ucBuffer = nullptr; //!< local unicast view of the post-attention hop's multicast workspace + uint32_t* flags = nullptr; //!< [4] protocol state of that workspace + int numRows = 0; + int rank = 0; + float eps = 1e-5F; + bool enablePdl = true; + //! true: read the protocol state and poll the transport before ``griddepcontrol.wait`` (relies on the producer + //! contract above); false: wait first, as postAttentionHop does, then read the state and poll. + bool earlyPoll = true; + //! Where the router-weight MMA fragments are loaded (same values and arithmetic in every case): 0 after the + //! barrier that ends the rank sum, 1 right after the rank sum (default), 2 right before the MMAs. + int routerLoad = 1; + cudaStream_t stream = nullptr; +}; + +using glm_moe_dsa::post_attention_hop::kHiddenDim; +using glm_moe_dsa::post_attention_hop::kLamportBufferBytes; +using glm_moe_dsa::post_attention_hop::kMaxRows; +using glm_moe_dsa::post_attention_hop::kNumExperts; +using glm_moe_dsa::post_attention_hop::kNumLamportBuffers; +using glm_moe_dsa::post_attention_hop::kNumRanks; +using glm_moe_dsa::post_attention_hop::kWorkspaceBytes; + +//! Launch the transport-input hop on ``params.stream``. Returns the launch status. +cudaError_t invokePostAttentionHopFromPartials(PostAttentionHopFromPartialsParams const& params); + +} // namespace kernels::glm_moe_dsa::attention_output_into_hop + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopRouting.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopRouting.cu new file mode 100644 index 000000000000..fe8cc14ea457 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopRouting.cu @@ -0,0 +1,889 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/glmMoeDsa/postAttentionHopRouting.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa::hop_routing +{ +namespace +{ +namespace cg = cooperative_groups; + +// The partials hop (glmMoeDsa/postAttentionHopFromPartials.cu, router variant) and top-8 routing. Everything up to and +// including the router-logit stores is the partials hop's code; the additions are marked "Routing". +// +// Work decomposition (the partials hop's). Every cluster has kSlices CTAs; CTA c owns hidden slice +// [c * kSlice, (c + 1) * kSlice). There are kG x kRowSplit clusters (kG = 4, kRowSplit = min(R, 3)). Cluster (g, h) +// collects and normalizes the rows r with r % kRowSplit == h and computes their logits for experts [64 g, 64 g + 64); +// CTA c of the cluster owns the 8 experts 64 g + 8 c + [0, 8) and writes their logits. +// +// Routing. The owner of logit (r, e) also writes the score bits of (r, e) into the gather buffer. Row +// r = h + j * kRowSplit (local row j of row group h) is finished by CTA 0 of cluster (g = j, h): its 256 threads poll +// the row's 256 scores, each warp ranks its 32 keys and keeps its top 8 (level 1), the 256 threads rank the 64 +// candidates (level 2, four lanes per candidate), and lanes 0..7 write the row's ids and weights. +constexpr int kSlices = 8; +constexpr int kSlice = kHiddenDim / kSlices; +constexpr int kEltsPerVec = 8; +constexpr int kVecPerSlice = kSlice / kEltsPerVec; +constexpr int kNativeWarpsPerSlice = kVecPerSlice / 32; +// Row stride of the normalized-slice tile in 16-byte vectors; 100 = 96 + 4 keeps the 16-byte B-fragment loads of two +// rows in one quarter-warp phase on disjoint banks. +constexpr int kNormedStride = kVecPerSlice + 4; +constexpr int kMmaChains = 4; +constexpr uint32_t kLamportSentinel = 0x80000000U; +// The empty gather slot: a routing score is never -0.0f (see postAttentionHopRouting.h). +constexpr uint32_t kGatherSentinel = 0x80000000U; +constexpr int kG = 4; +constexpr int kExpertWarps = kNumExperts / 32; +constexpr int kCandidates = kExpertWarps * kTopK; + +static_assert(kHiddenDim % (kSlices * kEltsPerVec * 32) == 0); +static_assert(kVecPerSlice == kNativeWarpsPerSlice * 32); +static_assert(kMaxRows == 8, "the router MMA maps rows onto the n8 dimension"); +static_assert(kNumExperts == 256 && kTopK == 8, "the finisher maps one expert per thread and 8 candidates per warp"); +static_assert(kCandidates * 4 == kNumExperts, "level 2 ranks each candidate on four lanes"); + +template +struct HopConfig +{ + static constexpr int kThreads = 384; + static constexpr int kWarps = kThreads / 32; + static constexpr int kExpertsPerCluster = kNumExperts / kG; + static constexpr int kExpertTiles = kExpertsPerCluster / 16; + static constexpr int kKGroups = kWarps / kExpertTiles; + static constexpr int kKPerGroup = kSlice / kKGroups; + static constexpr int kPairsPerGroup = kKPerGroup / 32; + static constexpr int kExpertsPerOwner = kExpertsPerCluster / kSlices; + static_assert(kExpertsPerCluster % 16 == 0 && kWarps % kExpertTiles == 0 && kSlice % (kKGroups * 32) == 0); + static_assert(kPairsPerGroup % (kMmaChains / 2) == 0); + static_assert(kThreads >= kNumExperts); +}; + +constexpr int kMaxRowsPerCtaAll = kMaxRows / 2; + +struct alignas(16) HopStorage +{ + uint4 gamma[kVecPerSlice]; + float vecSq[kMaxRowsPerCtaAll][kVecPerSlice]; + float warpSq[kMaxRowsPerCtaAll][kNativeWarpsPerSlice]; + float blockSq[kMaxRowsPerCtaAll][kSlices]; + // Completes when all slice sums of this CTA's rows have arrived (st.async complete_tx). + alignas(8) uint64_t normBar; + // Completes when all partial logits of the experts this CTA owns have arrived. + alignas(8) uint64_t routerBar; +}; + +template +struct alignas(16) RouterStorage +{ + using Cfg = HopConfig; + uint4 normed[kMaxRows][kNormedStride]; // local row on the MMA n dimension; unused rows stay zero + float partial[kSlices][Cfg::kKGroups][Cfg::kExpertsPerOwner][kMaxRows]; +}; + +// Routing: state of one CTA. +struct alignas(16) RouteStorage +{ + unsigned long long candKey[kCandidates]; // level-1 survivors: native 64-bit keys + float candSg[kCandidates]; // their score - bias + unsigned long long selKey[kTopK]; // level-2 result in rank order + float selSg[kTopK]; + float biasOwn[8]; // bias of the experts this CTA owns + uint32_t parity; // gather parity of this call + float biasAll[kNumExperts]; // finisher CTAs: bias of every expert +#if defined(TLLM_GLM_MOE_DSA_HOP_ROUTING_STAMPS) + unsigned long long stamp[kNumStamps]; +#endif +}; + +union BfVec +{ + uint4 packed; + __nv_bfloat16 elts[kEltsPerVec]; +}; + +__device__ __forceinline__ uint4 loadVolatile(void const* ptr) +{ + uint4 v; + asm volatile("ld.volatile.global.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(ptr) + : "memory"); + return v; +} + +__device__ __forceinline__ uint32_t loadVolatileU32(void const* ptr) +{ + uint32_t v; + asm volatile("ld.volatile.global.u32 %0, [%1];" : "=r"(v) : "l"(ptr) : "memory"); + return v; +} + +__device__ __forceinline__ uint4 loadConstant(void const* ptr) +{ + uint4 v; + asm volatile("ld.global.nc.v4.u32 {%0, %1, %2, %3}, [%4];" : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) : "l"(ptr)); + return v; +} + +__device__ __forceinline__ float loadBiasConstant(void const* bias, int expert) +{ + unsigned short raw; + asm volatile("ld.global.nc.u16 %0, [%1];" : "=h"(raw) : "l"(reinterpret_cast<__nv_bfloat16 const*>(bias) + expert)); + return __bfloat162float(__ushort_as_bfloat16(raw)); +} + +__device__ __forceinline__ void prefetchL2(void const* ptr) +{ + asm volatile("prefetch.global.L2 [%0];" ::"l"(ptr)); +} + +__device__ __forceinline__ void storeRelaxed(uint32_t* ptr, uint32_t v) +{ + asm volatile("st.relaxed.gpu.global.u32 [%0], %1;" ::"l"(ptr), "r"(v) : "memory"); +} + +__device__ __forceinline__ uint32_t loadRelaxed(uint32_t const* ptr) +{ + uint32_t v; + asm volatile("ld.relaxed.gpu.global.u32 %0, [%1];" : "=r"(v) : "l"(ptr) : "memory"); + return v; +} + +__device__ __forceinline__ bool isDirty(uint4 const& v) +{ + return v.x == kLamportSentinel || v.y == kLamportSentinel || v.z == kLamportSentinel || v.w == kLamportSentinel; +} + +__device__ __forceinline__ uint32_t sharedAddr(void const* ptr) +{ + return static_cast(__cvta_generic_to_shared(ptr)); +} + +// Address of the same shared-memory object in CTA `rank` of the cluster. +__device__ __forceinline__ uint32_t clusterAddr(uint32_t addr, uint32_t rank) +{ + uint32_t out; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(out) : "r"(addr), "r"(rank)); + return out; +} + +__device__ __forceinline__ void mbarrierInitExpectTx(uint64_t* bar, uint32_t bytes) +{ + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(sharedAddr(bar)) : "memory"); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(sharedAddr(bar)), "r"(bytes) + : "memory"); +} + +__device__ __forceinline__ void mbarrierWaitPhase0(uint64_t* bar) +{ + asm volatile( + "{\n" + ".reg .pred done;\n" + "ROUTING_MBAR_WAIT_%=:\n" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 done, [%0], 0;\n" + "@!done bra ROUTING_MBAR_WAIT_%=;\n" + "}\n" ::"r"(sharedAddr(bar)) + : "memory"); +} + +// Remote shared-memory stores that complete transactions on the destination CTA's mbarrier. +__device__ __forceinline__ void storeAsync(uint32_t remoteAddr, float v, uint32_t remoteBar) +{ + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.f32 [%0], %1, [%2];" ::"r"(remoteAddr), "f"(v), + "r"(remoteBar) + : "memory"); +} + +__device__ __forceinline__ void storeAsync(uint32_t remoteAddr, float a, float b, uint32_t remoteBar) +{ + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.v2.f32 [%0], {%1, %2}, [%3];" ::"r"(remoteAddr), + "f"(a), "f"(b), "r"(remoteBar) + : "memory"); +} + +__device__ __forceinline__ void mmaBf16( + float (&d)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t b0, uint32_t b1) +{ + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, " + "{%0, %1, %2, %3};" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)); +} + +// One warp's router-weight MMA fragments (the post-attention hop's register layout): experts tile * 16 + gid (+ 8) +// of this cluster's expert group, k pairs of this warp's k group within the CTA's slice. +template +__device__ __forceinline__ void loadRouterFragments(uint4 (&wLo)[Cfg::kPairsPerGroup], + uint4 (&wHi)[Cfg::kPairsPerGroup], void const* routerWeight, int expertGroup, int warp, int gid, int tig, + int sliceBase) +{ + int const tile = warp % Cfg::kExpertTiles; + int const kGroup = warp / Cfg::kExpertTiles; + auto const* w = reinterpret_cast<__nv_bfloat16 const*>(routerWeight); + int const expert = expertGroup * Cfg::kExpertsPerCluster + tile * 16 + gid; +#pragma unroll + for (int p = 0; p < Cfg::kPairsPerGroup; ++p) + { + int const k = sliceBase + kGroup * Cfg::kKPerGroup + p * 32 + tig * kEltsPerVec; + wLo[p] = loadConstant(w + static_cast(expert) * kHiddenDim + k); + wHi[p] = loadConstant(w + static_cast(expert + 8) * kHiddenDim + k); + } +} + +// Router-weight fragment load points, as in the partials hop (all orders give the same values and arithmetic). +constexpr int kRouterAfterBarrier = 0; +constexpr int kRouterAfterSum = 1; +constexpr int kRouterBeforeMma = 2; + +// Routing: the native routing score (RoutingKernel.cuh sigmoid_accurate plus the SigmoidBias preprocess). +__device__ __forceinline__ float routingScore(float logit, float bias) +{ + float const s = 0.5f * tanhf(0.5f * logit) + 0.5f; + return s + bias; +} + +// Routing: cub::Traits::TwiddleIn, the order-preserving map of FP32 bits onto unsigned integers. +__device__ __forceinline__ uint32_t twiddleIn(uint32_t bits) +{ + return bits ^ ((bits & 0x80000000U) != 0U ? 0xFFFFFFFFU : 0x80000000U); +} + +#if defined(TLLM_GLM_MOE_DSA_HOP_ROUTING_STAMPS) +__device__ __forceinline__ unsigned long long globalTimer() +{ + unsigned long long t; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t)); + return t; +} + +#define ROUTING_STAMP(field) atomicMax(&route.stamp[field], globalTimer()) +#else +#define ROUTING_STAMP(field) \ + do \ + { \ + } while (0) +#endif + +// Stamp fields (diagnostic builds). +constexpr int kStampEntry = 0; +constexpr int kStampWait = 1; +constexpr int kStampLogits = 2; +constexpr int kStampGather = 3; +constexpr int kStampPolled = 4; +constexpr int kStampLevel1 = 5; +constexpr int kStampRouted = 6; +constexpr int kStampEnd = 7; +static_assert(kStampEnd < kNumStamps); + +// The register cap of the partials hop, under which one CTA of it fits on an SM next to a CTA of the publishing +// o_proj. The routing code after the MMA phase would raise the allocation above it; the cap keeps the partials hop's +// allocation without spills. +constexpr int kMaxRegisters = 112; + +template +__global__ void __maxnreg__(kMaxRegisters) postAttnHopZ2Kernel(PostAttentionHopRoutingParams const params) +{ + static_assert(R >= 1 && R <= kMaxRows); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + using Cfg = HopConfig; + constexpr int kClusters = kG * kRowSplit; + constexpr int kGridCtas = kClusters * kSlices; + static_assert(kGridCtas <= kMaxGridCtas); + // Rows one CTA collects, normalizes and multiplies. + constexpr int kRowsPerCta = (R + kRowSplit - 1) / kRowSplit; + static_assert(kRowsPerCta <= kMaxRowsPerCtaAll && kRowsPerCta <= kG); + static_assert(kRowsPerCta * kVecPerSlice <= Cfg::kThreads && kRowsPerCta * kNativeWarpsPerSlice <= Cfg::kWarps); + static_assert(Cfg::kExpertsPerOwner * kRowsPerCta * kSlices <= Cfg::kThreads); + + __shared__ HopStorage hop; + __shared__ RouterStorage rs; + __shared__ RouteStorage route; + + cg::cluster_group cluster = cg::this_cluster(); + int const slice = static_cast(cluster.block_rank()); + int const clusterId = blockIdx.y; + int const expertGroup = clusterId % kG; + int const rowGroup = clusterId / kG; + int const tid = threadIdx.x; + int const lane = tid & 31; + int const warp = tid >> 5; + int const sliceBase = slice * kSlice; + uint32_t const ctaLinear = clusterId * kSlices + slice; + + // Rows this CTA collects and normalizes: local row j is global row rowGroup + j * kRowSplit. + constexpr int kRowStride = kRowSplit; + int const numRowsCta = (R - rowGroup + kRowSplit - 1) / kRowSplit; + // This thread's item: one 16-byte vector of one row slice. + int const itemRow = tid / kVecPerSlice; + int const itemVec = tid % kVecPerSlice; + bool const hasItem = itemRow < numRowsCta; + int const row = rowGroup + itemRow * kRowStride; + size_t const rowOff = static_cast(row) * kHiddenDim + sliceBase + itemVec * kEltsPerVec; + // The cluster that writes the normalized rows and the residual. + bool const writesOutputs = expertGroup == 0; + // Routing: this CTA finishes local row `expertGroup` of its row group, if that row exists. + bool const finisher = slice == 0 && expertGroup < numRowsCta; + +#if defined(TLLM_GLM_MOE_DSA_HOP_ROUTING_STAMPS) + if (tid < kNumStamps) + { + route.stamp[tid] = 0ULL; + } + __syncthreads(); + if (tid == 0) + { + ROUTING_STAMP(kStampEntry); + } +#endif + + // ---- Prologue: constant model weights only; may overlap the producer of `input` under PDL. ---- + if (tid < kVecPerSlice) + { + hop.gamma[tid] = loadConstant(reinterpret_cast(params.gamma) + slice * kVecPerSlice + tid); + } + // Routing: router bias of the owned experts and, on finishers, of every expert (a model parameter). + if (tid < Cfg::kExpertsPerOwner) + { + route.biasOwn[tid] = loadBiasConstant( + params.routingBias, expertGroup * Cfg::kExpertsPerCluster + slice * Cfg::kExpertsPerOwner + tid); + } + if (finisher && tid < kNumExperts) + { + route.biasAll[tid] = loadBiasConstant(params.routingBias, tid); + } + int const gid = lane >> 2; + int const tig = lane & 3; + // The router-weight fragments are loaded after the poll (not here), so they are not live across it; the prologue + // only asks L2 for the same lines (no registers held, no architectural effect). + { + int const tile = warp % Cfg::kExpertTiles; + int const kGroup = warp / Cfg::kExpertTiles; + auto const* w = reinterpret_cast<__nv_bfloat16 const*>(params.routerWeight); + int const expert = expertGroup * Cfg::kExpertsPerCluster + tile * 16 + gid; +#pragma unroll + for (int p = 0; p < Cfg::kPairsPerGroup; ++p) + { + int const k = sliceBase + kGroup * Cfg::kKPerGroup + p * 32 + tig * kEltsPerVec; + prefetchL2(w + static_cast(expert) * kHiddenDim + k); + prefetchL2(w + static_cast(expert + 8) * kHiddenDim + k); + } + } + // MMA columns of rows this CTA does not own stay zero. + for (int i = tid; i < kMaxRows * kNormedStride; i += Cfg::kThreads) + { + rs.normed[i / kNormedStride][i % kNormedStride] = make_uint4(0, 0, 0, 0); + } + + // Cluster-local completion barriers: every CTA of a cluster has the same rows, so each expects the 8 slice sums + // of its rows and the 8 x kKGroups partial tiles of the experts it owns. A split cluster barrier publishes the + // initialization; its wait precedes the first remote store, long after all CTAs arrived. + if (tid == 0) + { + mbarrierInitExpectTx(&hop.normBar, static_cast(numRowsCta * kSlices * sizeof(float))); + mbarrierInitExpectTx(&hop.routerBar, numRowsCta > 0 ? static_cast(sizeof(rs.partial)) : 0U); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + + // With earlyPoll no wait here (the protocol state and the transport may be read early, see + // postAttentionHopFromPartials.h). + auto const* residualIn = reinterpret_cast<__nv_bfloat16 const*>(params.residualIn); + uint4 res = make_uint4(0, 0, 0, 0); + if (!params.earlyPoll) + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + // Dependents may start their prologue; they still wait for this grid to complete. + asm volatile("griddepcontrol.launch_dependents;"); + if (hasItem) + { + res = *reinterpret_cast(residualIn + rowOff); + } + } + uint4 const flags = loadVolatile(params.flags); + uint32_t const cur = flags.x; + uint32_t const clearBytes = flags.y; + char* ucBase = reinterpret_cast(params.ucBuffer); + auto* curUc = reinterpret_cast<__nv_bfloat16*>(ucBase + cur * kLamportBufferBytes); + // Lamport layout of one buffer: [row][source rank][kHiddenDim]. + size_t const slotBase = static_cast(row) * kNumRanks * kHiddenDim + sliceBase + itemVec * kEltsPerVec; + + // ---- Collect every rank's partial, this rank's own included, from the transport. ---- + uint4 vals[kNumRanks]; + if (hasItem) + { +#pragma unroll + for (int s = 0; s < kNumRanks; ++s) + { + vals[s] = loadVolatile(curUc + slotBase + static_cast(s) * kHiddenDim); + } + } + + // ---- Reset the previous call's buffer to the sentinel (reused two calls from now), off the poll's path. ---- + { + uint32_t const dirty = (cur + kNumLamportBuffers - 1) % kNumLamportBuffers; + auto* dirtyBuf = reinterpret_cast(ucBase + dirty * kLamportBufferBytes); + uint32_t const numVec = clearBytes / sizeof(uint4); + uint4 const sentinel = make_uint4(kLamportSentinel, kLamportSentinel, kLamportSentinel, kLamportSentinel); + for (uint32_t j = ctaLinear * Cfg::kThreads + tid; j < numVec; j += kGridCtas * Cfg::kThreads) + { + dirtyBuf[j] = sentinel; + } + } + + if (hasItem) + { + bool pending = true; + while (pending) + { + pending = false; +#pragma unroll + for (int s = 0; s < kNumRanks; ++s) + { + if (isDirty(vals[s])) + { + pending = true; + vals[s] = loadVolatile(curUc + slotBase + static_cast(s) * kHiddenDim); + } + } + } + } + + // ---- earlyPoll: wait for the producer grid before reading the residual or writing any output. ---- + if (params.earlyPoll) + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + // Dependents only start their own prologue; they still wait for this grid to complete. + asm volatile("griddepcontrol.launch_dependents;"); + if (hasItem) + { + res = *reinterpret_cast(residualIn + rowOff); + } + } +#if defined(TLLM_GLM_MOE_DSA_HOP_ROUTING_STAMPS) + if (tid == 0) + { + ROUTING_STAMP(kStampWait); + } +#endif + + // Router-weight fragments (see kRouterLoad). + uint4 wLo[Cfg::kPairsPerGroup]; + uint4 wHi[Cfg::kPairsPerGroup]; + + // ---- Reduce in rank order, add the residual, and form per-vector sums of squares (native arithmetic). ---- + uint4 resOut = make_uint4(0, 0, 0, 0); + if (hasItem) + { + float acc[kEltsPerVec]; +#pragma unroll + for (int e = 0; e < kEltsPerVec; ++e) + { + acc[e] = 0.F; + } +#pragma unroll + for (int s = 0; s < kNumRanks; ++s) + { + BfVec x; + x.packed = vals[s]; +#pragma unroll + for (int e = 0; e < kEltsPerVec; ++e) + { + acc[e] += __bfloat162float(x.elts[e]); + } + } + BfVec rv; + rv.packed = res; + BfVec y; + float sq = 0.F; +#pragma unroll + for (int e = 0; e < kEltsPerVec; ++e) + { + y.elts[e] = __hadd(__float2bfloat16(acc[e]), rv.elts[e]); + sq += __bfloat162float(__hmul(y.elts[e], y.elts[e])); + } + resOut = y.packed; + hop.vecSq[itemRow][itemVec] = sq; + } + if constexpr (kRouterLoad == kRouterAfterSum) + { + if (numRowsCta > 0) + { + loadRouterFragments(wLo, wHi, params.routerWeight, expertGroup, warp, gid, tig, sliceBase); + } + } + if (tid == 0) + { + // Routing: the gather parity of this call (the routing state is read only after the dependency wait). It is + // kept in shared memory, published by the barrier below, so no register holds it across the MMA phase. + route.parity = loadVolatileU32(params.zstate) & 1U; + } + __syncthreads(); + if (tid == 0) + { + // Every thread of this CTA has read the protocol state and thread 0 the gather parity; CTA 0 advances both only + // after all CTAs report. + asm volatile("red.release.gpu.global.add.u32 [%0], %1;" ::"l"(params.flags + 2), "r"(1U) : "memory"); + } + // ---- Routing: reset the previous call's gather buffer (its readers belong to a completed grid), off the routing + // path. ---- + { + uint4* const other = reinterpret_cast(params.gather + (route.parity ^ 1U) * kGatherBufferWords); + uint4 const sentinel = make_uint4(kGatherSentinel, kGatherSentinel, kGatherSentinel, kGatherSentinel); + for (uint32_t j = ctaLinear * Cfg::kThreads + tid; j < kGatherBufferWords / 4; j += kGridCtas * Cfg::kThreads) + { + other[j] = sentinel; + } + } + if constexpr (kRouterLoad == kRouterAfterBarrier) + { + if (numRowsCta > 0) + { + loadRouterFragments(wLo, wHi, params.routerWeight, expertGroup, warp, gid, tig, sliceBase); + } + } + + // ---- RMS: warp butterfly over 32 vectors, then (w0 + w2) + w1 per slice, then the 8 slices in order. ---- + if (warp < numRowsCta * kNativeWarpsPerSlice) + { + int const j = warp / kNativeWarpsPerSlice; + int const q = warp % kNativeWarpsPerSlice; + float val = hop.vecSq[j][q * 32 + lane]; +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) + { + val = val + __shfl_xor_sync(0xffffffffU, val, mask, 32); + } + if (lane == 0) + { + hop.warpSq[j][q] = val; + } + } + __syncthreads(); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + if (tid < numRowsCta * kSlices) + { + int const j = tid / kSlices; + int const dst = tid % kSlices; + float const blockSum = (hop.warpSq[j][0] + hop.warpSq[j][2]) + hop.warpSq[j][1]; + storeAsync( + clusterAddr(sharedAddr(&hop.blockSq[j][slice]), dst), blockSum, clusterAddr(sharedAddr(&hop.normBar), dst)); + } + mbarrierWaitPhase0(&hop.normBar); + + if (hasItem) + { + float fullSum = 0.F; +#pragma unroll + for (int c = 0; c < kSlices; ++c) + { + fullSum += hop.blockSq[itemRow][c]; + } + float const rcpRms = rsqrtf(fullSum / static_cast(kHiddenDim) + params.eps); + BfVec y; + y.packed = resOut; + BfVec g; + g.packed = hop.gamma[itemVec]; + BfVec n; +#pragma unroll + for (int e = 0; e < kEltsPerVec; ++e) + { + n.elts[e] = __float2bfloat16(__bfloat162float(y.elts[e]) * rcpRms * __bfloat162float(g.elts[e])); + } + if (writesOutputs) + { + *reinterpret_cast(reinterpret_cast<__nv_bfloat16*>(params.normOut) + rowOff) = n.packed; + *reinterpret_cast(reinterpret_cast<__nv_bfloat16*>(params.residualOut) + rowOff) = y.packed; + } + rs.normed[itemRow][itemVec] = n.packed; + } + + if (numRowsCta > 0) + { + __syncthreads(); + int const tile = warp % Cfg::kExpertTiles; + int const kGroup = warp / Cfg::kExpertTiles; + if constexpr (kRouterLoad == kRouterBeforeMma) + { + loadRouterFragments(wLo, wHi, params.routerWeight, expertGroup, warp, gid, tig, sliceBase); + } + // Independent accumulator chains keep the MMA latency off a single dependency chain. + float acc[kMmaChains][4] = {}; +#pragma unroll + for (int p = 0; p < Cfg::kPairsPerGroup; ++p) + { + // One 16-byte load covers two k16 steps; the k-permutation inside a step is the same for both operands. + uint4 const b = rs.normed[gid][kGroup * (Cfg::kKPerGroup / kEltsPerVec) + p * 4 + tig]; + int const chain = (2 * p) % kMmaChains; + mmaBf16(acc[chain], wLo[p].x, wHi[p].x, wLo[p].y, wHi[p].y, b.x, b.y); + mmaBf16(acc[chain + 1], wLo[p].z, wHi[p].z, wLo[p].w, wHi[p].w, b.z, b.w); + } + float sum4[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + sum4[i] = (acc[0][i] + acc[1][i]) + (acc[2][i] + acc[3][i]); + } + // Expert e (cluster-local) is owned by CTA e / kExpertsPerOwner. + int const e0 = tile * 16 + gid; + int const e1 = e0 + 8; + uint32_t const owner0 = e0 / Cfg::kExpertsPerOwner; + uint32_t const owner1 = e1 / Cfg::kExpertsPerOwner; + storeAsync(clusterAddr(sharedAddr(&rs.partial[slice][kGroup][e0 % Cfg::kExpertsPerOwner][tig * 2]), owner0), + sum4[0], sum4[1], clusterAddr(sharedAddr(&hop.routerBar), owner0)); + storeAsync(clusterAddr(sharedAddr(&rs.partial[slice][kGroup][e1 % Cfg::kExpertsPerOwner][tig * 2]), owner1), + sum4[2], sum4[3], clusterAddr(sharedAddr(&hop.routerBar), owner1)); + // Also keeps this CTA resident until every partial addressed to it has landed. + mbarrierWaitPhase0(&hop.routerBar); + + // This CTA owns experts [slice * kExpertsPerOwner, (slice + 1) * kExpertsPerOwner) of the cluster. Lane + // groups of 8 sum over the 8 slices. + if (tid < Cfg::kExpertsPerOwner * numRowsCta * kSlices) + { + int const c = tid % kSlices; + int const ij = tid / kSlices; + int const ei = ij / numRowsCta; + int const j = ij % numRowsCta; + float sum = 0.F; +#pragma unroll + for (int kg = 0; kg < Cfg::kKGroups; ++kg) + { + sum += rs.partial[c][kg][ei][j]; + } +#pragma unroll + for (int mask = kSlices / 2; mask > 0; mask >>= 1) + { + sum = sum + __shfl_xor_sync(0xffffffffU, sum, mask, kSlices); + } + if (c == 0) + { + int const r = rowGroup + j * kRowStride; + int const e = expertGroup * Cfg::kExpertsPerCluster + slice * Cfg::kExpertsPerOwner + ei; + params.routerLogits[r * kNumExperts + e] = sum; + ROUTING_STAMP(kStampLogits); + // Routing: the score of (r, e) into this call's gather buffer. + float const logit = kLogitsIn ? params.logitsIn[r * kNumExperts + e] : sum; + storeRelaxed(params.gather + route.parity * kGatherBufferWords + r * kNumExperts + e, + __float_as_uint(routingScore(logit, route.biasOwn[ei]))); + ROUTING_STAMP(kStampGather); + } + } + } + + // ---- Routing: finisher CTAs select the top 8 of their row. ---- + if (finisher) + { + int const rf = rowGroup + expertGroup * kRowStride; + if (tid < kNumExperts) + { + uint32_t const* slot = params.gather + route.parity * kGatherBufferWords + rf * kNumExperts + tid; + uint32_t bits = loadRelaxed(slot); + while (bits == kGatherSentinel) + { + bits = loadRelaxed(slot); + } +#if defined(TLLM_GLM_MOE_DSA_HOP_ROUTING_STAMPS) + __syncwarp(); + if (lane == 0) + { + ROUTING_STAMP(kStampPolled); + } +#endif + float const sg = __uint_as_float(bits) - route.biasAll[tid]; + uint32_t const hi = twiddleIn(bits); + // Level 1: rank of this key among the warp's 32 (expert 32 * warp + i sits in lane i, so equal scores + // rank the lower lane first, as the native key does). + int rank = 0; +#pragma unroll + for (int i = 0; i < 32; ++i) + { + uint32_t const other = __shfl_sync(0xffffffffU, hi, i); + rank += (other > hi || (other == hi && i < lane)) ? 1 : 0; + } + if (rank < kTopK) + { + route.candKey[warp * kTopK + rank] + = (static_cast(hi) << 32) | static_cast(65535 - tid); + route.candSg[warp * kTopK + rank] = sg; + } + } + __syncthreads(); +#if defined(TLLM_GLM_MOE_DSA_HOP_ROUTING_STAMPS) + if (tid == 0) + { + ROUTING_STAMP(kStampLevel1); + } +#endif + if (tid < kNumExperts) + { + // Level 2: rank of the 64 candidates by the full native key (all keys are distinct). Four lanes share + // candidate tid / 4; lane q of them counts the larger keys among candidates q, q + 4, ..., q + 60. + int const ci = tid >> 2; + int const q = tid & 3; + unsigned long long const key = route.candKey[ci]; + int rank = 0; +#pragma unroll + for (int i = 0; i < kCandidates / 4; ++i) + { + rank += route.candKey[i * 4 + q] > key ? 1 : 0; + } + rank += __shfl_xor_sync(0xffffffffU, rank, 1); + rank += __shfl_xor_sync(0xffffffffU, rank, 2); + if (q == 0 && rank < kTopK) + { + route.selKey[rank] = key; + route.selSg[rank] = route.candSg[ci]; + } + } + __syncthreads(); + if (warp == 0) + { + if (lane < kTopK) + { + float s[kTopK]; +#pragma unroll + for (int k = 0; k < kTopK; ++k) + { + s[k] = route.selSg[k]; + } + // The native warp sum over lanes 0..7 (zeros above): butterfly offsets 16, 8, 4, 2, 1. + float const total = ((s[0] + s[4]) + (s[2] + s[6])) + ((s[1] + s[5]) + (s[3] + s[7])); + float const weight = route.selSg[lane] * params.routedScaling / total; + int const expert = 65535 - static_cast(route.selKey[lane] & 0xFFFFULL); + params.topkIds[rf * kTopK + lane] = expert; + params.topkWeights[rf * kTopK + lane] = __bfloat162float(__float2bfloat16(weight)); + ROUTING_STAMP(kStampRouted); + } + } + } + + // ---- CTA 0 advances the protocol state and the gather parity once every CTA has read them. ---- + if (tid == 0 && slice == 0 && clusterId == 0) + { + uint32_t arrived = 0; + do + { + asm volatile("ld.acquire.gpu.global.u32 %0, [%1];" : "=r"(arrived) : "l"(params.flags + 2) : "memory"); + } while (arrived < static_cast(kGridCtas)); + uint4 const next = make_uint4((cur + 1) % kNumLamportBuffers, + static_cast(R * kNumRanks * kHiddenDim * sizeof(__nv_bfloat16)), 0U, 0U); + asm volatile("st.relaxed.gpu.global.v4.u32 [%0], {%1, %2, %3, %4};" ::"l"(params.flags), "r"(next.x), + "r"(next.y), "r"(next.z), "r"(next.w) + : "memory"); + storeRelaxed(params.zstate, route.parity ^ 1U); + } +#if defined(TLLM_GLM_MOE_DSA_HOP_ROUTING_STAMPS) + __syncthreads(); + if (tid == 0) + { + ROUTING_STAMP(kStampEnd); + } + __syncthreads(); + if (tid < kNumStamps && params.stamps != nullptr) + { + params.stamps[ctaLinear * kNumStamps + tid] = route.stamp[tid]; + } +#endif +#else + // Clusters, distributed shared memory and PDL need SM90+; the launcher is only reached on supported devices. + __trap(); +#endif +} + +template +cudaError_t launchTyped(PostAttentionHopRoutingParams const& params) +{ + using Cfg = HopConfig; + cudaLaunchAttribute attrs[2]; + attrs[0].id = cudaLaunchAttributeClusterDimension; + attrs[0].val.clusterDim.x = kSlices; + attrs[0].val.clusterDim.y = 1; + attrs[0].val.clusterDim.z = 1; + attrs[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[1].val.programmaticStreamSerializationAllowed = params.enablePdl ? 1 : 0; + + cudaLaunchConfig_t config{}; + config.gridDim = dim3(kSlices, kG * kRowSplit, 1); + config.blockDim = dim3(Cfg::kThreads, 1, 1); + config.dynamicSmemBytes = 0; + config.stream = params.stream; + config.attrs = attrs; + config.numAttrs = 2; + return cudaLaunchKernelEx(&config, postAttnHopZ2Kernel, params); +} + +// 4 expert groups x min(R, 3) row groups (at most 12 clusters of 8 CTAs, each CTA <= 3 row slices). +template +cudaError_t launchAt(PostAttentionHopRoutingParams const& params) +{ + switch (params.numRows) + { + case 1: return launchTyped<1, 1, kRouterLoad, kLogitsIn>(params); + case 2: return launchTyped<2, 2, kRouterLoad, kLogitsIn>(params); + case 3: return launchTyped<3, 3, kRouterLoad, kLogitsIn>(params); + case 4: return launchTyped<4, 3, kRouterLoad, kLogitsIn>(params); + case 5: return launchTyped<5, 3, kRouterLoad, kLogitsIn>(params); + case 6: return launchTyped<6, 3, kRouterLoad, kLogitsIn>(params); + case 7: return launchTyped<7, 3, kRouterLoad, kLogitsIn>(params); + case 8: return launchTyped<8, 3, kRouterLoad, kLogitsIn>(params); + default: return cudaErrorInvalidValue; + } +} + +} // namespace + +bool stampsCompiled() +{ +#if defined(TLLM_GLM_MOE_DSA_HOP_ROUTING_STAMPS) + return true; +#else + return false; +#endif +} + +cudaError_t invokePostAttentionHopRouting(PostAttentionHopRoutingParams const& params) +{ + if (params.rank < 0 || params.rank >= kNumRanks || params.flags == nullptr || params.ucBuffer == nullptr + || params.routerWeight == nullptr || params.routingBias == nullptr || params.routerLogits == nullptr + || params.topkIds == nullptr || params.topkWeights == nullptr || params.gather == nullptr + || params.zstate == nullptr) + { + return cudaErrorInvalidValue; + } + if (params.logitsIn != nullptr) + { + // Test instance: the default router-weight load point only. + return params.routerLoad == kRouterAfterSum ? launchAt(params) : cudaErrorInvalidValue; + } + switch (params.routerLoad) + { + case kRouterAfterBarrier: return launchAt(params); + case kRouterAfterSum: return launchAt(params); + case kRouterBeforeMma: return launchAt(params); + default: return cudaErrorInvalidValue; + } +} + +} // namespace kernels::glm_moe_dsa::hop_routing + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopRouting.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopRouting.h new file mode 100644 index 000000000000..cb19a3f7f947 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/postAttentionHopRouting.h @@ -0,0 +1,116 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" +#include "tensorrt_llm/kernels/glmMoeDsa/postAttentionHop.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa::hop_routing +{ + +//! Pre-MoE hop with routing: the transport-input hop (postAttentionHopFromPartials.h) with router logits, plus +//! the native DeepSeek-V3 top-8 routing of every row computed from those logits. +//! +//! Normalized hidden, residual and router logits are produced exactly as by +//! ``glm_moe_dsa::attention_output_into_hop::invokePostAttentionHopFromPartials`` (same work decomposition, arithmetic +//! and protocol on the same post-attention hop workspace), so they are bit-identical to that hop's. +//! +//! Routing (the native ``SigmoidBias`` + ``ScaledSumNormalize`` routing of the TRT-LLM-Gen routing runner with one +//! expert group): every logit owner computes ``score = 0.5f * tanhf(0.5f * logit) + 0.5f + float(bias[e])`` and stores +//! its bits into the gather buffer; one finisher CTA per row collects the row's 256 scores and selects the 8 largest +//! native keys ``(TwiddleIn(score) << 32) | (65535 - e)`` (ties go to the lower expert) in descending order, and writes +//! ``topkIds[r][k]`` and ``topkWeights[r][k] = bf16((score_k - bias_k) * routedScaling / sum)`` (stored as FP32), where +//! ``sum`` adds the eight ``score_k - bias_k`` in the native warp-butterfly order ((s0+s4)+(s2+s6))+((s1+s5)+(s3+s7)). +//! +//! Gather protocol (private to this kernel): ``gather`` holds two buffers of ``kMaxRows x kNumExperts`` 32-bit words +//! and ``zstate[0]`` the parity ``p`` of the buffer this call uses. A score is never -0.0f, so the word 0x80000000 +//! marks an empty slot. Each call resets buffer ``p ^ 1`` (the previous call's) to that sentinel and CTA 0 flips the +//! parity together with the hop's state advance, after every CTA has read both states. Calls of other hop kinds on the +//! same hop workspace do not touch the routing state. +//! +//! Memory-order contract with the producer (programmatic dependent launch): as +//! ``glm_moe_dsa::attention_output_into_hop``. Before ``griddepcontrol.wait`` the kernel reads only the RMSNorm +//! weight, the router weight (L2 prefetch) and the router bias, which no kernel writes while serving; with +//! ``earlyPoll`` it also reads the protocol state of the post-attention hop workspace and polls the transport early +//! (see postAttentionHopFromPartials.h). The routing state, the gather buffers and every output are read or written +//! only after the wait. Finishers spin on gather words that every CTA writes before it waits on anything outside its +//! own cluster, so the wait-for graph is acyclic. +struct PostAttentionHopRoutingParams +{ + void const* residualIn = nullptr; //!< [numRows, kHiddenDim] bf16 + void const* gamma = nullptr; //!< [kHiddenDim] bf16 RMSNorm weight + void const* routerWeight = nullptr; //!< [kNumExperts, kHiddenDim] bf16 + void const* routingBias = nullptr; //!< [kNumExperts] bf16 score-correction bias + //! Test only: when set, the routing scores are computed from these [numRows, kNumExperts] FP32 logits instead of + //! the computed router logits (which are still written to ``routerLogits``). Read after the dependency wait. + float const* logitsIn = nullptr; + void* normOut = nullptr; //!< [numRows, kHiddenDim] bf16 + void* residualOut = nullptr; //!< [numRows, kHiddenDim] bf16 + float* routerLogits = nullptr; //!< [numRows, kNumExperts] fp32 + int32_t* topkIds = nullptr; //!< [numRows, kTopK] int32, descending native key order + float* topkWeights = nullptr; //!< [numRows, kTopK] fp32 holding bf16-rounded weights + void* ucBuffer = nullptr; //!< local unicast view of the post-attention hop's multicast workspace + uint32_t* flags = nullptr; //!< [4] protocol state of that workspace + uint32_t* gather = nullptr; //!< [kGatherWords] routing gather buffers + uint32_t* zstate = nullptr; //!< [4] routing state; word 0 is the parity + //! Diagnostic builds only (``TLLM_GLM_MOE_DSA_HOP_ROUTING_STAMPS``): [gridCtas, kNumStamps] globaltimer values. + unsigned long long* stamps = nullptr; + int numRows = 0; + int rank = 0; + float eps = 1e-5F; + float routedScaling = 2.5F; + bool enablePdl = true; + bool earlyPoll = false; + int routerLoad = 1; + cudaStream_t stream = nullptr; +}; + +using glm_moe_dsa::post_attention_hop::kHiddenDim; +using glm_moe_dsa::post_attention_hop::kLamportBufferBytes; +using glm_moe_dsa::post_attention_hop::kMaxRows; +using glm_moe_dsa::post_attention_hop::kNumExperts; +using glm_moe_dsa::post_attention_hop::kNumLamportBuffers; +using glm_moe_dsa::post_attention_hop::kNumRanks; +using glm_moe_dsa::post_attention_hop::kWorkspaceBytes; + +constexpr int kTopK = 8; +//! Words of one gather buffer and of both. +constexpr int kGatherBufferWords = kMaxRows * kNumExperts; +constexpr int kGatherWords = 2 * kGatherBufferWords; +constexpr int kStateWords = 4; +//! Diagnostic stamp fields per CTA: entry, wait, last logit store, last gather store, poll complete, level-1 done, +//! routing stored, end. +constexpr int kNumStamps = 8; +//! Upper bound of the grid size (4 expert groups x 3 row groups x 8 CTAs). +constexpr int kMaxGridCtas = 96; + +//! Whether this build records diagnostic stamps. +bool stampsCompiled(); + +//! Launch the routing hop on ``params.stream``. Returns the launch status. +cudaError_t invokePostAttentionHopRouting(PostAttentionHopRoutingParams const& params); + +} // namespace kernels::glm_moe_dsa::hop_routing + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/ropeKvAppend.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/ropeKvAppend.cu new file mode 100644 index 000000000000..e13d352ba57a --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/ropeKvAppend.cu @@ -0,0 +1,514 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/glmMoeDsa/ropeKvAppend.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/kernels/decoderMaskedMultiheadAttentionUtils.h" + +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa::rope_kv_append +{ + +namespace +{ + +constexpr int32_t kRopeThreads = 256; +constexpr int32_t kBridgeThreads = 128; +constexpr int32_t kMaxGroupSize = 8; +constexpr int32_t kEltsPerVec = 8; // bf16 elements per 16-byte load +constexpr int32_t kNopeVecs = kKvLoraRank / kEltsPerVec; // 64 +constexpr int32_t kRopeVecs = kRopeDim / kEltsPerVec; // 8 +constexpr int32_t kTotalVecsPerHead = kNopeVecs + kRopeVecs; // 72 +constexpr int32_t kIdxBytesPerItem = 4; + +__device__ __forceinline__ void pdlWait() +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif +} + +__device__ __forceinline__ void pdlLaunchDependents() +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); +#endif +} + +//! Same arithmetic as quantCopy<__nv_bfloat16, 8> in mlaKernels.cu: bf16 -> fp32, scale, saturating E4M3. +__device__ __forceinline__ uint2 quantizeBf16x8(uint4 const& src, float scale) +{ + __nv_bfloat162 const* in = reinterpret_cast<__nv_bfloat162 const*>(&src); + uint2 out; + __nv_fp8x2_e4m3* dst = reinterpret_cast<__nv_fp8x2_e4m3*>(&out); +#pragma unroll + for (int32_t j = 0; j < 4; ++j) + { + float2 v = __bfloat1622float2(in[j]); + v.x *= scale; + v.y *= scale; + dst[j] = __nv_fp8x2_e4m3(v); + } + return out; +} + +//! GPT-J RoPE on 8 bf16 values (4 adjacent pairs), identical to the per-pair transform of the native kernel. +__device__ __forceinline__ void rotateBf16x8(uint4& data, float2 const (&coef)[4]) +{ + __nv_bfloat162* pairs = reinterpret_cast<__nv_bfloat162*>(&data); +#pragma unroll + for (int32_t j = 0; j < 4; ++j) + { + pairs[j] = mmha::rotary_embedding_transform(pairs[j], coef[j]); + } +} + +//! Native unravel of a flat byte index over the strided 4-D indexer cache view (indexerKCacheScatter.cu); +//! 32-bit arithmetic when the index fits, same result. +__device__ __forceinline__ int64_t idxFlatToOffset(int64_t flat, int32_t const (&dims)[4], int64_t const (&strides)[4]) +{ + int32_t i0, i1, i2, i3; + if (flat >= 0 && flat <= 0x7fffffff) + { + uint32_t f = static_cast(flat); + i3 = static_cast(f % static_cast(dims[3])); + f /= static_cast(dims[3]); + i2 = static_cast(f % static_cast(dims[2])); + f /= static_cast(dims[2]); + i1 = static_cast(f % static_cast(dims[1])); + f /= static_cast(dims[1]); + i0 = static_cast(f); + } + else + { + i3 = flat % dims[3]; + flat /= dims[3]; + i2 = flat % dims[2]; + flat /= dims[2]; + i1 = flat % dims[1]; + flat /= dims[1]; + i0 = flat; + } + return i0 * strides[0] + i1 * strides[1] + i2 * strides[2] + i3 * strides[3]; +} + +enum class RopeItem : int32_t +{ + kNone = 0, + kQNope, + kQRope, + kKRope, + kKNope, + kIdxData, + kIdxScale, +}; + +//! One thread per 16-byte work item of one token. Per token the item space is +//! [0, H*64) q_nope FP8 quantization (head = i / 64, vec = i % 64) +//! [H*64, H*72) q_pe RoPE + FP8 (head, vec) +//! [H*72, H*72+8) k_pe RoPE + FP8 cache write (vec) +//! [H*72+8, H*72+72) compressed-KV FP8 cache write (vec) +//! then idxHeadBytes/4 indexer K data items (4 bytes each) and one indexer scale item. +//! Per-step metadata (cache lengths, KV block offsets, indexer slot mappings) and the item's source data are read +//! only after griddepcontrol.wait (see the header). +__global__ void __launch_bounds__(kRopeThreads) ropeAppendScatterKernel(RopeAppendParams const p) +{ + int32_t const token = blockIdx.x; + int32_t const item = blockIdx.y * kRopeThreads + threadIdx.x; + int32_t const qNopeEnd = p.numHeads * kNopeVecs; + int32_t const qRopeEnd = qNopeEnd + p.numHeads * kRopeVecs; + int32_t const kRopeEnd = qRopeEnd + kRopeVecs; + int32_t const kNopeEnd = kRopeEnd + kNopeVecs; + int32_t const idxDataEnd = kNopeEnd + p.idxHeadBytes / kIdxBytesPerItem; + + bool const ropeToken = token < p.numTokens; + bool const idxToken = token < p.idxNumTokens; + + RopeItem kind = RopeItem::kNone; + int32_t head = 0; + int32_t vec = 0; + if (item < qNopeEnd) + { + kind = ropeToken ? RopeItem::kQNope : RopeItem::kNone; + head = item / kNopeVecs; + vec = item % kNopeVecs; + } + else if (item < qRopeEnd) + { + kind = ropeToken ? RopeItem::kQRope : RopeItem::kNone; + head = (item - qNopeEnd) / kRopeVecs; + vec = (item - qNopeEnd) % kRopeVecs; + } + else if (item < kRopeEnd) + { + kind = ropeToken ? RopeItem::kKRope : RopeItem::kNone; + vec = item - qRopeEnd; + } + else if (item < kNopeEnd) + { + kind = ropeToken ? RopeItem::kKNope : RopeItem::kNone; + vec = item - kRopeEnd; + } + else if (item < idxDataEnd) + { + kind = idxToken ? RopeItem::kIdxData : RopeItem::kNone; + vec = item - kNopeEnd; + } + else if (item == idxDataEnd) + { + kind = idxToken ? RopeItem::kIdxScale : RopeItem::kNone; + } + + int32_t const batchIdx = token / p.seqLen; + int32_t const localIdx = token % p.seqLen; + bool const needPos = kind == RopeItem::kQRope || kind == RopeItem::kKRope || kind == RopeItem::kKNope; + bool const needCoef = kind == RopeItem::kQRope || kind == RopeItem::kKRope; + bool const needBlock = kind == RopeItem::kKRope || kind == RopeItem::kKNope; + bool const isIdx = kind == RopeItem::kIdxData || kind == RopeItem::kIdxScale; + + // Values that never change while kernels of the step run (the quantization scales) may be read before the wait; + // per-step metadata and producer outputs may not. + float scale = 1.F; + if (needBlock) + { + scale = p.quantScaleKv ? p.quantScaleKv[0] : 1.F; + } + else if (kind == RopeItem::kQNope || kind == RopeItem::kQRope) + { + scale = p.quantScaleQ ? p.quantScaleQ[0] : 1.F; + } + auto loadCoef = [&](int32_t pos, float2(&c)[4]) + { + float2 const* src = p.cosSin + static_cast(kRopeDim) * pos + (vec * kEltsPerVec) / 2; +#pragma unroll + for (int32_t j = 0; j < 4; ++j) + { + c[j] = src[j]; + } + }; + auto idxOffset = [&](int64_t slotData, int64_t slotScale) -> int64_t + { + if (slotData < 0 || slotScale < 0) + { + return -1; + } + return kind == RopeItem::kIdxData ? idxFlatToOffset(slotData + vec * kIdxBytesPerItem, p.idxDims, p.idxStrides) + : idxFlatToOffset(slotScale, p.idxDims, p.idxStrides); + }; + // Source data of this item; written by producers, so it is read after the wait. + uint4 vdata = make_uint4(0, 0, 0, 0); + uint32_t wdata = 0; + auto loadSource = [&]() + { + switch (kind) + { + case RopeItem::kQNope: + vdata = *reinterpret_cast( + p.fusedQ + (static_cast(token) * p.numHeads + head) * kLatentDim + vec * kEltsPerVec); + break; + case RopeItem::kQRope: + vdata = *reinterpret_cast(p.qPe + static_cast(token) * p.qPeStride + + static_cast(p.qPeLd) * head + vec * kEltsPerVec); + break; + case RopeItem::kKRope: + vdata = *reinterpret_cast( + p.latent + static_cast(token) * kLatentDim + kKvLoraRank + vec * kEltsPerVec); + break; + case RopeItem::kKNope: + vdata = *reinterpret_cast( + p.latent + static_cast(token) * kLatentDim + vec * kEltsPerVec); + break; + case RopeItem::kIdxData: + wdata = *reinterpret_cast( + p.idxK + static_cast(token) * p.idxHeadBytes + vec * kIdxBytesPerItem); + break; + case RopeItem::kIdxScale: + wdata = *reinterpret_cast(p.idxScale + static_cast(token) * kIdxBytesPerItem); + break; + default: break; + } + }; + + int32_t kvPos = 0; + float2 coef[4]; + __nv_fp8_e4m3* kBlock = nullptr; + int64_t idxDst = -1; + pdlWait(); + // Issue the source-data load first so it overlaps the metadata loads and the address math below. + loadSource(); + if (needPos) + { + kvPos = p.kvLens[batchIdx] - p.seqLen + localIdx; + } + if (needCoef) + { + loadCoef(kvPos, coef); + } + if (needBlock) + { + kBlock = reinterpret_cast<__nv_fp8_e4m3*>(p.kvCache.getKBlockPtr(batchIdx, kvPos)); + } + if (isIdx) + { + idxDst = idxOffset(p.idxSlotData[token], p.idxSlotScale[token]); + } + + switch (kind) + { + case RopeItem::kQNope: + { + size_t const off = (static_cast(token) * p.numHeads + head) * kLatentDim + vec * kEltsPerVec; + *reinterpret_cast(p.quantQ + off) = quantizeBf16x8(vdata, scale); + break; + } + case RopeItem::kQRope: + { + uint4 data = vdata; + rotateBf16x8(data, coef); + size_t const dst + = (static_cast(token) * p.numHeads + head) * kLatentDim + kKvLoraRank + vec * kEltsPerVec; + *reinterpret_cast(p.quantQ + dst) = quantizeBf16x8(data, scale); + break; + } + case RopeItem::kKRope: + { + uint4 data = vdata; + rotateBf16x8(data, coef); + int32_t const inBlockIdx = p.kvCache.getKVLocalIdx(kvPos, 0, kTotalVecsPerHead, kNopeVecs + vec); + *reinterpret_cast(kBlock + inBlockIdx * kEltsPerVec) = quantizeBf16x8(data, scale); + break; + } + case RopeItem::kKNope: + { + int32_t const inBlockIdx = p.kvCache.getKVLocalIdx(kvPos, 0, kTotalVecsPerHead, vec); + *reinterpret_cast(kBlock + inBlockIdx * kEltsPerVec) = quantizeBf16x8(vdata, scale); + if (vec == 0 && localIdx == 0) + { + p.cuQSeqlens[batchIdx + 1] = p.numHeads * p.seqLen * (batchIdx + 1); + } + break; + } + case RopeItem::kIdxData: + case RopeItem::kIdxScale: + { + if (idxDst >= 0) + { + *reinterpret_cast(p.idxCache + idxDst) = wdata; + } + break; + } + default: break; + } + + if (blockIdx.x == 0 && blockIdx.y == 0) + { + if (threadIdx.x == 0) + { + p.fmhaTileCounter[0] = 0; + p.cuQSeqlens[0] = 0; + float const dequantQ = p.dequantScaleQ ? p.dequantScaleQ[0] : 1.F; + float const dequantKv = p.dequantScaleKv ? p.dequantScaleKv[0] : 1.F; + float const quantO = p.quantScaleO ? p.quantScaleO[0] : 1.F; + if (p.bmm1Scale) + { + constexpr float kLog2e = 1.4426950408889634074F; + float const bmm1 = dequantQ * dequantKv * p.hostBmm1Scale; + p.bmm1Scale[0] = bmm1; + p.bmm1Scale[1] = bmm1 * kLog2e; + } + if (p.bmm2Scale) + { + p.bmm2Scale[0] = quantO * dequantKv; + } + } + if (threadIdx.x < 32) + { + // Exclusive prefix sum of the cache lengths -> cuKvSeqlens[0..batch]. + int32_t const lane = threadIdx.x; + int32_t const batch = p.numTokens / p.seqLen; + int32_t carry = 0; + for (int32_t base = 0; base <= batch; base += 32) + { + int32_t const idx = base + lane; + int32_t const len = idx < batch ? p.kvLens[idx] : 0; + int32_t incl = len; +#pragma unroll + for (int32_t off = 1; off < 32; off <<= 1) + { + int32_t const n = __shfl_up_sync(0xffffffffU, incl, off); + if (lane >= off) + { + incl += n; + } + } + if (idx <= batch) + { + p.cuKvSeqlens[idx] = incl - len + carry; + } + carry += __shfl_sync(0xffffffffU, incl, 31); + } + } + } + + pdlLaunchDependents(); +} + +__device__ __forceinline__ int32_t remapOne( + int32_t tok, int32_t layerId, int32_t blockSize, int32_t strideFactor, int32_t base) +{ + if (tok < 0 || base < 0) + { + return -1; + } + int32_t const inblockOff = tok % blockSize + layerId * blockSize; + return base * strideFactor + inblockOff; +} + +//! One thread per (row, column) element, like the native remap, so the scattered block-table reads spread over +//! many SMs; the group members' layer offsets are hoisted out of the store loop. Everything is read after the wait. +__global__ void __launch_bounds__(kBridgeThreads) topkBridgeKernel(TopkBridgeParams const p) +{ + int32_t const rowSlot = blockIdx.x; + int32_t const col = blockIdx.y * kBridgeThreads + threadIdx.x; + bool const isOutRow = rowSlot < p.numRows; + + int32_t req = 0; + int32_t srcRow = rowSlot; + int32_t layerIds[kMaxGroupSize]; + pdlWait(); + if (isOutRow) + { + req = p.reqIdx[rowSlot] - p.reqIdxBias; +#pragma unroll + for (int32_t g = 0; g < kMaxGroupSize; ++g) + { + layerIds[g] = g < p.groupSize ? (p.layerIds ? p.layerIds[g] : p.layerId) : 0; + } + } + else + { + int32_t const i = rowSlot - p.numRows; + if (p.numAccepted == nullptr) + { + srcRow = i * p.nextN + p.nextN - 1; + } + else + { + int64_t const acc = p.numAcceptedIs64 ? static_cast(p.numAccepted)[i] + : static_cast(static_cast(p.numAccepted)[i]); + int64_t off = acc - 1; + off = off < 0 ? 0 : (off > p.nextN - 1 ? p.nextN - 1 : off); + srcRow = static_cast(static_cast(i) * p.nextN + off); + } + } + + if (col < p.numTopk) + { + int32_t const tok = p.topk[static_cast(srcRow) * p.topkStride0 + col]; + if (isOutRow) + { + int32_t base = -1; + if (tok >= 0) + { + int32_t const blockId = tok / p.blockSize; + if (blockId < p.maxBlocksPerReq) + { + base = p.blockTable[req * p.btStride0 + blockId * p.btStride1]; + } + } + int32_t* dst = p.out + static_cast(rowSlot) * p.outStride0 + col; +#pragma unroll + for (int32_t g = 0; g < kMaxGroupSize; ++g) + { + if (g < p.groupSize) + { + dst[g * p.outStrideG] = remapOne(tok, layerIds[g], p.blockSize, p.strideFactor, base); + } + } + } + bool const doCopy = p.copyDst != nullptr + && ((p.copyMode == 0 && isOutRow && rowSlot < p.copyRows) || (p.copyMode == 1 && !isOutRow)); + if (doCopy) + { + int32_t const copyRow = p.copyRowStart + (isOutRow ? rowSlot : rowSlot - p.numRows); + p.copyDst[static_cast(copyRow) * p.copyStride0 + col] = tok; + } + } + + pdlLaunchDependents(); +} + +template +void launchWithOptionalPdl(Kernel kernel, dim3 grid, dim3 block, bool enablePdl, cudaStream_t stream, Params const& p) +{ + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = enablePdl ? 1 : 0; + config.numAttrs = 1; + config.attrs = attrs; + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, kernel, p)); +} + +} // namespace + +void invokeRopeAppendScatter(RopeAppendParams const& params, bool enablePdl, cudaStream_t stream) +{ + int32_t const gridTokens = std::max(params.numTokens, params.idxNumTokens); + if (gridTokens == 0) + { + return; + } + TLLM_CHECK_WITH_INFO(params.seqLen > 0, "seqLen must be positive"); + TLLM_CHECK_WITH_INFO(params.cosSinRows > 0, "the RoPE table must have at least one position"); + TLLM_CHECK_WITH_INFO(params.idxHeadBytes % kIdxBytesPerItem == 0, "indexer head bytes must be a multiple of 4"); + int32_t const itemsPerToken + = params.numHeads * kTotalVecsPerHead + kTotalVecsPerHead + params.idxHeadBytes / kIdxBytesPerItem + 1; + dim3 const grid(gridTokens, (itemsPerToken + kRopeThreads - 1) / kRopeThreads); + dim3 const block(kRopeThreads); + launchWithOptionalPdl(ropeAppendScatterKernel, grid, block, enablePdl, stream, params); +} + +void invokeTopkBridge(TopkBridgeParams const& params, bool enablePdl, cudaStream_t stream) +{ + int32_t const extraRows = params.copyMode == 1 ? params.copyRows : 0; + int32_t const rows = params.numRows + extraRows; + if (rows == 0 || params.numTopk == 0) + { + return; + } + TLLM_CHECK_WITH_INFO(params.groupSize >= 1 && params.groupSize <= kMaxGroupSize, + "glm_moe_dsa_topk_index_bridge: group sizes must be 1..%d, got %d", kMaxGroupSize, params.groupSize); + dim3 const grid(rows, (params.numTopk + kBridgeThreads - 1) / kBridgeThreads); + dim3 const block(kBridgeThreads); + launchWithOptionalPdl(topkBridgeKernel, grid, block, enablePdl, stream, params); +} + +} // namespace kernels::glm_moe_dsa::rope_kv_append + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/ropeKvAppend.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/ropeKvAppend.h new file mode 100644 index 000000000000..8d4ea666f26b --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/ropeKvAppend.h @@ -0,0 +1,136 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" +#include "tensorrt_llm/kernels/kvCacheUtils.h" + +#include +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa::rope_kv_append +{ + +//! MLA absorption-mode latent width and RoPE width supported by the fused kernels. +constexpr int32_t kKvLoraRank = 512; +constexpr int32_t kRopeDim = 64; +constexpr int32_t kLatentDim = kKvLoraRank + kRopeDim; + +//! Memory order of both kernels: they may start early as programmatic dependents (PDL) and read everything a +//! producer writes, including per-step metadata (cache lengths, KV block offsets, indexer slot mappings, request +//! indices, MTP accepted counts), only after griddepcontrol.wait. Only the quantization scales, which no kernel of the +//! step writes, are read before the wait. + +//! Generation-phase MLA RoPE + FP8 query quantization + FP8 latent KV append, optionally fused with the DSA +//! indexer K-cache scatter. With an FP8 MLA cache the outputs are byte-identical to +//! applyMLARopeAndAssignQKVKernelGeneration (absorption mode, rope_append, no helix) followed by +//! indexerKCacheScatterUnifiedKernel. +struct RopeAppendParams +{ + // Query side. + __nv_bfloat16 const* fusedQ{nullptr}; //!< [numTokens, numHeads, kLatentDim]; the first kKvLoraRank are read. + __nv_bfloat16 const* qPe{nullptr}; //!< [numTokens, numHeads, kRopeDim] with strides (qPeStride, qPeLd, 1). + int32_t qPeLd{0}; + int32_t qPeStride{0}; + uint8_t* quantQ{nullptr}; //!< FP8 E4M3 [numTokens, numHeads, kLatentDim]. + float const* quantScaleQ{nullptr}; //!< Optional; 1.0 when null. + + // KV side. + __nv_bfloat16 const* latent{nullptr}; //!< [numTokens, kLatentDim] rows (packed). + float2 const* cosSin{nullptr}; //!< Rotary table addressed as cosSin[kRopeDim * position + d / 2]. + int32_t cosSinRows{0}; //!< Positions in the rotary table. + int32_t const* kvLens{nullptr}; //!< [batch] cache lengths after this step's append. + float const* quantScaleKv{nullptr}; //!< Optional; 1.0 when null. + KVBlockArray kvCache; + + // Scheduler buffers consumed by the FP8 MLA generation FMHA. + int32_t* cuQSeqlens{nullptr}; + int32_t* cuKvSeqlens{nullptr}; + uint32_t* fmhaTileCounter{nullptr}; + float* bmm1Scale{nullptr}; //!< Optional [2]. + float* bmm2Scale{nullptr}; //!< Optional [1]. + float const* dequantScaleQ{nullptr}; //!< Optional; 1.0 when null. + float const* dequantScaleKv{nullptr}; //!< Optional; 1.0 when null. + float const* quantScaleO{nullptr}; //!< Optional; 1.0 when null. + float hostBmm1Scale{1.F}; + + int32_t numTokens{0}; + int32_t seqLen{1}; + int32_t numHeads{0}; + + // Optional DSA indexer K-cache scatter (skipped when idxNumTokens == 0). + uint8_t const* idxK{nullptr}; //!< [idxNumTokens, idxHeadBytes] bytes. + uint8_t const* idxScale{nullptr}; //!< [idxNumTokens, 4] bytes. + uint8_t* idxCache{nullptr}; //!< Strided 4-D byte view of one layer's indexer pool. + int64_t const* idxSlotData{nullptr}; + int64_t const* idxSlotScale{nullptr}; + int32_t idxNumTokens{0}; + int32_t idxHeadBytes{128}; + int32_t idxDims[4]{0, 0, 0, 0}; + int64_t idxStrides[4]{0, 0, 0, 0}; +}; + +//! Fused bridge from per-row local top-k token ids to global KV-pool indices, optionally writing the native +//! top-k reuse copy in the same pass. Output values are identical to convertReqIndexToGlobal[Grouped]Kernel. +struct TopkBridgeParams +{ + int32_t const* topk{nullptr}; //!< Source rows [srcRows, numTopk] with row stride topkStride0, unit column stride. + int64_t topkStride0{0}; + int32_t numRows{0}; //!< Rows remapped to global indices. + int32_t numTopk{0}; + + int32_t const* reqIdx{nullptr}; //!< [numRows] request index per row, before subtracting reqIdxBias. + int32_t reqIdxBias{0}; + int32_t const* blockTable{nullptr}; + int64_t btStride0{0}; + int64_t btStride1{1}; + int32_t maxBlocksPerReq{0}; + int32_t blockSize{0}; + int32_t strideFactor{0}; + + int32_t const* layerIds{nullptr}; //!< Optional [groupSize]; when null the single layerId is used. + int32_t layerId{0}; + int32_t groupSize{1}; + int32_t* out{nullptr}; //!< [groupSize, numRows, numTopk] with strides (outStrideG, outStride0, 1). + int64_t outStrideG{0}; + int64_t outStride0{0}; + + //! Copy rows: dst[copyRowStart + i, :numTopk] = topk[srcRow(i), :] for i in [0, copyRows). + //! srcRow(i) = i (copyMode 0) or the MTP last-accepted row i * nextN + clamp(numAccepted[i] - 1, 0, nextN - 1) + //! (copyMode 1; i * nextN + nextN - 1 when numAccepted is null). + int32_t* copyDst{nullptr}; + int64_t copyStride0{0}; + int32_t copyRows{0}; + int32_t copyRowStart{0}; + int32_t copyMode{0}; + int32_t nextN{1}; + void const* numAccepted{nullptr}; //!< Already offset to the first generation request. + bool numAcceptedIs64{false}; +}; + +void invokeRopeAppendScatter(RopeAppendParams const& params, bool enablePdl, cudaStream_t stream); + +void invokeTopkBridge(TopkBridgeParams const& params, bool enablePdl, cudaStream_t stream); + +} // namespace kernels::glm_moe_dsa::rope_kv_append + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/sparseMlaDecode.cu b/cpp/tensorrt_llm/kernels/glmMoeDsa/sparseMlaDecode.cu new file mode 100644 index 000000000000..b9c0a82f04ec --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/sparseMlaDecode.cu @@ -0,0 +1,714 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/envUtils.h" +#include "tensorrt_llm/kernels/glmMoeDsa/sparseMlaDecode.h" + +#include +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa::sparse_mla_decode +{ +namespace +{ + +// Work split: one CTA handles one query row and 128 consecutive positions of its index list. +// The 16 CTAs of a row form one thread-block cluster. Each CTA pushes its unnormalized partial +// output and softmax statistics into the shared memory of the CTA that finalizes those output +// dims with asynchronous stores that complete on that CTA's mbarrier. +int constexpr kKeysPerCta = 128; +int constexpr kNumSplits = kTopK / kKeysPerCta; +int constexpr kNumWarps = 8; +int constexpr kThreads = kNumWarps * 32; +int constexpr kKeysPerWarp = kKeysPerCta / kNumWarps; +int constexpr kDimsPerWarp = kHeadDimV / kNumWarps; +int constexpr kDimsPerCtaOut = kHeadDimV / kNumSplits; +[[maybe_unused]] int constexpr kQkSteps = kHeadDimQk / 32; +[[maybe_unused]] int constexpr kPvSteps = kKeysPerCta / 32; +[[maybe_unused]] int constexpr kPvTilesPerWarp = kDimsPerWarp / 16; +[[maybe_unused]] int constexpr kQChunks = kNumHeads * kHeadDimQk / 16; +// Probabilities are scaled to the e4m3 range before rounding, as in the native kernel. +[[maybe_unused]] float constexpr kProbScale = 448.F; + +// Shared-memory pitches, padded so that ldmatrix and fragment loads are bank-conflict free: K rows +// shift by 20 banks (eight ldmatrix rows hit disjoint 16-byte bank groups); the 64-bit query +// fragment loads of rows g = 0-3 cover banks 8g..8g+7 (pitch = 8 banks mod 32); the 32-bit +// probability fragment loads of rows g = 0-7 cover banks 4g..4g+3. +int constexpr kKRowBytes = kHeadDimQk + 16; +int constexpr kQRowHalves = kHeadDimQk + 16; +int constexpr kPRowHalves = kKeysPerCta + 8; + +int constexpr kKTileOffset = 0; +int constexpr kQTileOffset = kKTileOffset + kKeysPerCta * kKRowBytes; +int constexpr kPTileOffset = kQTileOffset + kNumHeads * kQRowHalves * 2; +// [source split][head][dim within this CTA's 32 output dims] FP32 partial outputs. +int constexpr kRecvOffset = kPTileOffset + kNumHeads * kPRowHalves * 2; +// [max of heads 0-7, sum of heads 0-7][source split] FP32. +int constexpr kRecvStatsOffset = kRecvOffset + kNumSplits * kNumHeads * kDimsPerCtaOut * 4; +int constexpr kRedMaxOffset = kRecvStatsOffset + kNumSplits * 2 * kNumHeads * 4; +int constexpr kRedSumOffset = kRedMaxOffset + kNumWarps * kNumHeads * 4; +// mbarriers: one per warp for its row gather, one for the pushed partials of this CTA. +int constexpr kBarrierOffset = kRedSumOffset + kNumWarps * kNumHeads * 4; +int constexpr kSmemBytes = kBarrierOffset + (kNumWarps + 1) * 8; +// Bytes each CTA receives: 16 sources x (8 heads x 32 dims + 16 statistics) FP32 values. +[[maybe_unused]] uint32_t constexpr kRecvBytes = kNumSplits * (kNumHeads * kDimsPerCtaOut + 2 * kNumHeads) * 4; + +static_assert(kTopK % kKeysPerCta == 0, "the index list must split evenly"); +static_assert(kNumSplits <= 16, "the merge cluster holds at most 16 CTAs"); +static_assert(kKeysPerWarp == 16, "one m16 tile of keys per warp"); +static_assert(kNumHeads == 8, "the eight heads form the n8 side of the MMAs"); +static_assert(kDimsPerCtaOut == 32 && kDimsPerWarp == 2 * kDimsPerCtaOut, "each warp feeds two merge CTAs"); +static_assert(kDimsPerCtaOut * kNumHeads == kThreads, "one merged output element per thread"); +static_assert(kKRowBytes % 16 == 0 && kQRowHalves % 8 == 0 && kPRowHalves % 8 == 0, "16-byte aligned shared rows"); +static_assert( + kKRowBytes / 4 % 32 == 20 && kQRowHalves / 2 % 32 == 8 && kPRowHalves / 2 % 32 == 4, "bank-conflict-free pitches"); +static_assert(kBarrierOffset % 8 == 0, "mbarrier alignment"); +static_assert(kQChunks <= 2 * kThreads, "two query chunks per thread at most"); +static_assert(2 * kNumHeads * kNumSplits == kThreads, "one statistics store per thread"); +static_assert( + kNumWarps == kNumHeads && (kNumSplits & (kNumSplits - 1)) == 0, "merge: one warp per head, one lane per split"); + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + +__device__ __forceinline__ uint32_t smemAddr(void const* ptr) +{ + return static_cast(__cvta_generic_to_shared(ptr)); +} + +__device__ __forceinline__ void mbarrierInit(uint32_t bar, uint32_t count) +{ + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(bar), "r"(count) : "memory"); +} + +__device__ __forceinline__ void fenceBarrierInit() +{ + asm volatile("fence.mbarrier_init.release.cluster;\n" ::: "memory"); +} + +__device__ __forceinline__ void mbarrierArriveExpectTx(uint32_t bar, uint32_t bytes) +{ + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"(bar), "r"(bytes) : "memory"); +} + +__device__ __forceinline__ void mbarrierArrive(uint32_t bar) +{ + asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0];\n" ::"r"(bar) : "memory"); +} + +__device__ __forceinline__ void mbarrierWait(uint32_t bar, uint32_t parity) +{ + uint32_t done = 0; + while (!done) + { + asm volatile( + "{\n" + ".reg .pred p;\n" + "mbarrier.try_wait.parity.shared::cta.b64 p, [%1], %2;\n" + "selp.u32 %0, 1, 0, p;\n" + "}\n" + : "=r"(done) + : "r"(bar), "r"(parity) + : "memory"); + } +} + +// The same wait with cluster-scope acquire: the arrivals and transaction bytes it waits for come +// from asynchronous stores issued by other CTAs of the cluster. +__device__ __forceinline__ void mbarrierWaitCluster(uint32_t bar, uint32_t parity) +{ + uint32_t done = 0; + while (!done) + { + asm volatile( + "{\n" + ".reg .pred p;\n" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 p, [%1], %2;\n" + "selp.u32 %0, 1, 0, p;\n" + "}\n" + : "=r"(done) + : "r"(bar), "r"(parity) + : "memory"); + } +} + +__device__ __forceinline__ void bulkCopyG2S(uint32_t dst, void const* src, uint32_t bytes, uint32_t bar) +{ + asm volatile("cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];\n" ::"r"(dst), + "l"(src), "r"(bytes), "r"(bar) + : "memory"); +} + +__device__ __forceinline__ void gridDependencyWait() +{ + asm volatile("griddepcontrol.wait;\n" ::: "memory"); +} + +__device__ __forceinline__ void gridDependencyLaunchDependents() +{ + asm volatile("griddepcontrol.launch_dependents;\n" ::: "memory"); +} + +__device__ __forceinline__ void clusterArriveRelaxed() +{ + asm volatile("barrier.cluster.arrive.relaxed.aligned;\n" ::: "memory"); +} + +__device__ __forceinline__ void clusterWait() +{ + asm volatile("barrier.cluster.wait.acquire.aligned;\n" ::: "memory"); +} + +__device__ __forceinline__ uint32_t mapToRank(uint32_t localAddr, uint32_t rank) +{ + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;\n" : "=r"(remote) : "r"(localAddr), "r"(rank)); + return remote; +} + +// Asynchronous stores into a peer CTA's shared memory that complete-tx on the peer's mbarrier. +__device__ __forceinline__ void storeAsync(uint32_t remoteAddr, float value, uint32_t remoteBar) +{ + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.f32 [%0], %1, [%2];\n" ::"r"(remoteAddr), + "f"(value), "r"(remoteBar) + : "memory"); +} + +__device__ __forceinline__ void storeAsync2(uint32_t remoteAddr, float a, float b, uint32_t remoteBar) +{ + asm volatile( + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v2.f32 [%0], {%1, %2}, [%3];\n" ::"r"(remoteAddr), + "f"(a), "f"(b), "r"(remoteBar) + : "memory"); +} + +__device__ __forceinline__ void ldmatrixX4(uint32_t (&r)[4], uint32_t addr) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(addr)); +} + +__device__ __forceinline__ void ldmatrixX4Trans(uint32_t (&r)[4], uint32_t addr) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(addr)); +} + +// D += A * B with A 16x16 f16 (row), B 16x8 f16 (col), FP32 accumulation. +__device__ __forceinline__ void mmaF16( + float (&d)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t b0, uint32_t b1) +{ + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, " + "{%0, %1, %2, %3};\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)); +} + +__device__ __forceinline__ float fastExp2(float x) +{ + float y; + asm volatile("ex2.approx.ftz.f32 %0, %1;\n" : "=f"(y) : "f"(x)); + return y; +} + +// Two e4m3 values (low byte first) to two FP16 values; exact. +__device__ __forceinline__ uint32_t e4m3x2ToHalf2(uint32_t pair) +{ + __half2_raw const h = __nv_cvt_fp8x2_to_halfraw2(static_cast<__nv_fp8x2_storage_t>(pair & 0xffffU), __NV_E4M3); + return static_cast(h.x) | (static_cast(h.y) << 16); +} + +// e4m3(448 * p) as an exact FP16 value. +__device__ __forceinline__ __half roundProbability(float p) +{ + __nv_fp8_storage_t const q = __nv_cvt_float_to_fp8(p * kProbScale, __NV_SATFINITE, __NV_E4M3); + return __half(__nv_cvt_fp8_to_halfraw(q, __NV_E4M3)); +} + +#endif // defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + +#ifdef GLM_MOE_DSA_SPARSE_MLA_DECODE_TIMING +// Development timeline: thread 0 of each CTA records the SM clock at phase boundaries and the +// global timer at entry and exit. +__device__ unsigned long long gTimeline[kTimelineMaxCtas][kTimelineSlots]; +#endif + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) +#ifdef GLM_MOE_DSA_SPARSE_MLA_DECODE_TIMING +__device__ __forceinline__ void markTimeline(int slot) +{ + if (threadIdx.x == 0 && blockIdx.x < kTimelineMaxCtas) + { + gTimeline[blockIdx.x][slot] = clock64(); + } +} + +__device__ __forceinline__ void markGlobalTime(int slot) +{ + if (threadIdx.x == 0 && blockIdx.x < kTimelineMaxCtas) + { + unsigned long long t; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t)); + gTimeline[blockIdx.x][slot] = t; + } +} +#else +__device__ __forceinline__ void markTimeline(int) {} + +__device__ __forceinline__ void markGlobalTime(int) {} +#endif +#endif // defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + +__global__ void __launch_bounds__(kThreads) sparseMlaDecodeKernel(SparseMlaDecodeParams const params) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + extern __shared__ __align__(128) uint8_t smem[]; + + int const tid = threadIdx.x; + int const warp = tid >> 5; + int const lane = tid & 31; + int const groupId = lane >> 2; + int const tig = lane & 3; + int const row = blockIdx.x / kNumSplits; + int const split = blockIdx.x % kNumSplits; + + uint8_t* kTile = smem + kKTileOffset; + __half* qTile = reinterpret_cast<__half*>(smem + kQTileOffset); + __half* pTile = reinterpret_cast<__half*>(smem + kPTileOffset); + float* recv = reinterpret_cast(smem + kRecvOffset); + float* recvStats = reinterpret_cast(smem + kRecvStatsOffset); + float* redMax = reinterpret_cast(smem + kRedMaxOffset); + float* redSum = reinterpret_cast(smem + kRedSumOffset); + uint32_t const gatherBar = smemAddr(smem + kBarrierOffset + warp * 8); + uint32_t const recvBar = smemAddr(smem + kBarrierOffset + kNumWarps * 8); + + markGlobalTime(0); + markTimeline(1); + if (lane == 0) + { + mbarrierInit(gatherBar, kKeysPerWarp); + } + if (tid == 0) + { + mbarrierInit(recvBar, 1); + } + // The initializations must be visible to the bulk-copy engine and to the peers' async stores. + fenceBarrierInit(); + __syncthreads(); + clusterArriveRelaxed(); + + // Everything below reads tensors written by preceding kernels in the stream. The kernel can be + // resident before those kernels finish, so these tensors are read with coherent loads, not + // through the non-coherent read-only path. + gridDependencyWait(); + markTimeline(2); + + // The query does not depend on the index list: start its loads first. + uint8_t const* qRow = static_cast(params.q) + static_cast(row) * kNumHeads * kHeadDimQk; + uint4 qRaw[2]; +#pragma unroll + for (int i = 0; i < 2; ++i) + { + int const chunk = tid + i * kThreads; + if (chunk < kQChunks) + { + qRaw[i] = reinterpret_cast(qRow)[chunk]; + } + } + + int32_t poolRow = -1; + if (lane < kKeysPerWarp) + { + int const seq = row / params.seqLenQ; + int const rowInSeq = row - seq * params.seqLenQ; + int const key = split * kKeysPerCta + warp * kKeysPerWarp + lane; + int32_t const listed = params.indices[static_cast(row) * params.indicesRowStride + key]; + int const usable = params.seqLensKv[seq] - (params.seqLenQ - 1 - rowInSeq); + poolRow = key < usable ? listed : -1; + } + uint32_t const validBits = __ballot_sync(0xffffffffU, poolRow >= 0); + markTimeline(3); + + // Gather this warp's 16 rows with one 576-byte bulk copy each; unusable rows are zeroed so that + // zero probabilities never meet non-finite bytes. + if (lane < kKeysPerWarp) + { + uint8_t* dst = kTile + (warp * kKeysPerWarp + lane) * kKRowBytes; + if (poolRow >= 0) + { + mbarrierArriveExpectTx(gatherBar, kHeadDimQk); + bulkCopyG2S(smemAddr(dst), + static_cast(params.kvPool) + static_cast(poolRow) * kHeadDimQk, kHeadDimQk, + gatherBar); + } + else + { +#pragma unroll 4 + for (int c = 0; c < kHeadDimQk / 16; ++c) + { + *reinterpret_cast(dst + c * 16) = make_uint4(0U, 0U, 0U, 0U); + } + mbarrierArrive(gatherBar); + } + } + + // FP16 copy of the query, shared by all warps. +#pragma unroll + for (int i = 0; i < 2; ++i) + { + int const chunk = tid + i * kThreads; + if (chunk < kQChunks) + { + int const head = chunk / (kHeadDimQk / 16); + int const col = (chunk - head * (kHeadDimQk / 16)) * 16; + uint4 lo; + uint4 hi; + lo.x = e4m3x2ToHalf2(qRaw[i].x); + lo.y = e4m3x2ToHalf2(qRaw[i].x >> 16); + lo.z = e4m3x2ToHalf2(qRaw[i].y); + lo.w = e4m3x2ToHalf2(qRaw[i].y >> 16); + hi.x = e4m3x2ToHalf2(qRaw[i].z); + hi.y = e4m3x2ToHalf2(qRaw[i].z >> 16); + hi.z = e4m3x2ToHalf2(qRaw[i].w); + hi.w = e4m3x2ToHalf2(qRaw[i].w >> 16); + uint4* dst = reinterpret_cast(qTile + head * kQRowHalves + col); + dst[0] = lo; + dst[1] = hi; + } + } + float const scaleLog2 = params.bmm1Scale[1]; + float const outScale = params.bmm2Scale[0]; + + gridDependencyLaunchDependents(); + __syncthreads(); + + // Every peer has initialized its barriers once phase 0 of the cluster barrier completes; waiting + // here overlaps that with the gather. Pushes into peers start after the softmax. + clusterWait(); + mbarrierWait(gatherBar, 0); + __syncwarp(); + markTimeline(4); + + // S^T = K * Q^T for this warp's 16 keys: keys on the M side, heads on the N side, FP16 MMAs on + // converted e4m3 keys. Within each 16-value k chunk the MMA k-slot order is permuted so that + // the two halves of an ldmatrix register map straight onto the A fragment: slots 2t, 2t+1 hold + // k = 4t, 4t+1 and slots 2t+8, 2t+9 hold k = 4t+2, 4t+3. The FP16 query is read in that order. + float sAcc[4]; + { + float acc[2][4] = {}; + uint32_t const aAddr = smemAddr(kTile + (warp * kKeysPerWarp + (lane & 15)) * kKRowBytes + (lane >> 4) * 16); + __half const* qFrag = qTile + groupId * kQRowHalves + 4 * tig; +#pragma unroll + for (int kk = 0; kk < kQkSteps; ++kk) + { + uint32_t r[4]; + ldmatrixX4(r, aAddr + kk * 32); + uint2 const b0 = *reinterpret_cast(qFrag + kk * 32); + uint2 const b1 = *reinterpret_cast(qFrag + kk * 32 + 16); + mmaF16(acc[0], e4m3x2ToHalf2(r[0]), e4m3x2ToHalf2(r[1]), e4m3x2ToHalf2(r[0] >> 16), + e4m3x2ToHalf2(r[1] >> 16), b0.x, b0.y); + mmaF16(acc[1], e4m3x2ToHalf2(r[2]), e4m3x2ToHalf2(r[3]), e4m3x2ToHalf2(r[2] >> 16), + e4m3x2ToHalf2(r[3] >> 16), b1.x, b1.y); + } +#pragma unroll + for (int i = 0; i < 4; ++i) + { + sAcc[i] = acc[0][i] + acc[1][i]; + } + } + markTimeline(5); + + // Fragment element (key, head): 0 -> (groupId, 2*tig), 1 -> (groupId, 2*tig+1), 2/3 -> key + 8. + bool const valid0 = (validBits >> groupId) & 1U; + bool const valid1 = (validBits >> (groupId + 8)) & 1U; + float s[4]; + s[0] = valid0 ? sAcc[0] * scaleLog2 : -INFINITY; + s[1] = valid0 ? sAcc[1] * scaleLog2 : -INFINITY; + s[2] = valid1 ? sAcc[2] * scaleLog2 : -INFINITY; + s[3] = valid1 ? sAcc[3] * scaleLog2 : -INFINITY; + + float maxA = fmaxf(s[0], s[2]); + float maxB = fmaxf(s[1], s[3]); +#pragma unroll + for (int offset = 4; offset < 32; offset <<= 1) + { + maxA = fmaxf(maxA, __shfl_xor_sync(0xffffffffU, maxA, offset)); + maxB = fmaxf(maxB, __shfl_xor_sync(0xffffffffU, maxB, offset)); + } + if (lane < 4) + { + redMax[warp * kNumHeads + 2 * tig] = maxA; + redMax[warp * kNumHeads + 2 * tig + 1] = maxB; + } + __syncthreads(); + + float headMaxA = -INFINITY; + float headMaxB = -INFINITY; +#pragma unroll + for (int w = 0; w < kNumWarps; ++w) + { + headMaxA = fmaxf(headMaxA, redMax[w * kNumHeads + 2 * tig]); + headMaxB = fmaxf(headMaxB, redMax[w * kNumHeads + 2 * tig + 1]); + } + float p[4]; + p[0] = s[0] == -INFINITY ? 0.F : fastExp2(s[0] - headMaxA); + p[1] = s[1] == -INFINITY ? 0.F : fastExp2(s[1] - headMaxB); + p[2] = s[2] == -INFINITY ? 0.F : fastExp2(s[2] - headMaxA); + p[3] = s[3] == -INFINITY ? 0.F : fastExp2(s[3] - headMaxB); + + float sumA = p[0] + p[2]; + float sumB = p[1] + p[3]; +#pragma unroll + for (int offset = 4; offset < 32; offset <<= 1) + { + sumA += __shfl_xor_sync(0xffffffffU, sumA, offset); + sumB += __shfl_xor_sync(0xffffffffU, sumB, offset); + } + if (lane < 4) + { + redSum[warp * kNumHeads + 2 * tig] = sumA; + redSum[warp * kNumHeads + 2 * tig + 1] = sumB; + } + { + int const key = warp * kKeysPerWarp + groupId; + pTile[(2 * tig) * kPRowHalves + key] = roundProbability(p[0]); + pTile[(2 * tig + 1) * kPRowHalves + key] = roundProbability(p[1]); + pTile[(2 * tig) * kPRowHalves + key + 8] = roundProbability(p[2]); + pTile[(2 * tig + 1) * kPRowHalves + key + 8] = roundProbability(p[3]); + } + __syncthreads(); + + // This CTA's softmax statistics to every CTA of the cluster: thread t sends value t / 16 + // (maxima of heads 0-7, then sums) to rank t % 16. + { + int const index = tid / kNumSplits; + int const dest = tid % kNumSplits; + int const head = index & (kNumHeads - 1); + bool const isMax = index < kNumHeads; + float value = isMax ? -INFINITY : 0.F; +#pragma unroll + for (int w = 0; w < kNumWarps; ++w) + { + value = isMax ? fmaxf(value, redMax[w * kNumHeads + head]) : value + redSum[w * kNumHeads + head]; + } + storeAsync(mapToRank(smemAddr(recvStats + index * kNumSplits + split), dest), value, mapToRank(recvBar, dest)); + } + markTimeline(6); + + // O^T = V^T * P^T for this warp's 64 output dims, one 16-dim tile at a time; each finished tile + // is pushed to the CTA that finalizes its dims while the next tile is computed. Transposing the + // e4m3 V rows with b16 ldmatrix yields keys (2t, 2t+1) x dims (2g, 2g+1) per register; a byte + // permute splits it into the two dims, so M row g of a tile is dim 2g and M row g+8 is dim 2g+1. + { + uint32_t pFrag[kPvSteps][2][2]; + __half const* pRow = pTile + groupId * kPRowHalves + 2 * tig; +#pragma unroll + for (int ks = 0; ks < kPvSteps; ++ks) + { +#pragma unroll + for (int half = 0; half < 2; ++half) + { + pFrag[ks][half][0] = *reinterpret_cast(pRow + ks * 32 + half * 16); + pFrag[ks][half][1] = *reinterpret_cast(pRow + ks * 32 + half * 16 + 8); + } + } + uint32_t const vAddr = smemAddr(kTile + lane * kKRowBytes + warp * kDimsPerWarp); + uint32_t const recvBase = smemAddr(recv); + float const invProbScale = 1.F / kProbScale; +#pragma unroll + for (int mt = 0; mt < kPvTilesPerWarp; ++mt) + { + // Independent accumulators per MMA keep the tensor pipe busy within one tile. + float part[kPvSteps][2][4] = {}; +#pragma unroll + for (int ks = 0; ks < kPvSteps; ++ks) + { + uint32_t r[4]; + ldmatrixX4Trans(r, vAddr + ks * 32 * kKRowBytes + mt * 16); + uint32_t dimEven[4]; + uint32_t dimOdd[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) + { + uint32_t const pairs = __byte_perm(r[j], 0, 0x3120); + dimEven[j] = e4m3x2ToHalf2(pairs); + dimOdd[j] = e4m3x2ToHalf2(pairs >> 16); + } + mmaF16(part[ks][0], dimEven[0], dimOdd[0], dimEven[1], dimOdd[1], pFrag[ks][0][0], pFrag[ks][0][1]); + mmaF16(part[ks][1], dimEven[2], dimOdd[2], dimEven[3], dimOdd[3], pFrag[ks][1][0], pFrag[ks][1][1]); + } + float o[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + o[i] = ((part[0][0][i] + part[0][1][i]) + (part[1][0][i] + part[1][1][i])) + + ((part[2][0][i] + part[2][1][i]) + (part[3][0][i] + part[3][1][i])); + } + int const dim = warp * kDimsPerWarp + mt * 16 + 2 * groupId; + uint32_t const dest = dim / kDimsPerCtaOut; + int const localDim = dim - static_cast(dest) * kDimsPerCtaOut; + uint32_t const addrA + = recvBase + ((split * kNumHeads + 2 * tig) * kDimsPerCtaOut + localDim) * sizeof(float); + uint32_t const addrB = addrA + kDimsPerCtaOut * sizeof(float); + uint32_t const remoteBar = mapToRank(recvBar, dest); + storeAsync2(mapToRank(addrA, dest), o[0] * invProbScale, o[2] * invProbScale, remoteBar); + storeAsync2(mapToRank(addrB, dest), o[1] * invProbScale, o[3] * invProbScale, remoteBar); + } + } + markTimeline(7); + markTimeline(8); + if (tid == 0) + { + mbarrierArriveExpectTx(recvBar, kRecvBytes); + } + markTimeline(9); + mbarrierWaitCluster(recvBar, 0); + markTimeline(10); + + // Finalize 32 of the 512 output dims of all eight heads from the pushed partials: warp h + // handles head h, lane i < 16 the statistics of source split i, every lane one output dim. + { + int const head = warp; + float const peerMax = lane < kNumSplits ? recvStats[head * kNumSplits + lane] : -INFINITY; + float const peerSum = lane < kNumSplits ? recvStats[(kNumHeads + head) * kNumSplits + lane] : 0.F; + float globalMax = peerMax; +#pragma unroll + for (int offset = kNumSplits / 2; offset > 0; offset >>= 1) + { + globalMax = fmaxf(globalMax, __shfl_xor_sync(0xffffffffU, globalMax, offset)); + } + float const weight = peerMax == -INFINITY ? 0.F : fastExp2(peerMax - globalMax); + float sum = weight * peerSum; +#pragma unroll + for (int offset = kNumSplits / 2; offset > 0; offset >>= 1) + { + sum += __shfl_xor_sync(0xffffffffU, sum, offset); + } + sum = __shfl_sync(0xffffffffU, sum, 0); + float acc = 0.F; +#pragma unroll + for (int i = 0; i < kNumSplits; ++i) + { + acc += __shfl_sync(0xffffffffU, weight, i) * recv[(i * kNumHeads + head) * kDimsPerCtaOut + lane]; + } + float const result = sum > 0.F ? acc * (outScale / sum) : 0.F; + params.out[static_cast(row) * kNumHeads * kHeadDimV + head * kHeadDimV + split * kDimsPerCtaOut + lane] + = __float2bfloat16_rn(result); + } + markTimeline(11); + markGlobalTime(12); +#else + __trap(); +#endif +} + +void prepareKernelAttributes() +{ + TLLM_CUDA_CHECK( + cudaFuncSetAttribute(sparseMlaDecodeKernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemBytes)); + TLLM_CUDA_CHECK(cudaFuncSetAttribute(sparseMlaDecodeKernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1)); +} + +cudaLaunchAttribute clusterAttribute() +{ + cudaLaunchAttribute attr; + attr.id = cudaLaunchAttributeClusterDimension; + attr.val.clusterDim.x = kNumSplits; + attr.val.clusterDim.y = 1; + attr.val.clusterDim.z = 1; + return attr; +} + +} // namespace + +#ifdef GLM_MOE_DSA_SPARSE_MLA_DECODE_TIMING +void copySparseMlaDecodeTimeline(unsigned long long* host) +{ + TLLM_CUDA_CHECK(cudaMemcpyFromSymbol(host, gTimeline, sizeof(gTimeline))); +} +#endif + +size_t getSparseMlaDecodeSmemBytes() +{ + return static_cast(kSmemBytes); +} + +bool isSparseMlaDecodeDeviceSupported(int device) +{ + int major = 0; + int maxSmem = 0; + TLLM_CUDA_CHECK(cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, device)); + TLLM_CUDA_CHECK(cudaDeviceGetAttribute(&maxSmem, cudaDevAttrMaxSharedMemoryPerBlockOptin, device)); + int constexpr kBlackwellDatacenterMajor = 10; + if (major != kBlackwellDatacenterMajor || maxSmem < kSmemBytes) + { + return false; + } + int currentDevice = 0; + TLLM_CUDA_CHECK(cudaGetDevice(¤tDevice)); + TLLM_CUDA_CHECK(cudaSetDevice(device)); + prepareKernelAttributes(); + cudaLaunchConfig_t config{}; + config.gridDim = dim3(kNumSplits); + config.blockDim = dim3(kThreads); + config.dynamicSmemBytes = kSmemBytes; + cudaLaunchAttribute attr = clusterAttribute(); + config.attrs = &attr; + config.numAttrs = 1; + int numClusters = 0; + cudaError_t const status = cudaOccupancyMaxActiveClusters(&numClusters, sparseMlaDecodeKernel, &config); + TLLM_CUDA_CHECK(cudaSetDevice(currentDevice)); + if (status != cudaSuccess) + { + // Clear the error recorded by the failed query. + cudaGetLastError(); + return false; + } + return numClusters > 0; +} + +void invokeSparseMlaDecode(SparseMlaDecodeParams const& params, cudaStream_t stream) +{ + if (params.numRows <= 0) + { + return; + } + prepareKernelAttributes(); + cudaLaunchConfig_t config{}; + config.gridDim = dim3(params.numRows * kNumSplits); + config.blockDim = dim3(kThreads); + config.dynamicSmemBytes = kSmemBytes; + config.stream = stream; + cudaLaunchAttribute attrs[2]; + attrs[0] = clusterAttribute(); + attrs[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[1].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL() ? 1 : 0; + config.attrs = attrs; + config.numAttrs = 2; + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, sparseMlaDecodeKernel, params)); +} + +} // namespace kernels::glm_moe_dsa::sparse_mla_decode + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/glmMoeDsa/sparseMlaDecode.h b/cpp/tensorrt_llm/kernels/glmMoeDsa/sparseMlaDecode.h new file mode 100644 index 000000000000..18dadfa6a556 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/glmMoeDsa/sparseMlaDecode.h @@ -0,0 +1,108 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::glm_moe_dsa::sparse_mla_decode +{ + +// Sparse MLA generation attention over an FP8 (e4m3) latent KV pool. +// +// Each query row attends to the pool rows listed in its `kTopK`-entry index list. A pool row +// holds the 512-dim latent followed by the 64-dim rotary part, 576 e4m3 values in total; the +// value vector is the first 512 values of the same row. Eight query heads share the row. +// +// Rules of the TRT-LLM generation FMHA for this layout, with the one difference marked below: +// * Index-list position j of query row r is used only when j < min(kTopK, kvLen - (seqLenQ - 1 - r % seqLenQ)), +// where kvLen counts every token of the sequence including this step's rows. Repeated pool rows +// count once per occurrence. +// * Difference: a negative entry at a used position is skipped. The native kernel does not check +// it and reads the pool row that -1 addresses. The native producers (indexer top-k, local-to-global +// remap) write -1 only at positions past the causal length, so the outputs agree on their inputs. +// * Scores are FP32 dot products of the e4m3 query and key times bmm1Scale[1] (log2 domain). +// * Probabilities relative to a running maximum are multiplied by 448 and rounded to e4m3 before +// the P*V product; row sums use the unrounded FP32 probabilities. +// * A row without a usable entry produces zeros. +// +// Programmatic dependent launch (PDL), active when TRTLLM_ENABLE_PDL allows it: +// * The kernel is launched with programmatic stream serialization. Before griddepcontrol.wait it +// only initializes its own shared-memory barriers and arrives on its cluster barrier; it accesses +// no global memory there, so it may start while the preceding kernel still writes its inputs. +// * After the wait it reads q, indices, seqLensKv, bmm1Scale and bmm2Scale and issues the +// asynchronous gathers of the pool rows. It then executes griddepcontrol.launch_dependents, before +// those gathers complete and before any compute, so the next kernel of the stream can become +// resident while this one still runs. +// * Its only global write is out, stored at the end of each CTA. +// * Obligation of a PDL dependent: execute griddepcontrol.wait (cudaGridDependencySynchronize) +// before reading out and before writing any buffer read here (q, indices, seqLensKv, the scales, +// the pool rows). Work before that wait may access only memory this kernel neither reads nor +// writes, such as weights streamed into shared memory or L2; such work competes with this +// kernel's gathers for memory bandwidth and SM resources. + +int constexpr kNumHeads = 8; +int constexpr kHeadDimQk = 576; +int constexpr kHeadDimV = 512; +int constexpr kTopK = 2048; + +struct SparseMlaDecodeParams +{ + // [numRows, kNumHeads, kHeadDimQk] e4m3 query, rows contiguous. + void const* q; + // Base of the FP8 latent pool; pool row i starts at kvPool + i * kHeadDimQk bytes. + void const* kvPool; + // [numRows, indicesRowStride] int32 pool rows. + int32_t const* indices; + int64_t indicesRowStride; + // [2] device scales; element 1 is the dequantization scale times the softmax scale times log2(e). + float const* bmm1Scale; + // [1] device scale applied to the normalized P*V output. + float const* bmm2Scale; + // [numRows / seqLenQ] KV lengths including this step's rows. + int32_t const* seqLensKv; + int32_t seqLenQ; + // [numRows, kNumHeads, kHeadDimV] BF16 output. + __nv_bfloat16* out; + int32_t numRows; +}; + +// Dynamic shared memory needed by one CTA. +size_t getSparseMlaDecodeSmemBytes(); + +// Whether the current device can launch the kernel (architecture and 16-CTA clusters). +bool isSparseMlaDecodeDeviceSupported(int device); + +void invokeSparseMlaDecode(SparseMlaDecodeParams const& params, cudaStream_t stream); + +#ifdef GLM_MOE_DSA_SPARSE_MLA_DECODE_TIMING +// Development builds only: per-CTA phase timeline of the most recent launches. +int constexpr kTimelineMaxCtas = 4096; +int constexpr kTimelineSlots = 13; +void copySparseMlaDecodeTimeline(unsigned long long* host); +#endif + +} // namespace kernels::glm_moe_dsa::sparse_mla_decode + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index 43485b14f4ce..e6c701f37c9a 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -173,8 +173,10 @@ add_th_common_family( mhcOp.cpp compressorOp.cpp) +add_th_common_family(th_common_glm_moe_dsa) + # GLM bindings keep their existing family targets; source lists live with the -# model. +# model. The GLM-5.2 GB300 target's torch ops go to th_common_glm_moe_dsa. add_subdirectory(glmMoeDsa) add_library( @@ -185,7 +187,8 @@ add_library( $ $ $ - $) + $ + $) set_property(TARGET th_common PROPERTY POSITION_INDEPENDENT_CODE ON) set_target_properties(th_common PROPERTIES CXX_STANDARD 20 CXX_STANDARD_REQUIRED ON) diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/CMakeLists.txt b/cpp/tensorrt_llm/thop/glmMoeDsa/CMakeLists.txt index cad8b5ad52a0..3ed89d970f96 100644 --- a/cpp/tensorrt_llm/thop/glmMoeDsa/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/glmMoeDsa/CMakeLists.txt @@ -18,3 +18,22 @@ target_sources( glmMnnvlHopOp.cpp glmDraftSampleGlueOp.cpp) target_sources(th_common_decoding PRIVATE glmVerifyHeadAcceptOp.cpp) + +# torch.ops.trtllm.glm_moe_dsa_* ops of the GLM-5.2 GB300 target, one file per +# kernel of cpp/tensorrt_llm/kernels/glmMoeDsa/'s own library. The parent +# defines th_common_glm_moe_dsa. +target_sources( + th_common_glm_moe_dsa + PRIVATE crossRankTopkClusterMergeOp.cpp + crossRankTopkGlobalEmitOp.cpp + crossRankTopkOp.cpp + denseMlpOp.cpp + draftLiveRowHopOp.cpp + draftRoutedMoeOp.cpp + draftTailOp.cpp + indexerLayerRopeKvAppendOp.cpp + postAttentionHopFromPartialsOp.cpp + postAttentionHopOp.cpp + postAttentionHopRoutingOp.cpp + ropeKvAppendOp.cpp + sparseMlaDecodeOp.cpp) diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/crossRankTopkClusterMergeOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsa/crossRankTopkClusterMergeOp.cpp new file mode 100644 index 000000000000..d53886b3d03f --- /dev/null +++ b/cpp/tensorrt_llm/thop/glmMoeDsa/crossRankTopkClusterMergeOp.cpp @@ -0,0 +1,108 @@ +/* + * 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/common/config.h" +#include "tensorrt_llm/kernels/glmMoeDsa/crossRankTopkClusterMerge.h" + +#include +#include + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext::glm_moe_dsa +{ +namespace +{ +namespace tk = tensorrt_llm::kernels::glm_moe_dsa::cross_rank_topk_cluster_merge; + +void crossRankTopkClusterMerge(at::Tensor const& logits, at::Tensor const& kvLens, at::Tensor& out, at::Tensor& state, + int64_t mcWorkspace, int64_t ucWorkspace, int64_t rank, int64_t worldSize, int64_t nextN, + std::optional const& debugTimes) +{ + TORCH_CHECK(logits.is_cuda() && kvLens.is_cuda() && out.is_cuda() && state.is_cuda(), "CUDA tensors required"); + TORCH_CHECK(logits.scalar_type() == at::kFloat && logits.dim() == 2 && logits.stride(1) == 1, + "logits must be a row-major [rows, local_capacity] float32 view"); + TORCH_CHECK(kvLens.scalar_type() == at::kInt && kvLens.numel() >= 1, "kv_lens must be int32 with >= 1 element"); + TORCH_CHECK(out.scalar_type() == at::kInt && out.dim() == 2 && out.is_contiguous() && out.size(1) == tk::kTopK, + "out must be a contiguous int32 [rows, ", tk::kTopK, "] tensor"); + TORCH_CHECK(out.size(0) == logits.size(0), "out and logits must have the same rows"); + TORCH_CHECK(logits.size(0) <= tk::kMaxRows, "at most ", tk::kMaxRows, " rows"); + TORCH_CHECK(state.scalar_type() == at::kInt && state.numel() >= tk::kStateWords && state.is_contiguous(), + "state must be a contiguous int32 tensor of >= ", tk::kStateWords, " words"); + TORCH_CHECK(worldSize >= 1 && worldSize <= tk::kMaxRanks && rank >= 0 && rank < worldSize, "bad rank/world"); + TORCH_CHECK(nextN >= 1 && nextN == logits.size(0), "next_n must equal the number of rows (one request)"); + TORCH_CHECK(logits.size(1) <= tk::cpTopkMaxLocalCapacity(static_cast(worldSize)), "local capacity ", + logits.size(1), " exceeds the kernel limit ", tk::cpTopkMaxLocalCapacity(static_cast(worldSize))); + TORCH_CHECK(mcWorkspace != 0 && ucWorkspace != 0 && (mcWorkspace % 16) == 0 && (ucWorkspace % 16) == 0, + "workspace addresses must be non-null and 16-byte aligned"); + + tk::CpTopkParams params{}; + params.logits = logits.data_ptr(); + params.logitsStride = logits.stride(0); + params.localCapacity = static_cast(logits.size(1)); + params.kvLens = kvLens.data_ptr(); + params.out = out.data_ptr(); + params.numRows = static_cast(logits.size(0)); + params.nextN = static_cast(nextN); + params.rank = static_cast(rank); + params.worldSize = static_cast(worldSize); + params.mcWorkspace = reinterpret_cast(mcWorkspace); + params.ucWorkspace = reinterpret_cast(ucWorkspace); + params.state = state.data_ptr(); + if (debugTimes.has_value()) + { + TORCH_CHECK(debugTimes->is_cuda() && debugTimes->scalar_type() == at::kLong && debugTimes->is_contiguous() + && debugTimes->numel() >= logits.size(0) * tk::kSourcesPerRank * tk::kDebugSlots, + "debug_times must be a contiguous int64 CUDA tensor of rows x ", tk::kSourcesPerRank, " x ", + tk::kDebugSlots); + params.debugTimes = debugTimes->data_ptr(); + } + tk::invokeCpTopk(params, at::cuda::getCurrentCUDAStream(logits.get_device())); +} + +//! [workspace words, top-k, max ranks, max rows, page tokens, sources per rank, first-round slot capacity, max local +//! capacity, state words, Lamport buffers, default timeout (ms), resident clusters of kMaxRows-row launches] +std::vector crossRankTopkClusterMergeWorkspaceInfo(int64_t worldSize) +{ + return {tk::kWorkspaceWords, tk::kTopK, tk::kMaxRanks, tk::kMaxRows, tk::kPageTokens, tk::kSourcesPerRank, + tk::kSourceCapacity, tk::cpTopkMaxLocalCapacity(static_cast(worldSize)), tk::kStateWords, + tk::kNumLamportBuffers, tk::cpTopkDefaultTimeoutMs(), + tk::cpTopkMaxActiveClusters(static_cast(worldSize))}; +} + +} // namespace +} // namespace torch_ext::glm_moe_dsa + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "glm_moe_dsa_cross_rank_topk_cluster_merge(Tensor logits, Tensor kv_lens, Tensor(a!) out, Tensor(b!) state, " + "int mc_workspace, " + "int uc_workspace, int rank, int world_size, int next_n, Tensor(c!)? debug_times=None) -> ()"); + m.def("glm_moe_dsa_cross_rank_topk_cluster_merge_workspace_info(int world_size) -> int[]", + &tensorrt_llm::torch_ext::glm_moe_dsa::crossRankTopkClusterMergeWorkspaceInfo); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl( + "glm_moe_dsa_cross_rank_topk_cluster_merge", &tensorrt_llm::torch_ext::glm_moe_dsa::crossRankTopkClusterMerge); +} diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/crossRankTopkGlobalEmitOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsa/crossRankTopkGlobalEmitOp.cpp new file mode 100644 index 000000000000..0796f5bd96ee --- /dev/null +++ b/cpp/tensorrt_llm/thop/glmMoeDsa/crossRankTopkGlobalEmitOp.cpp @@ -0,0 +1,187 @@ +/* + * 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/common/config.h" +#include "tensorrt_llm/kernels/glmMoeDsa/crossRankTopkGlobalEmit.h" + +#include +#include +#include + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext::glm_moe_dsa +{ +namespace +{ +namespace tk = tensorrt_llm::kernels::glm_moe_dsa::cross_rank_topk_global_emit; + +//! The cluster-merge cross-rank top-k (same arguments as glm_moe_dsa_cross_rank_topk_cluster_merge) that also writes +//! the outputs of the top-k index bridge (trtllm::glm_moe_dsa_topk_index_bridge with the same bridge arguments, applied +//! to `out`): returns the pool rows, [G, rows, top-k] with layer_ids (G = its length) or [rows, top-k] without, and +//! writes the copy rows into copy_dst. +at::Tensor crossRankTopkGlobalEmit(at::Tensor const& logits, at::Tensor const& kvLens, at::Tensor& out, + at::Tensor& state, int64_t mcWorkspace, int64_t ucWorkspace, int64_t rank, int64_t worldSize, int64_t nextN, + at::Tensor const& reqIdx, int64_t reqIdxBias, at::Tensor const& blockTable, int64_t blockSize, int64_t strideFactor, + std::optional const& layerIds, int64_t layerId, std::optional const& copyDst, + int64_t copyRowStart, int64_t copyRows, int64_t copyMode, int64_t copyNextN, + std::optional const& numAccepted, std::optional const& debugTimes) +{ + TORCH_CHECK(logits.is_cuda() && kvLens.is_cuda() && out.is_cuda() && state.is_cuda(), "CUDA tensors required"); + TORCH_CHECK(logits.scalar_type() == at::kFloat && logits.dim() == 2 && logits.stride(1) == 1, + "logits must be a row-major [rows, local_capacity] float32 view"); + TORCH_CHECK(kvLens.scalar_type() == at::kInt && kvLens.numel() >= 1, "kv_lens must be int32 with >= 1 element"); + TORCH_CHECK(out.scalar_type() == at::kInt && out.dim() == 2 && out.is_contiguous() && out.size(1) == tk::kTopK, + "out must be a contiguous int32 [rows, ", tk::kTopK, "] tensor"); + TORCH_CHECK(out.size(0) == logits.size(0), "out and logits must have the same rows"); + TORCH_CHECK(logits.size(0) <= tk::kMaxRows, "at most ", tk::kMaxRows, " rows"); + TORCH_CHECK(state.scalar_type() == at::kInt && state.numel() >= tk::kStateWords && state.is_contiguous(), + "state must be a contiguous int32 tensor of >= ", tk::kStateWords, " words"); + TORCH_CHECK(worldSize >= 1 && worldSize <= tk::kMaxRanks && rank >= 0 && rank < worldSize, "bad rank/world"); + TORCH_CHECK(nextN >= 1 && nextN == logits.size(0), "next_n must equal the number of rows (one request)"); + TORCH_CHECK(logits.size(1) <= tk::cpTopkMaxLocalCapacity(static_cast(worldSize)), "local capacity ", + logits.size(1), " exceeds the kernel limit ", tk::cpTopkMaxLocalCapacity(static_cast(worldSize))); + TORCH_CHECK(mcWorkspace != 0 && ucWorkspace != 0 && (mcWorkspace % 16) == 0 && (ucWorkspace % 16) == 0, + "workspace addresses must be non-null and 16-byte aligned"); + + // The bridge inputs, with the rules of trtllm::glm_moe_dsa_topk_index_bridge (topk = out, num_rows = rows). + int64_t const rows = out.size(0); + TORCH_CHECK(reqIdx.is_cuda() && blockTable.is_cuda(), "req_idx and block_table must be CUDA tensors"); + TORCH_CHECK(reqIdx.scalar_type() == at::kInt && reqIdx.dim() == 1 && reqIdx.is_contiguous(), + "req_idx must be packed int32"); + TORCH_CHECK(reqIdx.size(0) >= rows, "req_idx shorter than the rows"); + TORCH_CHECK(blockTable.scalar_type() == at::kInt && blockTable.dim() == 2, "block_table must be int32 2-D"); + TORCH_CHECK(blockSize == tk::kPageTokens, "block_size must be ", tk::kPageTokens, " (the CP page)"); + TORCH_CHECK(copyMode == 0 || copyMode == 1, "copy_mode must be 0 or 1"); + TORCH_CHECK(copyNextN >= 1 && copyRows >= 0 && copyRowStart >= 0, "bad copy rows"); + TORCH_CHECK(copyMode == 1 || copyRows <= rows, "copy_rows exceeds the rows"); + TORCH_CHECK(copyMode == 0 || copyRows * copyNextN <= rows, "copy_rows * copy_next_n exceeds the rows"); + + int32_t groupSize = 1; + int32_t const* layerIdsPtr = nullptr; + if (layerIds.has_value()) + { + TORCH_CHECK(layerIds->is_cuda() && layerIds->scalar_type() == at::kInt && layerIds->dim() == 1 + && layerIds->is_contiguous(), + "layer_ids must be packed int32 on the GPU"); + groupSize = static_cast(layerIds->size(0)); + TORCH_CHECK(groupSize >= 1 && groupSize <= tk::kMaxGroupSize, + "glm_moe_dsa_cross_rank_topk_global_emit: remap group sizes must be 1..", tk::kMaxGroupSize, ", got ", + groupSize); + layerIdsPtr = layerIds->data_ptr(); + } + int64_t const topK = tk::kTopK; + at::Tensor pool = layerIds.has_value() ? at::empty({static_cast(groupSize), rows, topK}, out.options()) + : at::empty({rows, topK}, out.options()); + + tk::CpTopkParams params{}; + params.logits = logits.data_ptr(); + params.logitsStride = logits.stride(0); + params.localCapacity = static_cast(logits.size(1)); + params.kvLens = kvLens.data_ptr(); + params.out = out.data_ptr(); + params.numRows = static_cast(rows); + params.nextN = static_cast(nextN); + params.rank = static_cast(rank); + params.worldSize = static_cast(worldSize); + params.mcWorkspace = reinterpret_cast(mcWorkspace); + params.ucWorkspace = reinterpret_cast(ucWorkspace); + params.state = state.data_ptr(); + + tk::CpTopkBridgeParams& b = params.bridge; + b.out = pool.data_ptr(); + b.outStrideG = layerIds.has_value() ? pool.stride(0) : 0; + b.outStride0 = layerIds.has_value() ? pool.stride(1) : pool.stride(0); + b.groupSize = groupSize; + b.layerIds = layerIdsPtr; + b.layerId = static_cast(layerId); + b.reqIdx = reqIdx.data_ptr(); + b.reqIdxBias = static_cast(reqIdxBias); + b.blockTable = blockTable.data_ptr(); + b.btStride0 = blockTable.stride(0); + b.btStride1 = blockTable.stride(1); + b.maxBlocksPerReq = static_cast(blockTable.size(1)); + b.strideFactor = static_cast(strideFactor); + if (copyDst.has_value() && copyRows > 0) + { + auto const& dst = copyDst.value(); + TORCH_CHECK(dst.is_cuda() && dst.scalar_type() == at::kInt && dst.dim() == 2 && dst.stride(1) == 1 + && dst.size(1) >= tk::kTopK && dst.size(0) >= copyRowStart + copyRows, + "copy_dst must be int32 [>= start + rows, >= top-k] with unit column stride"); + b.copyDst = dst.data_ptr(); + b.copyStride0 = dst.stride(0); + b.copyRows = static_cast(copyRows); + b.copyRowStart = static_cast(copyRowStart); + b.copyMode = static_cast(copyMode); + b.copyNextN = static_cast(copyNextN); + if (copyMode == 1 && numAccepted.has_value()) + { + auto const& acc = numAccepted.value(); + TORCH_CHECK(acc.is_cuda() && (acc.scalar_type() == at::kInt || acc.scalar_type() == at::kLong) + && acc.dim() == 1 && acc.is_contiguous() && acc.size(0) >= copyRows, + "num_accepted must be packed int32/int64 with at least copy_rows entries"); + b.numAccepted = acc.data_ptr(); + b.numAcceptedIs64 = acc.scalar_type() == at::kLong; + } + } + if (debugTimes.has_value()) + { + TORCH_CHECK(debugTimes->is_cuda() && debugTimes->scalar_type() == at::kLong && debugTimes->is_contiguous() + && debugTimes->numel() >= logits.size(0) * tk::kSourcesPerRank * tk::kDebugSlots, + "debug_times must be a contiguous int64 CUDA tensor of rows x ", tk::kSourcesPerRank, " x ", + tk::kDebugSlots); + params.debugTimes = debugTimes->data_ptr(); + } + tk::invokeCpTopk(params, at::cuda::getCurrentCUDAStream(logits.get_device())); + return pool; +} + +//! [workspace words, top-k, max ranks, max rows, page tokens, sources per rank, first-round slot capacity, max local +//! capacity, state words, Lamport buffers, default timeout (ms), resident clusters of kMaxRows-row launches, max remap +//! group size] +std::vector crossRankTopkGlobalEmitWorkspaceInfo(int64_t worldSize) +{ + return {tk::kWorkspaceWords, tk::kTopK, tk::kMaxRanks, tk::kMaxRows, tk::kPageTokens, tk::kSourcesPerRank, + tk::kSourceCapacity, tk::cpTopkMaxLocalCapacity(static_cast(worldSize)), tk::kStateWords, + tk::kNumLamportBuffers, tk::cpTopkDefaultTimeoutMs(), + tk::cpTopkMaxActiveClusters(static_cast(worldSize)), tk::kMaxGroupSize}; +} + +} // namespace +} // namespace torch_ext::glm_moe_dsa + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "glm_moe_dsa_cross_rank_topk_global_emit(Tensor logits, Tensor kv_lens, Tensor(a!) out, Tensor(b!) state, int " + "mc_workspace, " + "int uc_workspace, int rank, int world_size, int next_n, Tensor req_idx, int req_idx_bias, " + "Tensor block_table, int block_size, int stride_factor, Tensor? layer_ids, int layer_id, " + "Tensor(c!)? copy_dst, int copy_row_start, int copy_rows, int copy_mode, int copy_next_n, " + "Tensor? num_accepted, Tensor(d!)? debug_times=None) -> Tensor"); + m.def("glm_moe_dsa_cross_rank_topk_global_emit_workspace_info(int world_size) -> int[]", + &tensorrt_llm::torch_ext::glm_moe_dsa::crossRankTopkGlobalEmitWorkspaceInfo); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("glm_moe_dsa_cross_rank_topk_global_emit", &tensorrt_llm::torch_ext::glm_moe_dsa::crossRankTopkGlobalEmit); +} diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/crossRankTopkOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsa/crossRankTopkOp.cpp new file mode 100644 index 000000000000..710b0e716b88 --- /dev/null +++ b/cpp/tensorrt_llm/thop/glmMoeDsa/crossRankTopkOp.cpp @@ -0,0 +1,104 @@ +/* + * 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/common/config.h" +#include "tensorrt_llm/kernels/glmMoeDsa/crossRankTopk.h" + +#include +#include + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext::glm_moe_dsa +{ +namespace +{ +namespace tk = tensorrt_llm::kernels::glm_moe_dsa::cross_rank_topk; + +void crossRankTopk(at::Tensor const& logits, at::Tensor const& kvLens, at::Tensor& out, at::Tensor& state, + int64_t mcWorkspace, int64_t ucWorkspace, int64_t rank, int64_t worldSize, int64_t nextN, + std::optional const& debugTimes) +{ + TORCH_CHECK(logits.is_cuda() && kvLens.is_cuda() && out.is_cuda() && state.is_cuda(), "CUDA tensors required"); + TORCH_CHECK(logits.scalar_type() == at::kFloat && logits.dim() == 2 && logits.stride(1) == 1, + "logits must be a row-major [rows, local_capacity] float32 view"); + TORCH_CHECK(kvLens.scalar_type() == at::kInt && kvLens.numel() >= 1, "kv_lens must be int32 with >= 1 element"); + TORCH_CHECK(out.scalar_type() == at::kInt && out.dim() == 2 && out.is_contiguous() && out.size(1) == tk::kTopK, + "out must be a contiguous int32 [rows, ", tk::kTopK, "] tensor"); + TORCH_CHECK(out.size(0) == logits.size(0), "out and logits must have the same rows"); + TORCH_CHECK(logits.size(0) <= tk::kMaxRows, "at most ", tk::kMaxRows, " rows"); + TORCH_CHECK(state.scalar_type() == at::kInt && state.numel() >= tk::kStateWords && state.is_contiguous(), + "state must be a contiguous int32 tensor of >= ", tk::kStateWords, " words"); + TORCH_CHECK(worldSize >= 1 && worldSize <= tk::kMaxRanks && rank >= 0 && rank < worldSize, "bad rank/world"); + TORCH_CHECK(nextN >= 1 && nextN == logits.size(0), "next_n must equal the number of rows (one request)"); + TORCH_CHECK(logits.size(1) <= tk::cpTopkMaxLocalCapacity(static_cast(worldSize)), "local capacity ", + logits.size(1), " exceeds the kernel limit ", tk::cpTopkMaxLocalCapacity(static_cast(worldSize))); + TORCH_CHECK(mcWorkspace != 0 && ucWorkspace != 0 && (mcWorkspace % 16) == 0 && (ucWorkspace % 16) == 0, + "workspace addresses must be non-null and 16-byte aligned"); + + tk::CpTopkParams params{}; + params.logits = logits.data_ptr(); + params.logitsStride = logits.stride(0); + params.localCapacity = static_cast(logits.size(1)); + params.kvLens = kvLens.data_ptr(); + params.out = out.data_ptr(); + params.numRows = static_cast(logits.size(0)); + params.nextN = static_cast(nextN); + params.rank = static_cast(rank); + params.worldSize = static_cast(worldSize); + params.mcWorkspace = reinterpret_cast(mcWorkspace); + params.ucWorkspace = reinterpret_cast(ucWorkspace); + params.state = state.data_ptr(); + if (debugTimes.has_value()) + { + TORCH_CHECK(debugTimes->is_cuda() && debugTimes->scalar_type() == at::kLong && debugTimes->is_contiguous() + && debugTimes->numel() >= logits.size(0) * tk::kSourcesPerRank * tk::kDebugSlots, + "debug_times must be a contiguous int64 CUDA tensor of rows x ", tk::kSourcesPerRank, " x ", + tk::kDebugSlots); + params.debugTimes = debugTimes->data_ptr(); + } + tk::invokeCpTopk(params, at::cuda::getCurrentCUDAStream(logits.get_device())); +} + +std::vector crossRankTopkWorkspaceInfo(int64_t worldSize) +{ + return {tk::kWorkspaceWords, tk::kTopK, tk::kMaxRanks, tk::kMaxRows, tk::kPageTokens, tk::kSourcesPerRank, + tk::kSourceCapacity, tk::cpTopkMaxLocalCapacity(static_cast(worldSize)), tk::kStateWords, + tk::kNumLamportBuffers, tk::cpTopkDefaultTimeoutMs()}; +} + +} // namespace +} // namespace torch_ext::glm_moe_dsa + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "glm_moe_dsa_cross_rank_topk(Tensor logits, Tensor kv_lens, Tensor(a!) out, Tensor(b!) state, int " + "mc_workspace, " + "int uc_workspace, int rank, int world_size, int next_n, Tensor(c!)? debug_times=None) -> ()"); + m.def("glm_moe_dsa_cross_rank_topk_workspace_info(int world_size) -> int[]", + &tensorrt_llm::torch_ext::glm_moe_dsa::crossRankTopkWorkspaceInfo); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("glm_moe_dsa_cross_rank_topk", &tensorrt_llm::torch_ext::glm_moe_dsa::crossRankTopk); +} diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/denseMlpOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsa/denseMlpOp.cpp new file mode 100644 index 000000000000..c4843dcb6e44 --- /dev/null +++ b/cpp/tensorrt_llm/thop/glmMoeDsa/denseMlpOp.cpp @@ -0,0 +1,130 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/envUtils.h" +#include "tensorrt_llm/kernels/glmMoeDsa/denseMlp.h" + +#include +#include +#include + +#include + +namespace tk = tensorrt_llm::kernels::glm_moe_dsa::dense_mlp; + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext::glm_moe_dsa +{ + +namespace +{ + +constexpr int64_t kVecAlignBytes = 16; + +void checkBf16Matrix(at::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda(), name, " must be a CUDA tensor"); + TORCH_CHECK(t.scalar_type() == at::kBFloat16, name, " must be BF16"); + TORCH_CHECK(t.dim() == 2, name, " must be 2-D"); + TORCH_CHECK(t.is_contiguous(), name, " must be contiguous"); + TORCH_CHECK(reinterpret_cast(t.data_ptr()) % kVecAlignBytes == 0, name, " must be 16-byte aligned"); +} + +void launchGateUp(at::Tensor const& x, at::Tensor const& gateUpWeight, at::Tensor& h, int64_t variant, bool enablePdl) +{ + auto const stream = at::cuda::getCurrentCUDAStream(x.get_device()); + TLLM_CUDA_CHECK(tk::invokeGateUpSwiglu(reinterpret_cast<__nv_bfloat16 const*>(x.data_ptr()), + reinterpret_cast<__nv_bfloat16 const*>(gateUpWeight.data_ptr()), reinterpret_cast<__nv_bfloat16*>(h.data_ptr()), + static_cast(x.size(0)), static_cast(x.size(1)), static_cast(h.size(1)), + static_cast(variant), enablePdl, stream)); +} + +void launchDownProj(at::Tensor const& h, at::Tensor const& downWeight, at::Tensor& out, int64_t variant, bool enablePdl) +{ + auto const stream = at::cuda::getCurrentCUDAStream(h.get_device()); + TLLM_CUDA_CHECK(tk::invokeDown(reinterpret_cast<__nv_bfloat16 const*>(h.data_ptr()), + reinterpret_cast<__nv_bfloat16 const*>(downWeight.data_ptr()), reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), + static_cast(h.size(0)), static_cast(h.size(1)), static_cast(out.size(1)), + static_cast(variant), enablePdl, stream)); +} + +} // namespace + +//! SwiGLU of the fused gate/up projection: x [M, H], gate_up_weight [2I, H] -> h [M, I]. +at::Tensor denseMlpGateUpSwiglu(at::Tensor const& x, at::Tensor const& gateUpWeight, int64_t variant) +{ + checkBf16Matrix(x, "x"); + checkBf16Matrix(gateUpWeight, "gate_up_weight"); + int64_t const numRows = x.size(0); + int64_t const hidden = x.size(1); + TORCH_CHECK(gateUpWeight.size(1) == hidden, "gate_up_weight must be [2 * inter, hidden]"); + TORCH_CHECK(gateUpWeight.size(0) % 2 == 0, "gate_up_weight rows must be even"); + int64_t const inter = gateUpWeight.size(0) / 2; + TORCH_CHECK(tk::isGateUpSupported(static_cast(numRows), static_cast(hidden), static_cast(inter), + static_cast(variant)), + "glm_moe_dsa_dense_mlp_gate_up_swiglu: unsupported shape or variant"); + at::Tensor h = at::empty({numRows, inter}, x.options()); + launchGateUp(x, gateUpWeight, h, variant, tensorrt_llm::common::getEnvEnablePDL()); + return h; +} + +//! Down projection: h [M, I], down_weight [H, I] -> out [M, H] (rank-local partial, no reduction). +at::Tensor denseMlpDown(at::Tensor const& h, at::Tensor const& downWeight, int64_t variant) +{ + checkBf16Matrix(h, "h"); + checkBf16Matrix(downWeight, "down_weight"); + int64_t const numRows = h.size(0); + int64_t const inter = h.size(1); + TORCH_CHECK(downWeight.size(1) == inter, "down_weight must be [hidden, inter]"); + int64_t const hidden = downWeight.size(0); + TORCH_CHECK(tk::isDownSupported(static_cast(numRows), static_cast(inter), static_cast(hidden), + static_cast(variant)), + "glm_moe_dsa_dense_mlp_down: unsupported shape or variant"); + at::Tensor out = at::empty({numRows, hidden}, h.options()); + launchDownProj(h, downWeight, out, variant, tensorrt_llm::common::getEnvEnablePDL()); + return out; +} + +//! Dense gated MLP without the output all-reduce: x [M, H] -> rank-local partial [M, H]. +at::Tensor denseMlp(at::Tensor const& x, at::Tensor const& gateUpWeight, at::Tensor const& downWeight, + int64_t gateUpVariant, int64_t downVariant) +{ + at::Tensor const h = denseMlpGateUpSwiglu(x, gateUpWeight, gateUpVariant); + return denseMlpDown(h, downWeight, downVariant); +} + +} // namespace torch_ext::glm_moe_dsa + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("glm_moe_dsa_dense_mlp_gate_up_swiglu(Tensor x, Tensor gate_up_weight, int variant=1) -> Tensor"); + m.def("glm_moe_dsa_dense_mlp_down(Tensor h, Tensor down_weight, int variant=0) -> Tensor"); + m.def( + "glm_moe_dsa_dense_mlp(Tensor x, Tensor gate_up_weight, Tensor down_weight, int gate_up_variant=1, " + "int down_variant=0) -> Tensor"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("glm_moe_dsa_dense_mlp_gate_up_swiglu", &tensorrt_llm::torch_ext::glm_moe_dsa::denseMlpGateUpSwiglu); + m.impl("glm_moe_dsa_dense_mlp_down", &tensorrt_llm::torch_ext::glm_moe_dsa::denseMlpDown); + m.impl("glm_moe_dsa_dense_mlp", &tensorrt_llm::torch_ext::glm_moe_dsa::denseMlp); +} diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/draftLiveRowHopOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsa/draftLiveRowHopOp.cpp new file mode 100644 index 000000000000..b439f5d003fd --- /dev/null +++ b/cpp/tensorrt_llm/thop/glmMoeDsa/draftLiveRowHopOp.cpp @@ -0,0 +1,150 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/config.h" +#include "tensorrt_llm/kernels/glmMoeDsa/draftLiveRowHop.h" + +#include +#include +#include +#include + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext::glm_moe_dsa +{ + +namespace +{ + +namespace tk = tensorrt_llm::kernels::glm_moe_dsa::draft_live_rows; + +// Same switch as the native kernels: PDL is on unless TRTLLM_ENABLE_PDL=0. +bool pdlEnabled() +{ + static bool const enabled = []() + { + char const* env = std::getenv("TRTLLM_ENABLE_PDL"); + return !(env != nullptr && std::strcmp(env, "0") == 0); + }(); + return enabled; +} + +void checkRows(at::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda() && t.scalar_type() == at::kBFloat16, "[glm_moe_dsa_draft_live_row_hop] ", name, + " must be a CUDA bf16 tensor"); + TORCH_CHECK(t.is_contiguous(), "[glm_moe_dsa_draft_live_row_hop] ", name, " must be contiguous"); + TORCH_CHECK(t.dim() >= 1 && t.size(-1) == tk::kHiddenDim, "[glm_moe_dsa_draft_live_row_hop] ", name, + " must have hidden size ", tk::kHiddenDim); +} + +} // namespace + +//! Live-row post-attention hop: for the rows ``row_ids`` of ``input`` / ``residual`` only, the one-shot cross-rank +//! sum + residual + RMSNorm + router logits of the post-attention hop, on that hop's workspace. +//! Returns ``[normed, residual_out, router_logits]`` with ``row_ids.numel()`` rows (compact, in ``row_ids`` order). +//! See kernels/glmMoeDsa/draftLiveRowHop.h for the arithmetic and the workspace protocol. +std::vector draftLiveRowHop(at::Tensor const& input, at::Tensor const& residual, at::Tensor const& rowIds, + at::Tensor const& gamma, double eps, at::Tensor const& routerWeight, at::Tensor const& ucBuffer, int64_t mcPtr, + at::Tensor const& flags, int64_t rank, int64_t nRanks) +{ + checkRows(input, "input"); + checkRows(residual, "residual"); + TORCH_CHECK( + residual.sizes() == input.sizes(), "[glm_moe_dsa_draft_live_row_hop] residual must have the shape of input"); + checkRows(gamma, "gamma"); + TORCH_CHECK(gamma.dim() == 1, "[glm_moe_dsa_draft_live_row_hop] gamma must be 1-D"); + int64_t const numInputRows = input.numel() / tk::kHiddenDim; + TORCH_CHECK(numInputRows >= 1 && numInputRows <= tk::kMaxRows, + "[glm_moe_dsa_draft_live_row_hop] number of input rows ", numInputRows, " is outside the supported range [1, ", + tk::kMaxRows, "]"); + TORCH_CHECK(rowIds.is_cuda() && rowIds.scalar_type() == at::kLong && rowIds.dim() == 1 && rowIds.is_contiguous(), + "[glm_moe_dsa_draft_live_row_hop] row_ids must be a contiguous 1-D CUDA int64 tensor"); + int64_t const numRows = rowIds.numel(); + TORCH_CHECK(numRows >= 1 && numRows <= numInputRows, "[glm_moe_dsa_draft_live_row_hop] number of live rows ", + numRows, " is outside [1, ", numInputRows, "]"); + TORCH_CHECK(nRanks == tk::kNumRanks, "[glm_moe_dsa_draft_live_row_hop] only TP", tk::kNumRanks, + " is supported, got ", nRanks); + TORCH_CHECK(rank >= 0 && rank < nRanks, "[glm_moe_dsa_draft_live_row_hop] invalid rank ", rank); + TORCH_CHECK(ucBuffer.is_cuda() && ucBuffer.is_contiguous() + && static_cast(ucBuffer.numel() * ucBuffer.element_size()) >= tk::kWorkspaceBytes, + "[glm_moe_dsa_draft_live_row_hop] uc_buffer must be a contiguous CUDA tensor of at least ", tk::kWorkspaceBytes, + " bytes"); + TORCH_CHECK( + mcPtr != 0, "[glm_moe_dsa_draft_live_row_hop] mc_ptr must be the multicast address of uc_buffer's workspace"); + TORCH_CHECK(flags.is_cuda() && flags.is_contiguous() && flags.scalar_type() == at::kInt && flags.numel() >= 4, + "[glm_moe_dsa_draft_live_row_hop] flags must be a contiguous CUDA int32 tensor with at least 4 elements"); + TORCH_CHECK(routerWeight.is_cuda() && routerWeight.scalar_type() == at::kBFloat16 && routerWeight.is_contiguous() + && routerWeight.dim() == 2 && routerWeight.size(0) == tk::kNumExperts + && routerWeight.size(1) == tk::kHiddenDim, + "[glm_moe_dsa_draft_live_row_hop] router_weight must be a contiguous CUDA bf16 tensor of shape [", + tk::kNumExperts, ", ", tk::kHiddenDim, "]"); + TORCH_CHECK(input.get_device() == ucBuffer.get_device() && input.get_device() == flags.get_device() + && input.get_device() == rowIds.get_device(), + "[glm_moe_dsa_draft_live_row_hop] all tensors must be on the same device"); + + c10::cuda::CUDAGuard const deviceGuard(input.device()); + auto const rowsOpts = input.options(); + at::Tensor normOut = at::empty({numRows, tk::kHiddenDim}, rowsOpts); + at::Tensor residualOut = at::empty({numRows, tk::kHiddenDim}, rowsOpts); + at::Tensor routerLogits = at::empty({numRows, tk::kNumExperts}, rowsOpts.dtype(at::kFloat)); + + tk::LiveRowHopParams params; + params.input = input.const_data_ptr(); + params.residualIn = residual.const_data_ptr(); + params.rowIds = rowIds.const_data_ptr(); + params.gamma = gamma.const_data_ptr(); + params.routerWeight = routerWeight.const_data_ptr(); + params.normOut = normOut.mutable_data_ptr(); + params.residualOut = residualOut.mutable_data_ptr(); + params.routerLogits = routerLogits.mutable_data_ptr(); + params.ucBuffer = ucBuffer.mutable_data_ptr(); + params.mcBuffer = reinterpret_cast(mcPtr); + params.flags = reinterpret_cast(flags.mutable_data_ptr()); + params.numRows = static_cast(numRows); + params.numInputRows = static_cast(numInputRows); + params.rank = static_cast(rank); + params.eps = static_cast(eps); + params.enablePdl = pdlEnabled(); + params.stream = at::cuda::getCurrentCUDAStream(input.get_device()); + + cudaError_t const status = tk::invokeDraftLiveRowHop(params); + TORCH_CHECK( + status == cudaSuccess, "[glm_moe_dsa_draft_live_row_hop] kernel launch failed: ", cudaGetErrorString(status)); + return {normOut, residualOut, routerLogits}; +} + +} // namespace torch_ext::glm_moe_dsa + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "glm_moe_dsa_draft_live_row_hop(Tensor input, Tensor residual, Tensor row_ids, Tensor gamma, float eps, " + "Tensor router_weight, Tensor uc_buffer, int mc_ptr, Tensor flags, int rank, int nranks) -> Tensor[]"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("glm_moe_dsa_draft_live_row_hop", &tensorrt_llm::torch_ext::glm_moe_dsa::draftLiveRowHop); +} diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/draftRoutedMoeOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsa/draftRoutedMoeOp.cpp new file mode 100644 index 000000000000..15d02bced53a --- /dev/null +++ b/cpp/tensorrt_llm/thop/glmMoeDsa/draftRoutedMoeOp.cpp @@ -0,0 +1,246 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/config.h" +#include "tensorrt_llm/kernels/glmMoeDsa/draftRoutedMoe.h" + +#include +#include +#include +#include + +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext::glm_moe_dsa +{ + +namespace +{ + +namespace tk = tensorrt_llm::kernels::glm_moe_dsa::draft_routed_moe; + +#ifndef GLM_MOE_DSA_DRAFT_ROUTED_MOE_TIMELINE +#define GLM_MOE_DSA_DRAFT_ROUTED_MOE_TIMELINE 0 +#endif +//! Probe builds compile the kernels with their event record (see draftRoutedMoe.cu). +constexpr bool kTimelineBuild = GLM_MOE_DSA_DRAFT_ROUTED_MOE_TIMELINE != 0; + +void checkTensor(at::Tensor const& t, char const* name, at::ScalarType dtype, std::initializer_list shape, + c10::DeviceIndex device) +{ + TORCH_CHECK(t.is_cuda() && t.get_device() == device, "[glm_moe_dsa_draft_routed_moe] ", name, + " must be on the activation's device"); + TORCH_CHECK(t.scalar_type() == dtype, "[glm_moe_dsa_draft_routed_moe] ", name, " has dtype ", t.scalar_type()); + TORCH_CHECK(t.is_contiguous(), "[glm_moe_dsa_draft_routed_moe] ", name, " must be contiguous"); + TORCH_CHECK(t.sizes().vec() == std::vector(shape), "[glm_moe_dsa_draft_routed_moe] ", name, " has shape ", + t.sizes()); +} + +//! ``cuTensorMapEncodeTiled`` through the runtime's driver entry point (no link-time driver dependency). +PFN_cuTensorMapEncodeTiled_v12000 encodeTiledFn() +{ + static PFN_cuTensorMapEncodeTiled_v12000 fn = [] + { + void* ptr = nullptr; + cudaDriverEntryPointQueryResult query{}; + cudaError_t const status + = cudaGetDriverEntryPointByVersion("cuTensorMapEncodeTiled", &ptr, 12000, cudaEnableDefault, &query); + TORCH_CHECK(status == cudaSuccess && query == cudaDriverEntryPointSuccess && ptr != nullptr, + "[glm_moe_dsa_draft_routed_moe] cuTensorMapEncodeTiled is not available"); + return reinterpret_cast(ptr); + }(); + return fn; +} + +//! Tensor map of a BlockMajorK weight buffer viewed as [slabs][slabRows][64] BF16 (a slab: one K chunk of one +//! expert) with boxes of 64 x ``boxRows`` x ``boxSlabs`` and the 128-byte swizzle. Cached per (device, address, +//! shape, box): the expert buffers are fixed model parameters. +CUtensorMap const& weightMap( + int device, void const* base, uint64_t slabRows, uint64_t slabs, uint32_t boxRows, uint32_t boxSlabs) +{ + static std::mutex mutex; + static std::map, CUtensorMap> cache; + std::lock_guard lock(mutex); + auto const key = std::make_tuple(device, base, slabRows, slabs, boxRows, boxSlabs); + auto it = cache.find(key); + if (it != cache.end()) + { + return it->second; + } + CUtensorMap map{}; + cuuint64_t const rowBytes = static_cast(tk::kKChunk) * 2; + cuuint64_t const dims[3] = {static_cast(tk::kKChunk), slabRows, slabs}; + cuuint64_t const strides[2] = {rowBytes, rowBytes * slabRows}; + cuuint32_t const box[3] = {static_cast(tk::kKChunk), boxRows, boxSlabs}; + cuuint32_t const elementStrides[3] = {1, 1, 1}; + CUresult const result = encodeTiledFn()(&map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, const_cast(base), dims, + strides, box, elementStrides, CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, + CU_TENSOR_MAP_L2_PROMOTION_L2_256B, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TORCH_CHECK(result == CUDA_SUCCESS, + "[glm_moe_dsa_draft_routed_moe] cuTensorMapEncodeTiled failed: ", static_cast(result)); + return cache.emplace(key, map).first->second; +} + +} // namespace + +//! Routed experts of the MTP draft layer (see kernels/glmMoeDsa/draftRoutedMoe.h): native DeepSeek-V3 +//! routing of ``router_logits`` with ``routing_bias``, FC1 + SwiGLU + FC2 + finalize on the FlashInfer TRT-LLM-Gen +//! BF16 layouts of ``w3_w1`` / ``w2``. Returns the finalized routed output ``[R, 6144]`` (BF16). ``h_scratch`` and +//! ``route_scratch`` are written (declared mutated in the schema) and reused by every call; calls sharing them must +//! be ordered on one stream. ``debug_fc2`` (optional, ``[R, 8, 6144]`` BF16) +//! receives the FC2 row of every (token, slot). ``timeline`` (optional, int64 ``[2, kTimelineCtas, kTimelineSlots]``) +//! receives the per-CTA event record of the TMA-fed kernels; only probe builds +//! (``GLM_MOE_DSA_DRAFT_ROUTED_MOE_TIMELINE=1``) accept it. +at::Tensor draftRoutedMoe(at::Tensor const& x, at::Tensor const& routerLogits, at::Tensor const& routingBias, + at::Tensor const& w3w1, at::Tensor const& w2, at::Tensor const& hScratch, at::Tensor const& routeScratch, + double routedScaling, int64_t fc1Ctas, int64_t fc1Warps, int64_t fc2Ctas, int64_t fc2Warps, int64_t flags, + std::optional const& debugFc2, std::optional const& timeline) +{ + TORCH_CHECK(x.is_cuda() && x.scalar_type() == at::kBFloat16 && x.is_contiguous() && x.dim() == 2 + && x.size(1) == tk::kHidden, + "[glm_moe_dsa_draft_routed_moe] x must be a contiguous CUDA bf16 tensor [R, ", tk::kHidden, "]"); + int64_t const rows = x.size(0); + TORCH_CHECK(rows >= 1 && rows <= tk::kMaxRows, "[glm_moe_dsa_draft_routed_moe] rows ", rows, " outside [1, ", + tk::kMaxRows, "]"); + TORCH_CHECK(reinterpret_cast(x.const_data_ptr()) % 16 == 0, + "[glm_moe_dsa_draft_routed_moe] x must be 16-byte aligned"); + auto const device = x.get_device(); + checkTensor(routerLogits, "router_logits", at::kFloat, {rows, tk::kNumExperts}, device); + checkTensor(routingBias, "routing_bias", at::kBFloat16, {tk::kNumExperts}, device); + TORCH_CHECK(reinterpret_cast(routerLogits.const_data_ptr()) % 16 == 0 + && reinterpret_cast(routingBias.const_data_ptr()) % 16 == 0, + "[glm_moe_dsa_draft_routed_moe] router_logits and routing_bias must be 16-byte aligned"); + checkTensor(w3w1, "w3_w1", at::kBFloat16, {tk::kNumExperts, tk::kFc1Chunks, tk::kFc1Rows, tk::kKChunk}, device); + checkTensor(w2, "w2", at::kBFloat16, {tk::kNumExperts, tk::kFc2Chunks, tk::kHidden, tk::kKChunk}, device); + TORCH_CHECK(hScratch.is_cuda() && hScratch.get_device() == device && hScratch.scalar_type() == at::kBFloat16 + && hScratch.is_contiguous() && hScratch.numel() >= tk::kScratchHElems, + "[glm_moe_dsa_draft_routed_moe] h_scratch must be a contiguous CUDA bf16 tensor with at least ", + tk::kScratchHElems, " elements"); + TORCH_CHECK(routeScratch.is_cuda() && routeScratch.get_device() == device && routeScratch.scalar_type() == at::kInt + && routeScratch.is_contiguous() && routeScratch.numel() >= tk::kRouteWords, + "[glm_moe_dsa_draft_routed_moe] route_scratch must be a contiguous CUDA int32 tensor with at least ", + tk::kRouteWords, " elements"); + TORCH_CHECK( + fc1Warps == 1 || fc1Warps == 2 || fc1Warps == 4, "[glm_moe_dsa_draft_routed_moe] fc1_warps must be 1, 2 or 4"); + TORCH_CHECK(fc2Warps == 1 || fc2Warps == 2 || fc2Warps == 4 || (fc2Warps == 3 && (flags & tk::kFlagTma)), + "[glm_moe_dsa_draft_routed_moe] fc2_warps must be 1, 2 or 4 (or 3 with the TMA flag)"); + TORCH_CHECK(fc1Ctas >= 1 && fc1Ctas <= 1024 && fc2Ctas >= 1 && fc2Ctas <= 1024, + "[glm_moe_dsa_draft_routed_moe] grid sizes"); + if (debugFc2.has_value()) + { + checkTensor(debugFc2.value(), "debug_fc2", at::kBFloat16, {rows, tk::kTopK, tk::kHidden}, device); + } + if (timeline.has_value()) + { + TORCH_CHECK(kTimelineBuild, + "[glm_moe_dsa_draft_routed_moe] timeline needs a build with GLM_MOE_DSA_DRAFT_ROUTED_MOE_TIMELINE=1"); + checkTensor(timeline.value(), "timeline", at::kLong, {2, tk::kTimelineCtas, tk::kTimelineSlots}, device); + } + + c10::cuda::CUDAGuard const deviceGuard(x.device()); + at::Tensor out = at::empty({rows, tk::kHidden}, x.options()); + + tk::DraftRoutedMoeParams params{}; + params.logits = routerLogits.const_data_ptr(); + params.bias = reinterpret_cast<__nv_bfloat16 const*>(routingBias.const_data_ptr()); + params.x = reinterpret_cast<__nv_bfloat16 const*>(x.const_data_ptr()); + params.w13 = reinterpret_cast<__nv_bfloat16 const*>(w3w1.const_data_ptr()); + params.w2 = reinterpret_cast<__nv_bfloat16 const*>(w2.const_data_ptr()); + params.out = reinterpret_cast<__nv_bfloat16*>(out.mutable_data_ptr()); + params.h = reinterpret_cast<__nv_bfloat16*>(hScratch.mutable_data_ptr()); + params.route = routeScratch.mutable_data_ptr(); + params.debugFc2 = debugFc2.has_value() ? reinterpret_cast<__nv_bfloat16*>(debugFc2->mutable_data_ptr()) : nullptr; + params.timeline = timeline.has_value() ? timeline->mutable_data_ptr() : nullptr; + params.rows = static_cast(rows); + params.routedScaling = static_cast(routedScaling); + params.fc1Ctas = static_cast(fc1Ctas); + params.fc1Warps = static_cast(fc1Warps); + params.fc2Ctas = static_cast(fc2Ctas); + params.fc2Warps = static_cast(fc2Warps); + params.flags = static_cast(flags); + + tk::DraftRoutedMoeTensorMaps maps{}; + bool const tma = (flags & tk::kFlagTma) != 0; + if (tma) + { + int const f = static_cast(flags); + int const w1 = static_cast(fc1Warps); + maps.w13 = weightMap(device, w3w1.const_data_ptr(), tk::kFc1Rows, + static_cast(tk::kNumExperts) * tk::kFc1Chunks, static_cast(tk::fc1TmaBoxRows(f, w1)), + static_cast(tk::fc1TmaBoxSlabs(f, w1))); + maps.w2 = weightMap(device, w2.const_data_ptr(), tk::kHidden, + static_cast(tk::kNumExperts) * tk::kFc2Chunks, static_cast(16 * fc2Warps), + static_cast(tk::kFc2Chunks)); + // The triples path of the FC1 kernel (4 MMA warps) reads 64-row boxes; otherwise the map is not read. + maps.w13Triple = ((f & tk::kFlagFc1Triples) && w1 == 4) ? weightMap(device, w3w1.const_data_ptr(), tk::kFc1Rows, + static_cast(tk::kNumExperts) * tk::kFc1Chunks, + static_cast(tk::kTripleRows), static_cast(tk::kTripleStageChunks)) + : maps.w13; + // kFlagStreamX: the call's activation as [token][K chunk][64] (boxes of 8 tokens; rows beyond the call's are + // zero-filled and only feed discarded MMA columns). Encoded per call (``x`` changes between calls). + maps.x = maps.w13; + if (f & tk::kFlagStreamX) + { + cuuint64_t const dims[3] = {static_cast(tk::kKChunk), static_cast(tk::kFc1Chunks), + static_cast(rows)}; + cuuint64_t const strides[2] + = {static_cast(tk::kKChunk) * 2, static_cast(tk::kHidden) * 2}; + cuuint32_t const box[3] = {static_cast(tk::kKChunk), static_cast(tk::kXBoxChunks), + static_cast(tk::kMaxRows)}; + cuuint32_t const elementStrides[3] = {1, 1, 1}; + CUresult const result + = encodeTiledFn()(&maps.x, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, const_cast(x.const_data_ptr()), + dims, strides, box, elementStrides, CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, + CU_TENSOR_MAP_L2_PROMOTION_L2_128B, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TORCH_CHECK(result == CUDA_SUCCESS, + "[glm_moe_dsa_draft_routed_moe] cuTensorMapEncodeTiled (x) failed: ", static_cast(result)); + } + } + cudaError_t const status + = tk::invokeDraftRoutedMoe(params, tma ? &maps : nullptr, at::cuda::getCurrentCUDAStream(device)); + TORCH_CHECK( + status == cudaSuccess, "[glm_moe_dsa_draft_routed_moe] kernel launch failed: ", cudaGetErrorString(status)); + return out; +} + +} // namespace torch_ext::glm_moe_dsa + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "glm_moe_dsa_draft_routed_moe(Tensor x, Tensor router_logits, Tensor routing_bias, Tensor w3_w1, Tensor w2, " + "Tensor(a!) h_scratch, Tensor(b!) route_scratch, float routed_scaling, int fc1_ctas, int fc1_warps, " + "int fc2_ctas, int fc2_warps, int flags, Tensor(c!)? debug_fc2=None, Tensor(d!)? timeline=None) -> Tensor"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("glm_moe_dsa_draft_routed_moe", &tensorrt_llm::torch_ext::glm_moe_dsa::draftRoutedMoe); +} diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/draftTailOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsa/draftTailOp.cpp new file mode 100644 index 000000000000..72cda4e7ab47 --- /dev/null +++ b/cpp/tensorrt_llm/thop/glmMoeDsa/draftTailOp.cpp @@ -0,0 +1,168 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Torch ops of the draft tail (greedy draft-token selection over a +// vocab-sharded LM head). See kernels/glmMoeDsa/draftTail.h. + +#include "tensorrt_llm/kernels/glmMoeDsa/draftTail.h" + +#include +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext::glm_moe_dsa +{ + +namespace +{ +namespace tk = tensorrt_llm::kernels::glm_moe_dsa::draft_tail; + +void fillCommon(tk::DraftTailParams& p, int64_t rows, int64_t vocabLocal, at::Tensor& tokens, at::Tensor& scratch, + std::optional const& keysOut, int64_t vocabOffset, int64_t tpRank, int64_t tpSize, int64_t mode, + int64_t mcPtr, int64_t ucPtr, std::optional const& peerPtrs, int64_t timeoutCycles, c10::Device device) +{ + TORCH_CHECK(rows >= 1 && rows <= tk::kMaxRows, "glm_moe_dsa_draft_tail: rows must be in [1, ", tk::kMaxRows, "]"); + TORCH_CHECK(tokens.device() == device && tokens.scalar_type() == at::kInt && tokens.is_contiguous() + && tokens.numel() >= rows, + "glm_moe_dsa_draft_tail: tokens must be a contiguous int32 tensor with at least rows elements"); + TORCH_CHECK(scratch.device() == device && scratch.scalar_type() == at::kLong && scratch.is_contiguous() + && scratch.numel() >= tk::kScratchWords, + "glm_moe_dsa_draft_tail: scratch must be a contiguous int64 tensor with ", tk::kScratchWords, " elements"); + TORCH_CHECK(vocabOffset >= 0 && vocabOffset + vocabLocal <= 0xFFFFFFFFll, + "glm_moe_dsa_draft_tail: vocab range out of bounds"); + TORCH_CHECK(tpSize >= 1 && tpSize <= tk::kMaxRanks && tpRank >= 0 && tpRank < tpSize, + "glm_moe_dsa_draft_tail: bad tp rank"); + TORCH_CHECK(mode >= 0 && mode <= 4, "glm_moe_dsa_draft_tail: bad exchange mode"); + p.rows = static_cast(rows); + p.vocabLocal = static_cast(vocabLocal); + p.vocabOffset = vocabOffset; + p.tokens = tokens.data_ptr(); + p.scratch = reinterpret_cast(scratch.data_ptr()); + if (keysOut.has_value()) + { + TORCH_CHECK(keysOut->device() == device && keysOut->scalar_type() == at::kLong && keysOut->is_contiguous() + && keysOut->numel() >= rows, + "glm_moe_dsa_draft_tail: keys_out must be a contiguous int64 tensor with at least rows elements"); + p.keysOut = reinterpret_cast(keysOut->data_ptr()); + } + p.tpRank = static_cast(tpRank); + p.tpSize = static_cast(tpSize); + p.mode = static_cast(mode); + p.timeoutCycles = timeoutCycles; + if (p.mode != tk::ExchangeMode::kNone && tpSize > 1) + { + TORCH_CHECK(ucPtr != 0 && ucPtr % 16 == 0, "glm_moe_dsa_draft_tail: exchange needs the local unicast buffer"); + p.ucSlotsLocal = reinterpret_cast(ucPtr); + if (p.mode == tk::ExchangeMode::kMulticast || p.mode == tk::ExchangeMode::kMulticastDirect) + { + TORCH_CHECK( + mcPtr != 0 && mcPtr % 16 == 0, "glm_moe_dsa_draft_tail: multicast exchange needs the multicast buffer"); + p.mcSlots = reinterpret_cast(mcPtr); + } + else + { + TORCH_CHECK(peerPtrs.has_value() && peerPtrs->device() == device && peerPtrs->scalar_type() == at::kLong + && peerPtrs->is_contiguous() && peerPtrs->numel() >= tpSize, + "glm_moe_dsa_draft_tail: unicast exchange needs a device int64 tensor of tp_size peer addresses"); + p.ucSlotsPeers = reinterpret_cast(peerPtrs->data_ptr()); + } + } +} + +} // namespace + +void draftTail(at::Tensor const& hidden, at::Tensor const& weight, at::Tensor& tokens, at::Tensor& scratch, + std::optional const& keysOut, int64_t vocabOffset, int64_t tpRank, int64_t tpSize, int64_t mode, + int64_t mcPtr, int64_t ucPtr, std::optional const& peerPtrs, int64_t variant, int64_t numCtas, + int64_t keepPermille, bool pdl, int64_t timeoutCycles, std::optional const& logitsOut) +{ + TORCH_CHECK(hidden.is_cuda() && hidden.dim() == 2 && hidden.scalar_type() == at::kBFloat16, + "glm_moe_dsa_draft_tail: hidden must be a 2D BF16 CUDA tensor"); + TORCH_CHECK(weight.device() == hidden.device() && weight.dim() == 2 && weight.scalar_type() == at::kBFloat16 + && weight.is_contiguous(), + "glm_moe_dsa_draft_tail: weight must be a contiguous 2D BF16 tensor on the hidden device"); + TORCH_CHECK(hidden.size(1) == 6144 && weight.size(1) == 6144, "glm_moe_dsa_draft_tail: hidden size must be 6144"); + TORCH_CHECK(hidden.stride(1) == 1 && hidden.stride(0) % 8 == 0 + && reinterpret_cast(hidden.data_ptr()) % 16 == 0 + && reinterpret_cast(weight.data_ptr()) % 32 == 0, + "glm_moe_dsa_draft_tail: hidden rows must be 16-byte aligned and the weight 32-byte aligned"); + tk::DraftTailParams p; + fillCommon(p, hidden.size(0), weight.size(0), tokens, scratch, keysOut, vocabOffset, tpRank, tpSize, mode, mcPtr, + ucPtr, peerPtrs, timeoutCycles, hidden.device()); + p.hidden = reinterpret_cast<__nv_bfloat16 const*>(hidden.data_ptr()); + p.hiddenStride = hidden.stride(0); + p.weight = reinterpret_cast<__nv_bfloat16 const*>(weight.data_ptr()); + p.hiddenDim = static_cast(hidden.size(1)); + TORCH_CHECK( + keepPermille >= 0 && keepPermille <= 1000, "glm_moe_dsa_draft_tail: keep_permille must be in [0, 1000]"); + p.keepPermille = static_cast(keepPermille); + if (logitsOut.has_value()) + { + TORCH_CHECK(logitsOut->device() == hidden.device() && logitsOut->scalar_type() == at::kBFloat16 + && logitsOut->dim() == 2 && logitsOut->size(0) >= hidden.size(0) && logitsOut->size(1) == weight.size(0) + && logitsOut->stride(1) == 1, + "glm_moe_dsa_draft_tail: logits_out must be a BF16 [rows, vocab_local] tensor with unit last stride"); + p.logitsOut = reinterpret_cast<__nv_bfloat16*>(logitsOut->data_ptr()); + p.logitsOutStride = logitsOut->stride(0); + } + tk::FusedLaunchConfig config; + config.numCtas = static_cast(numCtas); + config.variant = static_cast(variant); + config.pdl = pdl; + tk::invokeDraftTailFused(p, config, at::cuda::getCurrentCUDAStream(hidden.get_device())); +} + +void draftTailFromLogits(at::Tensor const& logits, at::Tensor& tokens, at::Tensor& scratch, + std::optional const& keysOut, int64_t vocabOffset, int64_t tpRank, int64_t tpSize, int64_t mode, + int64_t mcPtr, int64_t ucPtr, std::optional const& peerPtrs, bool pdl, int64_t timeoutCycles) +{ + TORCH_CHECK(logits.is_cuda() && logits.dim() == 2 && logits.scalar_type() == at::kBFloat16 && logits.stride(1) == 1, + "glm_moe_dsa_draft_tail: logits must be a 2D BF16 CUDA tensor with unit last stride"); + tk::DraftTailParams p; + fillCommon(p, logits.size(0), logits.size(1), tokens, scratch, keysOut, vocabOffset, tpRank, tpSize, mode, mcPtr, + ucPtr, peerPtrs, timeoutCycles, logits.device()); + p.logits = reinterpret_cast<__nv_bfloat16 const*>(logits.data_ptr()); + p.logitsStride = logits.stride(0); + tk::invokeDraftTailFromLogits(p, pdl, at::cuda::getCurrentCUDAStream(logits.get_device())); +} + +} // namespace torch_ext::glm_moe_dsa + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "glm_moe_dsa_draft_tail(Tensor hidden, Tensor weight, Tensor(a!) tokens, Tensor(b!) scratch, " + "Tensor(c!)? keys_out, int vocab_offset, int tp_rank, int tp_size, int mode, int mc_ptr, int uc_ptr, " + "Tensor? peer_ptrs, int variant, int num_ctas, int keep_permille, bool pdl, int timeout_cycles, " + "Tensor(d!)? logits_out=None) -> ()"); + m.def( + "glm_moe_dsa_draft_tail_from_logits(Tensor logits, Tensor(a!) tokens, Tensor(b!) scratch, Tensor(c!)? " + "keys_out, " + "int vocab_offset, int tp_rank, int tp_size, int mode, int mc_ptr, int uc_ptr, Tensor? peer_ptrs, bool pdl, " + "int timeout_cycles) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("glm_moe_dsa_draft_tail", &tensorrt_llm::torch_ext::glm_moe_dsa::draftTail); + m.impl("glm_moe_dsa_draft_tail_from_logits", &tensorrt_llm::torch_ext::glm_moe_dsa::draftTailFromLogits); +} diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/indexerLayerRopeKvAppendOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsa/indexerLayerRopeKvAppendOp.cpp new file mode 100644 index 000000000000..d51ca99cf61d --- /dev/null +++ b/cpp/tensorrt_llm/thop/glmMoeDsa/indexerLayerRopeKvAppendOp.cpp @@ -0,0 +1,242 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/envUtils.h" +#include "tensorrt_llm/common/quantization.h" +#include "tensorrt_llm/kernels/glmMoeDsa/indexerLayerRopeKvAppend.h" +#include "tensorrt_llm/kernels/glmMoeDsa/ropeKvAppend.h" +#include "tensorrt_llm/kernels/kvCacheUtils.h" + +#include +#include + +#include +#include +#include +#include + +namespace tk = tensorrt_llm::kernels; +namespace tc = tensorrt_llm::common; +namespace rope_kv_append = tensorrt_llm::kernels::glm_moe_dsa::rope_kv_append; +namespace indexer_layer_front = tensorrt_llm::kernels::glm_moe_dsa::indexer_layer_front; + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext::glm_moe_dsa +{ + +namespace +{ + +void checkCuda(torch::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda(), name, " must be a CUDA tensor"); +} + +//! Paged MLA latent cache of one layer: the KVBlockArray that buildPagedKvCacheBuffers (thop/attentionOp.cpp) +//! builds for an FP8 MLA cache (kvFactor 1, one KV head, 8-bit elements, beam width 1, no FP4 scale pool). +tk::KVBlockArray buildFp8MlaKvBlockArray(torch::Tensor const& kvCacheBlockOffsets, + torch::Tensor const& hostKvCachePoolPointers, torch::Tensor const& hostKvCachePoolMapping, int64_t layerIdx, + int32_t batchSize, int64_t tokensPerBlock, int64_t attentionWindowSize, int64_t seqOffset) +{ + TORCH_CHECK(hostKvCachePoolMapping.device().is_cpu() && hostKvCachePoolMapping.dim() == 2 + && hostKvCachePoolMapping.size(1) >= 2 && hostKvCachePoolMapping.scalar_type() == torch::kInt32, + "host_kv_cache_pool_mapping must be a CPU int32 [layers, >=2] tensor"); + TORCH_CHECK(layerIdx >= 0 && layerIdx < hostKvCachePoolMapping.size(0), "layer index out of the pool mapping"); + auto const* mapping = hostKvCachePoolMapping.data_ptr(); + int32_t const poolIndex = mapping[layerIdx * hostKvCachePoolMapping.stride(0)]; + int32_t const layerIdxInCachePool + = mapping[layerIdx * hostKvCachePoolMapping.stride(0) + hostKvCachePoolMapping.stride(1)]; + + TORCH_CHECK(hostKvCachePoolPointers.device().is_cpu() && hostKvCachePoolPointers.dim() == 2 + && hostKvCachePoolPointers.scalar_type() == torch::kInt64 && poolIndex >= 0 + && poolIndex < hostKvCachePoolPointers.size(0) && hostKvCachePoolPointers.size(1) >= 2, + "host_kv_cache_pool_pointers must be a CPU int64 [pools, 2] tensor"); + TORCH_CHECK(kvCacheBlockOffsets.is_cuda() && kvCacheBlockOffsets.scalar_type() == torch::kInt32 + && kvCacheBlockOffsets.dim() >= 2 && poolIndex < kvCacheBlockOffsets.size(0) + && seqOffset < kvCacheBlockOffsets.size(1), + "kv_cache_block_offsets must be CUDA int32 [pools, seqs, 2, max_blocks]"); + + int64_t constexpr kCacheElemBits = 8; + int64_t const blockSize = tokensPerBlock * /*kv_head_num=*/1 * rope_kv_append::kLatentDim; + int64_t const bytesPerBlock = blockSize * kCacheElemBits / CHAR_BIT; + int64_t const intraPoolOffset = layerIdxInCachePool * /*kvFactor=*/1 * bytesPerBlock; + auto const sizePerToken = static_cast(rope_kv_append::kLatentDim * kCacheElemBits / 8); + auto const* pointers = hostKvCachePoolPointers.data_ptr(); + auto const rowOffset = poolIndex * hostKvCachePoolPointers.stride(0); + void* primaryPoolPtr = reinterpret_cast(pointers[rowOffset]) + intraPoolOffset; + void* secondaryPoolPtr + = reinterpret_cast(pointers[rowOffset + hostKvCachePoolPointers.stride(1)]) + intraPoolOffset; + auto* blockOffsets + = reinterpret_cast(static_cast(kvCacheBlockOffsets.data_ptr()) + + poolIndex * kvCacheBlockOffsets.stride(0) + seqOffset * kvCacheBlockOffsets.stride(1)); + auto const maxBlocksPerSequence = static_cast(kvCacheBlockOffsets.size(-1)); + auto const window = static_cast(attentionWindowSize); + return tk::KVBlockArray(batchSize, maxBlocksPerSequence, static_cast(tokensPerBlock), sizePerToken, window, + window, /*sinkTokenLen=*/0, /*canUseOneMoreBlock=*/false, primaryPoolPtr, secondaryPoolPtr, blockOffsets); +} + +} // namespace + +//! Variant of trtllm::glm_moe_dsa_rope_kv_append: same arguments, checks and written bytes; with early_trigger each +//! kernel thread triggers the programmatic dependents right after the global reads of its work item (otherwise at the +//! end of the kernel, as in trtllm::glm_moe_dsa_rope_kv_append). +void indexerLayerRopeKvAppend(torch::Tensor const& fused_q, torch::Tensor const& q_pe, + torch::Tensor const& latent_cache, torch::Tensor const& rotary_cos_sin, torch::Tensor const& cu_q_seqlens, + torch::Tensor const& cu_kv_seqlens, torch::Tensor const& fmha_scheduler_counter, + torch::Tensor const& mla_bmm1_scale, torch::Tensor const& mla_bmm2_scale, torch::Tensor const& quant_q_buffer, + torch::Tensor const& sequence_length, torch::Tensor const& host_context_lengths, int64_t num_contexts, + torch::Tensor const& kv_cache_block_offsets, torch::Tensor const& host_kv_cache_pool_pointers, + torch::Tensor const& host_kv_cache_pool_mapping, std::optional const& kv_cache_scale_orig_quant, + int64_t layer_idx, int64_t num_heads, int64_t tokens_per_block, int64_t attention_window_size, int64_t quant_mode, + double q_scaling, int64_t qk_nope_head_dim, int64_t qk_rope_head_dim, int64_t kv_lora_rank, + std::optional const& idx_k_fp8, std::optional const& idx_k_scale, + std::optional const& idx_k_cache, std::optional const& idx_slot_mapping_fp8, + std::optional const& idx_slot_mapping_scale, int64_t idx_num_tokens, bool early_trigger) +{ + TORCH_CHECK(kv_lora_rank == rope_kv_append::kKvLoraRank && qk_rope_head_dim == rope_kv_append::kRopeDim, + "glm_moe_dsa_indexer_layer_rope_kv_append: supports kv_lora_rank 512 and qk_rope_head_dim 64 only"); + auto const kvQuantMode = tc::QuantMode(static_cast(quant_mode)); + TORCH_CHECK(kvQuantMode.hasFp8KvCache() && !kvQuantMode.hasFp4KvCache(), + "glm_moe_dsa_indexer_layer_rope_kv_append: needs an FP8 MLA cache"); + checkCuda(fused_q, "fused_q"); + checkCuda(q_pe, "q_pe"); + checkCuda(latent_cache, "latent_cache"); + TORCH_CHECK(fused_q.scalar_type() == torch::kBFloat16 && q_pe.scalar_type() == torch::kBFloat16 + && latent_cache.scalar_type() == torch::kBFloat16, + "glm_moe_dsa_indexer_layer_rope_kv_append: needs BF16 fused_q, q_pe and latent_cache"); + TORCH_CHECK(rotary_cos_sin.scalar_type() == torch::kFloat32 && rotary_cos_sin.is_contiguous(), + "rotary_cos_sin must be packed float32"); + int64_t const cosSinRows = rotary_cos_sin.numel() / (2 * rope_kv_append::kRopeDim); + TORCH_CHECK(cosSinRows > 0 && cosSinRows <= INT32_MAX, "rotary_cos_sin has no complete position row"); + TORCH_CHECK(q_pe.dim() == 3 && q_pe.stride(2) == 1, "q_pe must be [tokens, heads, rope] with unit inner stride"); + TORCH_CHECK(latent_cache.dim() == 2 && latent_cache.size(1) == rope_kv_append::kLatentDim + && latent_cache.stride(1) == 1 && latent_cache.stride(0) == rope_kv_append::kLatentDim, + "latent_cache must be packed [tokens, 576]"); + TORCH_CHECK(fused_q.is_contiguous() && quant_q_buffer.is_contiguous(), "fused_q and quant_q_buffer must be packed"); + TORCH_CHECK(host_kv_cache_pool_mapping.dim() >= 1, "host_kv_cache_pool_mapping is required"); + + int32_t const numTokens = static_cast(latent_cache.size(0)); + int32_t const numSeqs = static_cast(host_context_lengths.size(0)); + int32_t const numGenerations = numSeqs - static_cast(num_contexts); + TORCH_CHECK(numGenerations > 0, "glm_moe_dsa_indexer_layer_rope_kv_append: needs at least one generation request"); + TORCH_CHECK(numTokens % numGenerations == 0, "MLA can only support input sequences with the same sequence length."); + TORCH_CHECK(fused_q.numel() >= static_cast(numTokens) * num_heads * rope_kv_append::kLatentDim, + "fused_q is smaller than [tokens, heads, 576]"); + TORCH_CHECK(quant_q_buffer.numel() * quant_q_buffer.element_size() + >= static_cast(numTokens) * num_heads * rope_kv_append::kLatentDim, + "quant_q_buffer is smaller than [tokens, heads, 576] bytes"); + TORCH_CHECK(mla_bmm1_scale.numel() >= 2 && mla_bmm2_scale.numel() >= 1, "bmm scale buffers too small"); + TORCH_CHECK(cu_q_seqlens.numel() >= numGenerations + 1 && cu_kv_seqlens.numel() >= numGenerations + 1, + "cu_seqlens buffers too small"); + + auto stream = at::cuda::getCurrentCUDAStream(latent_cache.get_device()); + rope_kv_append::RopeAppendParams params{}; + params.fusedQ = reinterpret_cast<__nv_bfloat16 const*>(fused_q.data_ptr()); + params.qPe = reinterpret_cast<__nv_bfloat16 const*>(q_pe.data_ptr()); + params.qPeLd = static_cast(q_pe.stride(1)); + params.qPeStride = static_cast(q_pe.stride(0)); + params.quantQ = reinterpret_cast(quant_q_buffer.data_ptr()); + params.quantScaleQ = nullptr; + params.latent = reinterpret_cast<__nv_bfloat16 const*>(latent_cache.data_ptr()); + params.cosSin = reinterpret_cast(rotary_cos_sin.data_ptr()); + params.cosSinRows = static_cast(cosSinRows); + params.kvLens = sequence_length.slice(0, num_contexts).data_ptr(); + params.quantScaleKv + = kv_cache_scale_orig_quant.has_value() ? kv_cache_scale_orig_quant->data_ptr() : nullptr; + params.kvCache = buildFp8MlaKvBlockArray(kv_cache_block_offsets, host_kv_cache_pool_pointers, + host_kv_cache_pool_mapping, layer_idx, numGenerations, tokens_per_block, attention_window_size, num_contexts); + params.cuQSeqlens = cu_q_seqlens.data_ptr(); + params.cuKvSeqlens = cu_kv_seqlens.data_ptr(); + params.fmhaTileCounter = reinterpret_cast(fmha_scheduler_counter.data_ptr()); + params.bmm1Scale = mla_bmm1_scale.data_ptr(); + params.bmm2Scale = mla_bmm2_scale.data_ptr(); + params.dequantScaleQ = nullptr; + params.dequantScaleKv = nullptr; + params.quantScaleO = nullptr; + // Same expression as the native op so the folded FMHA scales match bit for bit. + params.hostBmm1Scale = 1.F / (q_scaling * std::sqrt(static_cast(qk_nope_head_dim + qk_rope_head_dim))); + params.numTokens = numTokens; + params.seqLen = numTokens / numGenerations; + params.numHeads = static_cast(num_heads); + + if (idx_num_tokens > 0) + { + TORCH_CHECK(idx_k_fp8.has_value() && idx_k_scale.has_value() && idx_k_cache.has_value() + && idx_slot_mapping_fp8.has_value() && idx_slot_mapping_scale.has_value(), + "indexer scatter needs k_fp8, k_scale, k_cache and both slot mappings"); + auto const& kFp8 = idx_k_fp8.value(); + auto const& kScale = idx_k_scale.value(); + auto const& kCache = idx_k_cache.value(); + auto const& slotFp8 = idx_slot_mapping_fp8.value(); + auto const& slotScale = idx_slot_mapping_scale.value(); + checkCuda(kFp8, "k_fp8"); + checkCuda(kScale, "k_scale"); + checkCuda(kCache, "k_cache"); + TORCH_CHECK(kFp8.dim() == 2 && kFp8.element_size() == 1 && kFp8.is_contiguous(), "k_fp8 must be packed bytes"); + TORCH_CHECK(kScale.dim() == 2 && kScale.element_size() == 4 && kScale.is_contiguous() + && kScale.size(1) * kScale.element_size() == 4, + "k_scale must be one packed 4-byte word per token"); + TORCH_CHECK(kCache.dim() == 4 && kCache.size(2) == 1, "k_cache must be [blocks, block_size, 1, bytes]"); + TORCH_CHECK(slotFp8.scalar_type() == torch::kInt64 && slotScale.scalar_type() == torch::kInt64 + && slotFp8.is_contiguous() && slotScale.is_contiguous(), + "slot mappings must be packed int64"); + int32_t const headBytes = static_cast(kFp8.size(1)); + TORCH_CHECK(headBytes == 128 || headBytes == 64, "indexer k head bytes must be 128 or 64"); + TORCH_CHECK(kFp8.size(0) >= idx_num_tokens && kScale.size(0) >= idx_num_tokens + && slotFp8.numel() >= idx_num_tokens && slotScale.numel() >= idx_num_tokens, + "indexer inputs are shorter than idx_num_tokens"); + params.idxK = reinterpret_cast(kFp8.data_ptr()); + params.idxScale = reinterpret_cast(kScale.data_ptr()); + params.idxCache = reinterpret_cast(kCache.data_ptr()); + params.idxSlotData = slotFp8.data_ptr(); + params.idxSlotScale = slotScale.data_ptr(); + params.idxNumTokens = static_cast(idx_num_tokens); + params.idxHeadBytes = headBytes; + for (int32_t d = 0; d < 4; ++d) + { + params.idxDims[d] = static_cast(kCache.size(d)); + params.idxStrides[d] = kCache.stride(d); + } + } + + indexer_layer_front::invokeIndexerLayerRopeKvAppend(params, tc::getEnvEnablePDL(), early_trigger, stream); +} + +} // namespace torch_ext::glm_moe_dsa + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "glm_moe_dsa_indexer_layer_rope_kv_append(Tensor fused_q, Tensor q_pe, Tensor latent_cache, Tensor " + "rotary_cos_sin, " + "Tensor(a!) cu_q_seqlens, Tensor(b!) cu_kv_seqlens, Tensor(c!) fmha_scheduler_counter, " + "Tensor(d!) mla_bmm1_scale, Tensor(e!) mla_bmm2_scale, Tensor(f!) quant_q_buffer, Tensor sequence_length, " + "Tensor host_context_lengths, int num_contexts, Tensor kv_cache_block_offsets, " + "Tensor host_kv_cache_pool_pointers, Tensor host_kv_cache_pool_mapping, Tensor? kv_cache_scale_orig_quant, " + "int layer_idx, int num_heads, int tokens_per_block, int attention_window_size, int quant_mode, " + "float q_scaling, int qk_nope_head_dim, int qk_rope_head_dim, int kv_lora_rank, Tensor? idx_k_fp8, " + "Tensor? idx_k_scale, Tensor(g!)? idx_k_cache, Tensor? idx_slot_mapping_fp8, Tensor? idx_slot_mapping_scale, " + "int idx_num_tokens, bool early_trigger) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("glm_moe_dsa_indexer_layer_rope_kv_append", &tensorrt_llm::torch_ext::glm_moe_dsa::indexerLayerRopeKvAppend); +} diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/postAttentionHopFromPartialsOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsa/postAttentionHopFromPartialsOp.cpp new file mode 100644 index 000000000000..7d94a64871ff --- /dev/null +++ b/cpp/tensorrt_llm/thop/glmMoeDsa/postAttentionHopFromPartialsOp.cpp @@ -0,0 +1,150 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/config.h" +#include "tensorrt_llm/kernels/glmMoeDsa/postAttentionHopFromPartials.h" + +#include +#include +#include +#include + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext::glm_moe_dsa +{ + +namespace +{ + +namespace tk = tensorrt_llm::kernels::glm_moe_dsa::attention_output_into_hop; + +// Same switch as the native kernels: PDL is on unless TRTLLM_ENABLE_PDL=0. +bool pdlEnabled() +{ + static bool const enabled = []() + { + char const* env = std::getenv("TRTLLM_ENABLE_PDL"); + return !(env != nullptr && std::strcmp(env, "0") == 0); + }(); + return enabled; +} + +void checkRows(at::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda() && t.scalar_type() == at::kBFloat16, "[glm_moe_dsa_post_attention_hop_from_partials] ", + name, " must be a CUDA bf16 tensor"); + TORCH_CHECK(t.is_contiguous(), "[glm_moe_dsa_post_attention_hop_from_partials] ", name, " must be contiguous"); + TORCH_CHECK(t.dim() >= 1 && t.size(-1) == tk::kHiddenDim, "[glm_moe_dsa_post_attention_hop_from_partials] ", name, + " must have hidden size ", tk::kHiddenDim); +} + +} // namespace + +//! Post-attention hop whose rank-local partials were already published into the hop's transport by the attention +//! output tail. ``residual`` gives the row count. Returns ``[normed, residual_out, router_logits]``; ``router_logits`` +//! is an empty tensor when ``router_weight`` is not given. See kernels/glmMoeDsa/postAttentionHopFromPartials.h. +std::vector postAttentionHopFromPartials(at::Tensor const& residual, at::Tensor const& gamma, double eps, + std::optional const& routerWeight, at::Tensor const& ucBuffer, at::Tensor const& flags, int64_t rank, + int64_t nRanks, bool earlyPoll, int64_t routerLoad) +{ + checkRows(residual, "residual"); + checkRows(gamma, "gamma"); + TORCH_CHECK(gamma.dim() == 1, "[glm_moe_dsa_post_attention_hop_from_partials] gamma must be 1-D"); + int64_t const numRows = residual.numel() / tk::kHiddenDim; + TORCH_CHECK(numRows >= 1 && numRows <= tk::kMaxRows, + "[glm_moe_dsa_post_attention_hop_from_partials] number of rows ", numRows, + " is outside the supported range [1, ", tk::kMaxRows, "]"); + TORCH_CHECK(nRanks == tk::kNumRanks, "[glm_moe_dsa_post_attention_hop_from_partials] only TP", tk::kNumRanks, + " is supported, got ", nRanks); + TORCH_CHECK(rank >= 0 && rank < nRanks, "[glm_moe_dsa_post_attention_hop_from_partials] invalid rank ", rank); + TORCH_CHECK(routerLoad >= 0 && routerLoad <= 2, + "[glm_moe_dsa_post_attention_hop_from_partials] router_load must be 0, 1 or 2, got ", routerLoad); + TORCH_CHECK(ucBuffer.is_cuda() && ucBuffer.is_contiguous() + && static_cast(ucBuffer.numel() * ucBuffer.element_size()) >= tk::kWorkspaceBytes, + "[glm_moe_dsa_post_attention_hop_from_partials] uc_buffer must be a contiguous CUDA tensor of at least ", + tk::kWorkspaceBytes, " bytes"); + TORCH_CHECK(flags.is_cuda() && flags.is_contiguous() && flags.scalar_type() == at::kInt && flags.numel() >= 4, + "[glm_moe_dsa_post_attention_hop_from_partials] flags must be a contiguous CUDA int32 tensor with at least 4 " + "elements"); + TORCH_CHECK(residual.get_device() == ucBuffer.get_device() && residual.get_device() == flags.get_device(), + "[glm_moe_dsa_post_attention_hop_from_partials] all tensors must be on the same device"); + + at::Tensor routerLogits; + if (routerWeight.has_value()) + { + auto const& w = routerWeight.value(); + TORCH_CHECK(w.is_cuda() && w.scalar_type() == at::kBFloat16 && w.is_contiguous() && w.dim() == 2 + && w.size(0) == tk::kNumExperts && w.size(1) == tk::kHiddenDim, + "[glm_moe_dsa_post_attention_hop_from_partials] router_weight must be a contiguous CUDA bf16 tensor of " + "shape [", + tk::kNumExperts, ", ", tk::kHiddenDim, "]"); + routerLogits = at::empty({numRows, tk::kNumExperts}, residual.options().dtype(at::kFloat)); + } + else + { + routerLogits = at::empty({0}, residual.options().dtype(at::kFloat)); + } + + c10::cuda::CUDAGuard const deviceGuard(residual.device()); + at::Tensor normOut = at::empty_like(residual); + at::Tensor residualOut = at::empty_like(residual); + + tk::PostAttentionHopFromPartialsParams params; + params.residualIn = residual.const_data_ptr(); + params.gamma = gamma.const_data_ptr(); + params.routerWeight = routerWeight.has_value() ? routerWeight->const_data_ptr() : nullptr; + params.normOut = normOut.mutable_data_ptr(); + params.residualOut = residualOut.mutable_data_ptr(); + params.routerLogits = routerWeight.has_value() ? routerLogits.mutable_data_ptr() : nullptr; + params.ucBuffer = ucBuffer.mutable_data_ptr(); + params.flags = reinterpret_cast(flags.mutable_data_ptr()); + params.numRows = static_cast(numRows); + params.rank = static_cast(rank); + params.eps = static_cast(eps); + params.enablePdl = pdlEnabled(); + params.earlyPoll = earlyPoll; + params.routerLoad = static_cast(routerLoad); + params.stream = at::cuda::getCurrentCUDAStream(residual.get_device()); + + cudaError_t const status = tk::invokePostAttentionHopFromPartials(params); + TORCH_CHECK(status == cudaSuccess, + "[glm_moe_dsa_post_attention_hop_from_partials] kernel launch failed: ", cudaGetErrorString(status)); + return {normOut, residualOut, routerLogits}; +} + +} // namespace torch_ext::glm_moe_dsa + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "glm_moe_dsa_post_attention_hop_from_partials(Tensor residual, Tensor gamma, float eps, Tensor? router_weight, " + "Tensor uc_buffer, Tensor flags, int rank, int nranks, bool early_poll, int router_load=1) -> Tensor[]"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("glm_moe_dsa_post_attention_hop_from_partials", + &tensorrt_llm::torch_ext::glm_moe_dsa::postAttentionHopFromPartials); +} diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/postAttentionHopOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsa/postAttentionHopOp.cpp new file mode 100644 index 000000000000..de99e3968c6b --- /dev/null +++ b/cpp/tensorrt_llm/thop/glmMoeDsa/postAttentionHopOp.cpp @@ -0,0 +1,170 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/config.h" +#include "tensorrt_llm/kernels/glmMoeDsa/postAttentionHop.h" + +#include +#include +#include +#include + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext::glm_moe_dsa +{ + +namespace +{ + +namespace tk = tensorrt_llm::kernels::glm_moe_dsa::post_attention_hop; + +// Same switch as the native kernels: PDL is on unless TRTLLM_ENABLE_PDL=0. +bool pdlEnabled() +{ + static bool const enabled = []() + { + char const* env = std::getenv("TRTLLM_ENABLE_PDL"); + return !(env != nullptr && std::strcmp(env, "0") == 0); + }(); + return enabled; +} + +void checkRows(at::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda() && t.scalar_type() == at::kBFloat16, "[glm_moe_dsa_post_attention_hop] ", name, + " must be a CUDA bf16 tensor"); + TORCH_CHECK(t.is_contiguous(), "[glm_moe_dsa_post_attention_hop] ", name, " must be contiguous"); + TORCH_CHECK(t.dim() >= 1 && t.size(-1) == tk::kHiddenDim, "[glm_moe_dsa_post_attention_hop] ", name, + " must have hidden size ", tk::kHiddenDim); +} + +} // namespace + +//! Fused post-attention hop: one-shot cross-rank sum of ``input`` + residual + RMSNorm (+ router logits). +//! Returns ``[normed, residual_out, router_logits]``; ``router_logits`` is an empty tensor when ``router_weight`` is +//! not given. See kernels/glmMoeDsa/postAttentionHop.h for the arithmetic and the workspace protocol. +std::vector postAttentionHop(at::Tensor const& input, at::Tensor const& residual, at::Tensor const& gamma, + double eps, std::optional const& routerWeight, at::Tensor const& ucBuffer, int64_t mcPtr, + at::Tensor const& flags, int64_t rank, int64_t nRanks) +{ + checkRows(input, "input"); + checkRows(residual, "residual"); + TORCH_CHECK( + residual.sizes() == input.sizes(), "[glm_moe_dsa_post_attention_hop] residual must have the shape of input"); + checkRows(gamma, "gamma"); + TORCH_CHECK(gamma.dim() == 1, "[glm_moe_dsa_post_attention_hop] gamma must be 1-D"); + int64_t const numRows = input.numel() / tk::kHiddenDim; + TORCH_CHECK(numRows >= 1 && numRows <= tk::kMaxRows, "[glm_moe_dsa_post_attention_hop] number of rows ", numRows, + " is outside the supported range [1, ", tk::kMaxRows, "]"); + TORCH_CHECK(nRanks == tk::kNumRanks, "[glm_moe_dsa_post_attention_hop] only TP", tk::kNumRanks, + " is supported, got ", nRanks); + TORCH_CHECK(rank >= 0 && rank < nRanks, "[glm_moe_dsa_post_attention_hop] invalid rank ", rank); + TORCH_CHECK(ucBuffer.is_cuda() && ucBuffer.is_contiguous() + && static_cast(ucBuffer.numel() * ucBuffer.element_size()) >= tk::kWorkspaceBytes, + "[glm_moe_dsa_post_attention_hop] uc_buffer must be a contiguous CUDA tensor of at least ", tk::kWorkspaceBytes, + " bytes"); + TORCH_CHECK( + mcPtr != 0, "[glm_moe_dsa_post_attention_hop] mc_ptr must be the multicast address of uc_buffer's workspace"); + TORCH_CHECK(flags.is_cuda() && flags.is_contiguous() && flags.scalar_type() == at::kInt && flags.numel() >= 4, + "[glm_moe_dsa_post_attention_hop] flags must be a contiguous CUDA int32 tensor with at least 4 elements"); + TORCH_CHECK(input.get_device() == ucBuffer.get_device() && input.get_device() == flags.get_device(), + "[glm_moe_dsa_post_attention_hop] all tensors must be on the same device"); + + at::Tensor routerLogits; + if (routerWeight.has_value()) + { + auto const& w = routerWeight.value(); + TORCH_CHECK(w.is_cuda() && w.scalar_type() == at::kBFloat16 && w.is_contiguous() && w.dim() == 2 + && w.size(0) == tk::kNumExperts && w.size(1) == tk::kHiddenDim, + "[glm_moe_dsa_post_attention_hop] router_weight must be a contiguous CUDA bf16 tensor of shape [", + tk::kNumExperts, ", ", tk::kHiddenDim, "]"); + routerLogits = at::empty({numRows, tk::kNumExperts}, input.options().dtype(at::kFloat)); + } + else + { + routerLogits = at::empty({0}, input.options().dtype(at::kFloat)); + } + + c10::cuda::CUDAGuard const deviceGuard(input.device()); + at::Tensor normOut = at::empty_like(input); + at::Tensor residualOut = at::empty_like(residual); + + tk::PostAttnHopParams params; + params.input = input.const_data_ptr(); + params.residualIn = residual.const_data_ptr(); + params.gamma = gamma.const_data_ptr(); + params.routerWeight = routerWeight.has_value() ? routerWeight->const_data_ptr() : nullptr; + params.normOut = normOut.mutable_data_ptr(); + params.residualOut = residualOut.mutable_data_ptr(); + params.routerLogits = routerWeight.has_value() ? routerLogits.mutable_data_ptr() : nullptr; + params.ucBuffer = ucBuffer.mutable_data_ptr(); + params.mcBuffer = reinterpret_cast(mcPtr); + params.flags = reinterpret_cast(flags.mutable_data_ptr()); + params.numRows = static_cast(numRows); + params.rank = static_cast(rank); + params.eps = static_cast(eps); + params.enablePdl = pdlEnabled(); + params.stream = at::cuda::getCurrentCUDAStream(input.get_device()); + + cudaError_t const status = tk::invokePostAttentionHop(params); + TORCH_CHECK( + status == cudaSuccess, "[glm_moe_dsa_post_attention_hop] kernel launch failed: ", cudaGetErrorString(status)); + return {normOut, residualOut, routerLogits}; +} + +#ifdef TLLM_GLM_LL_W2_TIMING +//! Development builds only: phase timestamps of the most recent hop launch. +at::Tensor postAttentionHopTiming() +{ + at::Tensor out = at::empty({tk::kTimingMaxCtas, tk::kTimingStamps}, at::TensorOptions().dtype(at::kLong)); + auto stream = at::cuda::getCurrentCUDAStream(); + TORCH_CHECK(tk::copyPostAttnHopTiming(reinterpret_cast(out.data_ptr()), stream) + == cudaSuccess); + TORCH_CHECK(cudaStreamSynchronize(stream) == cudaSuccess); + return out; +} +#endif + +} // namespace torch_ext::glm_moe_dsa + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "glm_moe_dsa_post_attention_hop(Tensor input, Tensor residual, Tensor gamma, float eps, Tensor? router_weight, " + "Tensor uc_buffer, int mc_ptr, Tensor flags, int rank, int nranks) -> Tensor[]"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("glm_moe_dsa_post_attention_hop", &tensorrt_llm::torch_ext::glm_moe_dsa::postAttentionHop); +} + +#ifdef TLLM_GLM_LL_W2_TIMING +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("glm_moe_dsa_post_attention_hop_timing() -> Tensor", + &tensorrt_llm::torch_ext::glm_moe_dsa::postAttentionHopTiming); +} +#endif diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/postAttentionHopRoutingOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsa/postAttentionHopRoutingOp.cpp new file mode 100644 index 000000000000..0d71f2b34496 --- /dev/null +++ b/cpp/tensorrt_llm/thop/glmMoeDsa/postAttentionHopRoutingOp.cpp @@ -0,0 +1,189 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/config.h" +#include "tensorrt_llm/kernels/glmMoeDsa/postAttentionHopRouting.h" + +#include +#include +#include +#include + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext::glm_moe_dsa +{ + +namespace +{ + +namespace tk = tensorrt_llm::kernels::glm_moe_dsa::hop_routing; + +// Same switch as the native kernels: PDL is on unless TRTLLM_ENABLE_PDL=0. +bool routingPdlEnabled() +{ + static bool const enabled = []() + { + char const* env = std::getenv("TRTLLM_ENABLE_PDL"); + return !(env != nullptr && std::strcmp(env, "0") == 0); + }(); + return enabled; +} + +void checkBf16(at::Tensor const& t, char const* name, int64_t lastDim) +{ + TORCH_CHECK(t.is_cuda() && t.scalar_type() == at::kBFloat16, "[glm_moe_dsa_post_attention_hop_routing] ", name, + " must be a CUDA bf16 tensor"); + TORCH_CHECK(t.is_contiguous(), "[glm_moe_dsa_post_attention_hop_routing] ", name, " must be contiguous"); + TORCH_CHECK(t.dim() >= 1 && t.size(-1) == lastDim, "[glm_moe_dsa_post_attention_hop_routing] ", name, + " must have last dimension ", lastDim); +} + +void checkWords(at::Tensor const& t, char const* name, int64_t minWords) +{ + TORCH_CHECK(t.is_cuda() && t.is_contiguous() && t.scalar_type() == at::kInt && t.numel() >= minWords, + "[glm_moe_dsa_post_attention_hop_routing] ", name, " must be a contiguous CUDA int32 tensor with at least ", + minWords, " elements"); +} + +} // namespace + +//! Pre-MoE hop with top-8 routing: the transport-input hop with router logits, plus the native +//! DeepSeek-V3 routing of every row. Returns ``[normed, residual_out, router_logits, routing]`` where ``routing`` is +//! an int32 ``[2, R, 8]`` tensor: ``routing[0]`` holds the expert ids and ``routing[1]`` the FP32 bits of the +//! bf16-rounded weights, both in descending native key order. See kernels/glmMoeDsa/postAttentionHopRouting.h. +std::vector postAttentionHopRouting(at::Tensor const& residual, at::Tensor const& gamma, double eps, + at::Tensor const& routerWeight, at::Tensor const& routingBias, at::Tensor const& ucBuffer, at::Tensor const& flags, + at::Tensor const& gather, at::Tensor const& zstate, int64_t rank, int64_t nRanks, bool earlyPoll, + int64_t routerLoad, double routedScaling, std::optional const& logitsIn, + std::optional const& stamps) +{ + checkBf16(residual, "residual", tk::kHiddenDim); + checkBf16(gamma, "gamma", tk::kHiddenDim); + TORCH_CHECK(gamma.dim() == 1, "[glm_moe_dsa_post_attention_hop_routing] gamma must be 1-D"); + checkBf16(routerWeight, "router_weight", tk::kHiddenDim); + TORCH_CHECK(routerWeight.dim() == 2 && routerWeight.size(0) == tk::kNumExperts, + "[glm_moe_dsa_post_attention_hop_routing] router_weight must have shape [", tk::kNumExperts, ", ", + tk::kHiddenDim, "]"); + checkBf16(routingBias, "routing_bias", tk::kNumExperts); + TORCH_CHECK(routingBias.dim() == 1, "[glm_moe_dsa_post_attention_hop_routing] routing_bias must be 1-D"); + int64_t const numRows = residual.numel() / tk::kHiddenDim; + TORCH_CHECK(numRows >= 1 && numRows <= tk::kMaxRows, "[glm_moe_dsa_post_attention_hop_routing] number of rows ", + numRows, " is outside the supported range [1, ", tk::kMaxRows, "]"); + TORCH_CHECK(nRanks == tk::kNumRanks, "[glm_moe_dsa_post_attention_hop_routing] only TP", tk::kNumRanks, + " is supported, got ", nRanks); + TORCH_CHECK(rank >= 0 && rank < nRanks, "[glm_moe_dsa_post_attention_hop_routing] invalid rank ", rank); + TORCH_CHECK(routerLoad >= 0 && routerLoad <= 2, + "[glm_moe_dsa_post_attention_hop_routing] router_load must be 0, 1 or 2, got ", routerLoad); + TORCH_CHECK(ucBuffer.is_cuda() && ucBuffer.is_contiguous() + && static_cast(ucBuffer.numel() * ucBuffer.element_size()) >= tk::kWorkspaceBytes, + "[glm_moe_dsa_post_attention_hop_routing] uc_buffer must be a contiguous CUDA tensor of at least ", + tk::kWorkspaceBytes, " bytes"); + checkWords(flags, "flags", 4); + checkWords(gather, "gather", tk::kGatherWords); + checkWords(zstate, "zstate", tk::kStateWords); + auto const device = residual.get_device(); + TORCH_CHECK(gamma.get_device() == device && routerWeight.get_device() == device + && routingBias.get_device() == device && ucBuffer.get_device() == device && flags.get_device() == device + && gather.get_device() == device && zstate.get_device() == device, + "[glm_moe_dsa_post_attention_hop_routing] all tensors must be on the same device"); + if (logitsIn.has_value()) + { + auto const& l = logitsIn.value(); + TORCH_CHECK(l.is_cuda() && l.get_device() == device && l.scalar_type() == at::kFloat && l.is_contiguous() + && l.numel() == numRows * tk::kNumExperts, + "[glm_moe_dsa_post_attention_hop_routing] logits_in must be a contiguous CUDA fp32 tensor of shape [", + numRows, ", ", tk::kNumExperts, "]"); + TORCH_CHECK(routerLoad == 1, "[glm_moe_dsa_post_attention_hop_routing] logits_in needs router_load=1"); + } + if (stamps.has_value()) + { + auto const& s = stamps.value(); + TORCH_CHECK(tk::stampsCompiled(), + "[glm_moe_dsa_post_attention_hop_routing] stamps need a TLLM_GLM_MOE_DSA_HOP_ROUTING_STAMPS build"); + TORCH_CHECK(s.is_cuda() && s.get_device() == device && s.scalar_type() == at::kLong && s.is_contiguous() + && s.numel() >= tk::kMaxGridCtas * tk::kNumStamps, + "[glm_moe_dsa_post_attention_hop_routing] stamps must be a contiguous CUDA int64 tensor with at least ", + tk::kMaxGridCtas * tk::kNumStamps, " elements"); + } + + c10::cuda::CUDAGuard const deviceGuard(residual.device()); + at::Tensor normOut = at::empty_like(residual); + at::Tensor residualOut = at::empty_like(residual); + at::Tensor routerLogits = at::empty({numRows, tk::kNumExperts}, residual.options().dtype(at::kFloat)); + at::Tensor routing = at::empty({2, numRows, tk::kTopK}, residual.options().dtype(at::kInt)); + + tk::PostAttentionHopRoutingParams params; + params.residualIn = residual.const_data_ptr(); + params.gamma = gamma.const_data_ptr(); + params.routerWeight = routerWeight.const_data_ptr(); + params.routingBias = routingBias.const_data_ptr(); + params.logitsIn = logitsIn.has_value() ? logitsIn->const_data_ptr() : nullptr; + params.normOut = normOut.mutable_data_ptr(); + params.residualOut = residualOut.mutable_data_ptr(); + params.routerLogits = routerLogits.mutable_data_ptr(); + params.topkIds = routing.mutable_data_ptr(); + params.topkWeights = reinterpret_cast(routing.mutable_data_ptr() + numRows * tk::kTopK); + params.ucBuffer = ucBuffer.mutable_data_ptr(); + params.flags = reinterpret_cast(flags.mutable_data_ptr()); + params.gather = reinterpret_cast(gather.mutable_data_ptr()); + params.zstate = reinterpret_cast(zstate.mutable_data_ptr()); + params.stamps + = stamps.has_value() ? reinterpret_cast(stamps->mutable_data_ptr()) : nullptr; + params.numRows = static_cast(numRows); + params.rank = static_cast(rank); + params.eps = static_cast(eps); + params.routedScaling = static_cast(routedScaling); + params.enablePdl = routingPdlEnabled(); + params.earlyPoll = earlyPoll; + params.routerLoad = static_cast(routerLoad); + params.stream = at::cuda::getCurrentCUDAStream(device); + + cudaError_t const status = tk::invokePostAttentionHopRouting(params); + TORCH_CHECK(status == cudaSuccess, + "[glm_moe_dsa_post_attention_hop_routing] kernel launch failed: ", cudaGetErrorString(status)); + return {normOut, residualOut, routerLogits, routing}; +} + +} // namespace torch_ext::glm_moe_dsa + +TRTLLM_NAMESPACE_END + +// A diagnostic build (stamps) registers under its own name, so it can be loaded next to a build that has the op. +#if defined(TLLM_GLM_MOE_DSA_HOP_ROUTING_STAMPS) +#define GLM_MOE_DSA_HOP_ROUTING_OP_NAME "glm_moe_dsa_post_attention_hop_routing_diag" +#else +#define GLM_MOE_DSA_HOP_ROUTING_OP_NAME "glm_moe_dsa_post_attention_hop_routing" +#endif + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def(GLM_MOE_DSA_HOP_ROUTING_OP_NAME + "(Tensor residual, Tensor gamma, float eps, Tensor router_weight, Tensor routing_bias, " + "Tensor uc_buffer, Tensor flags, Tensor gather, Tensor zstate, int rank, int nranks, bool early_poll, " + "int router_load, float routed_scaling, Tensor? logits_in=None, Tensor? stamps=None) -> Tensor[]"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl(GLM_MOE_DSA_HOP_ROUTING_OP_NAME, &tensorrt_llm::torch_ext::glm_moe_dsa::postAttentionHopRouting); +} diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/ropeKvAppendOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsa/ropeKvAppendOp.cpp new file mode 100644 index 000000000000..4411b6f88af8 --- /dev/null +++ b/cpp/tensorrt_llm/thop/glmMoeDsa/ropeKvAppendOp.cpp @@ -0,0 +1,330 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/envUtils.h" +#include "tensorrt_llm/common/quantization.h" +#include "tensorrt_llm/kernels/glmMoeDsa/ropeKvAppend.h" +#include "tensorrt_llm/kernels/kvCacheUtils.h" + +#include +#include + +#include +#include +#include +#include + +namespace tk = tensorrt_llm::kernels; +namespace tc = tensorrt_llm::common; +namespace rope_kv_append = tensorrt_llm::kernels::glm_moe_dsa::rope_kv_append; + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext::glm_moe_dsa +{ + +namespace +{ + +void checkCuda(torch::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda(), name, " must be a CUDA tensor"); +} + +//! Paged MLA latent cache of one layer: the KVBlockArray that buildPagedKvCacheBuffers (thop/attentionOp.cpp) +//! builds for an FP8 MLA cache (kvFactor 1, one KV head, 8-bit elements, beam width 1, no FP4 scale pool). +tk::KVBlockArray buildFp8MlaKvBlockArray(torch::Tensor const& kvCacheBlockOffsets, + torch::Tensor const& hostKvCachePoolPointers, torch::Tensor const& hostKvCachePoolMapping, int64_t layerIdx, + int32_t batchSize, int64_t tokensPerBlock, int64_t attentionWindowSize, int64_t seqOffset) +{ + TORCH_CHECK(hostKvCachePoolMapping.device().is_cpu() && hostKvCachePoolMapping.dim() == 2 + && hostKvCachePoolMapping.size(1) >= 2 && hostKvCachePoolMapping.scalar_type() == torch::kInt32, + "host_kv_cache_pool_mapping must be a CPU int32 [layers, >=2] tensor"); + TORCH_CHECK(layerIdx >= 0 && layerIdx < hostKvCachePoolMapping.size(0), "layer index out of the pool mapping"); + auto const* mapping = hostKvCachePoolMapping.data_ptr(); + int32_t const poolIndex = mapping[layerIdx * hostKvCachePoolMapping.stride(0)]; + int32_t const layerIdxInCachePool + = mapping[layerIdx * hostKvCachePoolMapping.stride(0) + hostKvCachePoolMapping.stride(1)]; + + TORCH_CHECK(hostKvCachePoolPointers.device().is_cpu() && hostKvCachePoolPointers.dim() == 2 + && hostKvCachePoolPointers.scalar_type() == torch::kInt64 && poolIndex >= 0 + && poolIndex < hostKvCachePoolPointers.size(0) && hostKvCachePoolPointers.size(1) >= 2, + "host_kv_cache_pool_pointers must be a CPU int64 [pools, 2] tensor"); + TORCH_CHECK(kvCacheBlockOffsets.is_cuda() && kvCacheBlockOffsets.scalar_type() == torch::kInt32 + && kvCacheBlockOffsets.dim() >= 2 && poolIndex < kvCacheBlockOffsets.size(0) + && seqOffset < kvCacheBlockOffsets.size(1), + "kv_cache_block_offsets must be CUDA int32 [pools, seqs, 2, max_blocks]"); + + int64_t constexpr kCacheElemBits = 8; + int64_t const blockSize = tokensPerBlock * /*kv_head_num=*/1 * rope_kv_append::kLatentDim; + int64_t const bytesPerBlock = blockSize * kCacheElemBits / CHAR_BIT; + int64_t const intraPoolOffset = layerIdxInCachePool * /*kvFactor=*/1 * bytesPerBlock; + auto const sizePerToken = static_cast(rope_kv_append::kLatentDim * kCacheElemBits / 8); + auto const* pointers = hostKvCachePoolPointers.data_ptr(); + auto const rowOffset = poolIndex * hostKvCachePoolPointers.stride(0); + void* primaryPoolPtr = reinterpret_cast(pointers[rowOffset]) + intraPoolOffset; + void* secondaryPoolPtr + = reinterpret_cast(pointers[rowOffset + hostKvCachePoolPointers.stride(1)]) + intraPoolOffset; + auto* blockOffsets + = reinterpret_cast(static_cast(kvCacheBlockOffsets.data_ptr()) + + poolIndex * kvCacheBlockOffsets.stride(0) + seqOffset * kvCacheBlockOffsets.stride(1)); + auto const maxBlocksPerSequence = static_cast(kvCacheBlockOffsets.size(-1)); + auto const window = static_cast(attentionWindowSize); + return tk::KVBlockArray(batchSize, maxBlocksPerSequence, static_cast(tokensPerBlock), sizePerToken, window, + window, /*sinkTokenLen=*/0, /*canUseOneMoreBlock=*/false, primaryPoolPtr, secondaryPoolPtr, blockOffsets); +} + +} // namespace + +//! Generation MLA RoPE + FP8 query quantization + FP8 latent-cache append (the FP8-cache, absorption-mode, +//! non-helix subset of trtllm::mla_rope_generation), optionally fused with the DSA indexer K-cache scatter of +//! trtllm::indexer_k_cache_scatter_op. Argument meanings follow those two native ops. +void ropeKvAppend(torch::Tensor const& fused_q, torch::Tensor const& q_pe, torch::Tensor const& latent_cache, + torch::Tensor const& rotary_cos_sin, torch::Tensor const& cu_q_seqlens, torch::Tensor const& cu_kv_seqlens, + torch::Tensor const& fmha_scheduler_counter, torch::Tensor const& mla_bmm1_scale, + torch::Tensor const& mla_bmm2_scale, torch::Tensor const& quant_q_buffer, torch::Tensor const& sequence_length, + torch::Tensor const& host_context_lengths, int64_t num_contexts, torch::Tensor const& kv_cache_block_offsets, + torch::Tensor const& host_kv_cache_pool_pointers, torch::Tensor const& host_kv_cache_pool_mapping, + std::optional const& kv_cache_scale_orig_quant, int64_t layer_idx, int64_t num_heads, + int64_t tokens_per_block, int64_t attention_window_size, int64_t quant_mode, double q_scaling, + int64_t qk_nope_head_dim, int64_t qk_rope_head_dim, int64_t kv_lora_rank, + std::optional const& idx_k_fp8, std::optional const& idx_k_scale, + std::optional const& idx_k_cache, std::optional const& idx_slot_mapping_fp8, + std::optional const& idx_slot_mapping_scale, int64_t idx_num_tokens) +{ + TORCH_CHECK(kv_lora_rank == rope_kv_append::kKvLoraRank && qk_rope_head_dim == rope_kv_append::kRopeDim, + "glm_moe_dsa_rope_kv_append: supports kv_lora_rank 512 and qk_rope_head_dim 64 only"); + auto const kvQuantMode = tc::QuantMode(static_cast(quant_mode)); + TORCH_CHECK(kvQuantMode.hasFp8KvCache() && !kvQuantMode.hasFp4KvCache(), + "glm_moe_dsa_rope_kv_append: needs an FP8 MLA cache"); + checkCuda(fused_q, "fused_q"); + checkCuda(q_pe, "q_pe"); + checkCuda(latent_cache, "latent_cache"); + TORCH_CHECK(fused_q.scalar_type() == torch::kBFloat16 && q_pe.scalar_type() == torch::kBFloat16 + && latent_cache.scalar_type() == torch::kBFloat16, + "glm_moe_dsa_rope_kv_append: needs BF16 fused_q, q_pe and latent_cache"); + TORCH_CHECK(rotary_cos_sin.scalar_type() == torch::kFloat32 && rotary_cos_sin.is_contiguous(), + "rotary_cos_sin must be packed float32"); + int64_t const cosSinRows = rotary_cos_sin.numel() / (2 * rope_kv_append::kRopeDim); + TORCH_CHECK(cosSinRows > 0 && cosSinRows <= INT32_MAX, "rotary_cos_sin has no complete position row"); + TORCH_CHECK(q_pe.dim() == 3 && q_pe.stride(2) == 1, "q_pe must be [tokens, heads, rope] with unit inner stride"); + TORCH_CHECK(latent_cache.dim() == 2 && latent_cache.size(1) == rope_kv_append::kLatentDim + && latent_cache.stride(1) == 1 && latent_cache.stride(0) == rope_kv_append::kLatentDim, + "latent_cache must be packed [tokens, 576]"); + TORCH_CHECK(fused_q.is_contiguous() && quant_q_buffer.is_contiguous(), "fused_q and quant_q_buffer must be packed"); + TORCH_CHECK(host_kv_cache_pool_mapping.dim() >= 1, "host_kv_cache_pool_mapping is required"); + + int32_t const numTokens = static_cast(latent_cache.size(0)); + int32_t const numSeqs = static_cast(host_context_lengths.size(0)); + int32_t const numGenerations = numSeqs - static_cast(num_contexts); + TORCH_CHECK(numGenerations > 0, "glm_moe_dsa_rope_kv_append: needs at least one generation request"); + TORCH_CHECK(numTokens % numGenerations == 0, "MLA can only support input sequences with the same sequence length."); + TORCH_CHECK(fused_q.numel() >= static_cast(numTokens) * num_heads * rope_kv_append::kLatentDim, + "fused_q is smaller than [tokens, heads, 576]"); + TORCH_CHECK(quant_q_buffer.numel() * quant_q_buffer.element_size() + >= static_cast(numTokens) * num_heads * rope_kv_append::kLatentDim, + "quant_q_buffer is smaller than [tokens, heads, 576] bytes"); + TORCH_CHECK(mla_bmm1_scale.numel() >= 2 && mla_bmm2_scale.numel() >= 1, "bmm scale buffers too small"); + TORCH_CHECK(cu_q_seqlens.numel() >= numGenerations + 1 && cu_kv_seqlens.numel() >= numGenerations + 1, + "cu_seqlens buffers too small"); + + auto stream = at::cuda::getCurrentCUDAStream(latent_cache.get_device()); + rope_kv_append::RopeAppendParams params{}; + params.fusedQ = reinterpret_cast<__nv_bfloat16 const*>(fused_q.data_ptr()); + params.qPe = reinterpret_cast<__nv_bfloat16 const*>(q_pe.data_ptr()); + params.qPeLd = static_cast(q_pe.stride(1)); + params.qPeStride = static_cast(q_pe.stride(0)); + params.quantQ = reinterpret_cast(quant_q_buffer.data_ptr()); + params.quantScaleQ = nullptr; + params.latent = reinterpret_cast<__nv_bfloat16 const*>(latent_cache.data_ptr()); + params.cosSin = reinterpret_cast(rotary_cos_sin.data_ptr()); + params.cosSinRows = static_cast(cosSinRows); + params.kvLens = sequence_length.slice(0, num_contexts).data_ptr(); + params.quantScaleKv + = kv_cache_scale_orig_quant.has_value() ? kv_cache_scale_orig_quant->data_ptr() : nullptr; + params.kvCache = buildFp8MlaKvBlockArray(kv_cache_block_offsets, host_kv_cache_pool_pointers, + host_kv_cache_pool_mapping, layer_idx, numGenerations, tokens_per_block, attention_window_size, num_contexts); + params.cuQSeqlens = cu_q_seqlens.data_ptr(); + params.cuKvSeqlens = cu_kv_seqlens.data_ptr(); + params.fmhaTileCounter = reinterpret_cast(fmha_scheduler_counter.data_ptr()); + params.bmm1Scale = mla_bmm1_scale.data_ptr(); + params.bmm2Scale = mla_bmm2_scale.data_ptr(); + params.dequantScaleQ = nullptr; + params.dequantScaleKv = nullptr; + params.quantScaleO = nullptr; + // Same expression as the native op so the folded FMHA scales match bit for bit. + params.hostBmm1Scale = 1.F / (q_scaling * std::sqrt(static_cast(qk_nope_head_dim + qk_rope_head_dim))); + params.numTokens = numTokens; + params.seqLen = numTokens / numGenerations; + params.numHeads = static_cast(num_heads); + + if (idx_num_tokens > 0) + { + TORCH_CHECK(idx_k_fp8.has_value() && idx_k_scale.has_value() && idx_k_cache.has_value() + && idx_slot_mapping_fp8.has_value() && idx_slot_mapping_scale.has_value(), + "indexer scatter needs k_fp8, k_scale, k_cache and both slot mappings"); + auto const& kFp8 = idx_k_fp8.value(); + auto const& kScale = idx_k_scale.value(); + auto const& kCache = idx_k_cache.value(); + auto const& slotFp8 = idx_slot_mapping_fp8.value(); + auto const& slotScale = idx_slot_mapping_scale.value(); + checkCuda(kFp8, "k_fp8"); + checkCuda(kScale, "k_scale"); + checkCuda(kCache, "k_cache"); + TORCH_CHECK(kFp8.dim() == 2 && kFp8.element_size() == 1 && kFp8.is_contiguous(), "k_fp8 must be packed bytes"); + TORCH_CHECK(kScale.dim() == 2 && kScale.element_size() == 4 && kScale.is_contiguous() + && kScale.size(1) * kScale.element_size() == 4, + "k_scale must be one packed 4-byte word per token"); + TORCH_CHECK(kCache.dim() == 4 && kCache.size(2) == 1, "k_cache must be [blocks, block_size, 1, bytes]"); + TORCH_CHECK(slotFp8.scalar_type() == torch::kInt64 && slotScale.scalar_type() == torch::kInt64 + && slotFp8.is_contiguous() && slotScale.is_contiguous(), + "slot mappings must be packed int64"); + int32_t const headBytes = static_cast(kFp8.size(1)); + TORCH_CHECK(headBytes == 128 || headBytes == 64, "indexer k head bytes must be 128 or 64"); + TORCH_CHECK(kFp8.size(0) >= idx_num_tokens && kScale.size(0) >= idx_num_tokens + && slotFp8.numel() >= idx_num_tokens && slotScale.numel() >= idx_num_tokens, + "indexer inputs are shorter than idx_num_tokens"); + params.idxK = reinterpret_cast(kFp8.data_ptr()); + params.idxScale = reinterpret_cast(kScale.data_ptr()); + params.idxCache = reinterpret_cast(kCache.data_ptr()); + params.idxSlotData = slotFp8.data_ptr(); + params.idxSlotScale = slotScale.data_ptr(); + params.idxNumTokens = static_cast(idx_num_tokens); + params.idxHeadBytes = headBytes; + for (int32_t d = 0; d < 4; ++d) + { + params.idxDims[d] = static_cast(kCache.size(d)); + params.idxStrides[d] = kCache.stride(d); + } + } + + rope_kv_append::invokeRopeAppendScatter(params, tc::getEnvEnablePDL(), stream); +} + +//! Local top-k token ids -> global KV-pool indices (values of trtllm::convert_req_index_to_global and its grouped +//! variant), optionally fused with the native top-k reuse copy. Returns [group, rows, topk] when layer_ids is given, +//! otherwise [rows, topk]. +torch::Tensor topkIndexBridge(torch::Tensor const& topk, int64_t num_rows, torch::Tensor const& req_idx, + int64_t req_idx_bias, torch::Tensor const& block_table, int64_t block_size, int64_t stride_factor, + std::optional const& layer_ids, int64_t layer_id, std::optional const& copy_dst, + int64_t copy_row_start, int64_t copy_rows, int64_t copy_mode, int64_t next_n, + std::optional const& num_accepted) +{ + checkCuda(topk, "topk"); + checkCuda(req_idx, "req_idx"); + checkCuda(block_table, "block_table"); + TORCH_CHECK(topk.scalar_type() == torch::kInt32 && topk.dim() == 2 && topk.stride(1) == 1, + "topk must be int32 [rows, topk] with unit column stride"); + TORCH_CHECK(req_idx.scalar_type() == torch::kInt32 && req_idx.dim() == 1 && req_idx.is_contiguous(), + "req_idx must be packed int32"); + TORCH_CHECK(block_table.scalar_type() == torch::kInt32 && block_table.dim() == 2, "block_table must be int32 2-D"); + TORCH_CHECK(copy_mode == 0 || copy_mode == 1, "copy_mode must be 0 or 1"); + int64_t const numTopk = topk.size(1); + TORCH_CHECK(req_idx.size(0) >= num_rows, "req_idx shorter than num_rows"); + TORCH_CHECK(copy_mode == 1 || topk.size(0) >= num_rows, "topk shorter than num_rows"); + TORCH_CHECK(copy_mode == 0 || topk.size(0) >= copy_rows * next_n, "topk shorter than copy_rows * next_n"); + TORCH_CHECK(copy_mode == 1 || copy_rows <= num_rows, "copy_rows exceeds num_rows"); + + int32_t groupSize = 1; + int32_t const* layerIdsPtr = nullptr; + if (layer_ids.has_value()) + { + TORCH_CHECK(layer_ids->scalar_type() == torch::kInt32 && layer_ids->dim() == 1 && layer_ids->is_contiguous(), + "layer_ids must be packed int32"); + groupSize = static_cast(layer_ids->size(0)); + layerIdsPtr = layer_ids->data_ptr(); + } + torch::Tensor out = layer_ids.has_value() ? torch::empty({groupSize, num_rows, numTopk}, topk.options()) + : torch::empty({num_rows, numTopk}, topk.options()); + + rope_kv_append::TopkBridgeParams params{}; + params.topk = topk.data_ptr(); + params.topkStride0 = topk.stride(0); + params.numRows = static_cast(num_rows); + params.numTopk = static_cast(numTopk); + params.reqIdx = req_idx.data_ptr(); + params.reqIdxBias = static_cast(req_idx_bias); + params.blockTable = block_table.data_ptr(); + params.btStride0 = block_table.stride(0); + params.btStride1 = block_table.stride(1); + params.maxBlocksPerReq = static_cast(block_table.size(1)); + params.blockSize = static_cast(block_size); + params.strideFactor = static_cast(stride_factor); + params.layerIds = layerIdsPtr; + params.layerId = static_cast(layer_id); + params.groupSize = groupSize; + params.out = out.data_ptr(); + params.outStrideG = layer_ids.has_value() ? out.stride(0) : 0; + params.outStride0 = layer_ids.has_value() ? out.stride(1) : out.stride(0); + if (copy_dst.has_value() && copy_rows > 0) + { + auto const& dst = copy_dst.value(); + checkCuda(dst, "copy_dst"); + TORCH_CHECK(dst.scalar_type() == torch::kInt32 && dst.dim() == 2 && dst.stride(1) == 1 && dst.size(1) >= numTopk + && dst.size(0) >= copy_row_start + copy_rows, + "copy_dst must be int32 [>= start + rows, >= topk] with unit column stride"); + params.copyDst = dst.data_ptr(); + params.copyStride0 = dst.stride(0); + params.copyRows = static_cast(copy_rows); + params.copyRowStart = static_cast(copy_row_start); + params.copyMode = static_cast(copy_mode); + params.nextN = static_cast(next_n); + if (copy_mode == 1 && num_accepted.has_value()) + { + auto const& acc = num_accepted.value(); + checkCuda(acc, "num_accepted"); + TORCH_CHECK((acc.scalar_type() == torch::kInt32 || acc.scalar_type() == torch::kInt64) && acc.dim() == 1 + && acc.is_contiguous() && acc.size(0) >= copy_rows, + "num_accepted must be packed int32/int64 with at least copy_rows entries"); + params.numAccepted = acc.data_ptr(); + params.numAcceptedIs64 = acc.scalar_type() == torch::kInt64; + } + } + + auto stream = at::cuda::getCurrentCUDAStream(topk.get_device()); + rope_kv_append::invokeTopkBridge(params, tc::getEnvEnablePDL(), stream); + return out; +} + +} // namespace torch_ext::glm_moe_dsa + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "glm_moe_dsa_rope_kv_append(Tensor fused_q, Tensor q_pe, Tensor latent_cache, Tensor rotary_cos_sin, " + "Tensor(a!) cu_q_seqlens, Tensor(b!) cu_kv_seqlens, Tensor(c!) fmha_scheduler_counter, " + "Tensor(d!) mla_bmm1_scale, Tensor(e!) mla_bmm2_scale, Tensor(f!) quant_q_buffer, Tensor sequence_length, " + "Tensor host_context_lengths, int num_contexts, Tensor kv_cache_block_offsets, " + "Tensor host_kv_cache_pool_pointers, Tensor host_kv_cache_pool_mapping, Tensor? kv_cache_scale_orig_quant, " + "int layer_idx, int num_heads, int tokens_per_block, int attention_window_size, int quant_mode, " + "float q_scaling, int qk_nope_head_dim, int qk_rope_head_dim, int kv_lora_rank, Tensor? idx_k_fp8, " + "Tensor? idx_k_scale, Tensor(g!)? idx_k_cache, Tensor? idx_slot_mapping_fp8, Tensor? idx_slot_mapping_scale, " + "int idx_num_tokens) -> ()"); + m.def( + "glm_moe_dsa_topk_index_bridge(Tensor topk, int num_rows, Tensor req_idx, int req_idx_bias, Tensor " + "block_table, " + "int block_size, int stride_factor, Tensor? layer_ids, int layer_id, Tensor(a!)? copy_dst, " + "int copy_row_start, int copy_rows, int copy_mode, int next_n, Tensor? num_accepted) -> Tensor"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("glm_moe_dsa_rope_kv_append", &tensorrt_llm::torch_ext::glm_moe_dsa::ropeKvAppend); + m.impl("glm_moe_dsa_topk_index_bridge", &tensorrt_llm::torch_ext::glm_moe_dsa::topkIndexBridge); +} diff --git a/cpp/tensorrt_llm/thop/glmMoeDsa/sparseMlaDecodeOp.cpp b/cpp/tensorrt_llm/thop/glmMoeDsa/sparseMlaDecodeOp.cpp new file mode 100644 index 000000000000..a7dba3610644 --- /dev/null +++ b/cpp/tensorrt_llm/thop/glmMoeDsa/sparseMlaDecodeOp.cpp @@ -0,0 +1,121 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/glmMoeDsa/sparseMlaDecode.h" + +#include +#include +#include +#include + +#include + +namespace tk = tensorrt_llm::kernels::glm_moe_dsa::sparse_mla_decode; + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext::glm_moe_dsa +{ + +namespace +{ + +bool isFp8ByteTensor(torch::Tensor const& t) +{ + return t.scalar_type() == torch::kUInt8 || t.scalar_type() == torch::kFloat8_e4m3fn; +} + +} // namespace + +// Sparse MLA generation attention; see sparseMlaDecode.h. +void sparseMlaDecode(torch::Tensor const& q, int64_t kvPoolPtr, torch::Tensor const& indices, + torch::Tensor const& bmm1Scale, torch::Tensor const& bmm2Scale, torch::Tensor const& seqLensKv, int64_t seqLenQ, + torch::Tensor& out) +{ + TORCH_CHECK(q.is_cuda() && indices.is_cuda() && bmm1Scale.is_cuda() && bmm2Scale.is_cuda() && out.is_cuda(), + "glm_moe_dsa_sparse_mla_decode expects CUDA tensors"); + TORCH_CHECK(isFp8ByteTensor(q), "q must hold e4m3 bytes (uint8 or float8_e4m3fn)"); + TORCH_CHECK(q.is_contiguous(), "q must be contiguous"); + int64_t const numRows = q.size(0); + TORCH_CHECK(q.numel() == numRows * tk::kNumHeads * tk::kHeadDimQk, "q must be [rows, ", tk::kNumHeads, ", ", + tk::kHeadDimQk, "]"); + TORCH_CHECK(indices.scalar_type() == torch::kInt32 && indices.dim() == 2, "indices must be a 2-D int32 tensor"); + TORCH_CHECK(indices.size(0) == numRows, "indices must have one row per query row"); + TORCH_CHECK(indices.size(1) == tk::kTopK, "indices must hold ", tk::kTopK, " entries per row"); + TORCH_CHECK(indices.stride(1) == 1, "indices rows must be contiguous"); + TORCH_CHECK(bmm1Scale.scalar_type() == torch::kFloat32 && bmm1Scale.numel() >= 2, "bmm1_scale must be float32[2]"); + TORCH_CHECK(bmm2Scale.scalar_type() == torch::kFloat32 && bmm2Scale.numel() >= 1, "bmm2_scale must be float32[1]"); + TORCH_CHECK(out.scalar_type() == torch::kBFloat16 && out.is_contiguous(), "out must be contiguous bfloat16"); + TORCH_CHECK(out.numel() == numRows * tk::kNumHeads * tk::kHeadDimV, "out must be [rows, ", + tk::kNumHeads * tk::kHeadDimV, "]"); + TORCH_CHECK(kvPoolPtr != 0, "kv_pool_ptr must be a valid device address"); + + tk::SparseMlaDecodeParams params{}; + params.q = q.data_ptr(); + params.kvPool = reinterpret_cast(kvPoolPtr); + params.indices = indices.data_ptr(); + params.indicesRowStride = indices.stride(0); + params.bmm1Scale = bmm1Scale.data_ptr(); + params.bmm2Scale = bmm2Scale.data_ptr(); + TORCH_CHECK(seqLensKv.is_cuda() && seqLensKv.scalar_type() == torch::kInt32 && seqLensKv.is_contiguous(), + "seq_lens_kv must be a contiguous int32 CUDA tensor"); + TORCH_CHECK(seqLenQ > 0 && numRows % seqLenQ == 0, "rows must be a multiple of seq_len_q"); + TORCH_CHECK(seqLensKv.numel() >= numRows / seqLenQ, "seq_lens_kv must cover every sequence"); + params.seqLensKv = seqLensKv.data_ptr(); + params.seqLenQ = static_cast(seqLenQ); + params.out = reinterpret_cast<__nv_bfloat16*>(out.data_ptr()); + params.numRows = static_cast(numRows); + + auto const stream = at::cuda::getCurrentCUDAStream(q.get_device()); + tk::invokeSparseMlaDecode(params, stream); +} + +#ifdef GLM_MOE_DSA_SPARSE_MLA_DECODE_TIMING +torch::Tensor sparseMlaDecodeDebugTimeline() +{ + auto out = torch::empty({tk::kTimelineMaxCtas, tk::kTimelineSlots}, torch::dtype(torch::kInt64)); + tk::copySparseMlaDecodeTimeline(reinterpret_cast(out.data_ptr())); + return out; +} +#endif + +bool sparseMlaDecodeSupported(int64_t device) +{ + return tk::isSparseMlaDecodeDeviceSupported(static_cast(device)); +} + +} // namespace torch_ext::glm_moe_dsa + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "glm_moe_dsa_sparse_mla_decode(Tensor q, int kv_pool_ptr, Tensor indices, Tensor bmm1_scale, " + "Tensor bmm2_scale, Tensor seq_lens_kv, int seq_len_q, Tensor(a!) out) -> ()"); + m.def("glm_moe_dsa_sparse_mla_decode_supported(int device) -> bool", + &tensorrt_llm::torch_ext::glm_moe_dsa::sparseMlaDecodeSupported); +#ifdef GLM_MOE_DSA_SPARSE_MLA_DECODE_TIMING + m.def("glm_moe_dsa_sparse_mla_decode_debug_timeline() -> Tensor", + &tensorrt_llm::torch_ext::glm_moe_dsa::sparseMlaDecodeDebugTimeline); +#endif +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("glm_moe_dsa_sparse_mla_decode", &tensorrt_llm::torch_ext::glm_moe_dsa::sparseMlaDecode); +} diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md index 3956bcdf6725..ff37dd2861f8 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md @@ -132,6 +132,15 @@ catalog// entries more than one target calls does. It has one subdirectory per `ops/` category, plus `common/` for helpers that more than one category imports. +GLM-5.2's low-latency target, `glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8`, serves GLM-5.2 +NVFP4 decode on sm_103 with the attention heads and the experts tensor-parallel over eight +ranks; its routing module admits only that layout. Its fast paths are packages under the +target's own `kernels/`, one per fast path, with their C++ kernels and ops in +`cpp/tensorrt_llm/kernels/glmMoeDsa` and `cpp/tensorrt_llm/thop/glmMoeDsa`; the target's +`dispatch.py` selects them and lists the knobs, and its layers run the composition of +`layers.py`. The target is gated end to end by +`tests/integration/defs/accuracy/modeling_v2/test_glm_moe_dsa.py` (GSM8K under `require`). + An entry, shared or owned, is one `OpWrapper` subclass: `__call__` is the op call a target makes, `reference` is what it is supposed to compute in plain torch, `is_valid` names the inputs it accepts and answers wrongly, and `CELLS` diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py index 16d312d3e7b1..f08e3a6be5b8 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py @@ -32,6 +32,7 @@ NoTarget, modeling_v2_resolve, modeling_v2_resolve_on_all_ranks, + resolved_mode, step_contract_enabled, ) @@ -40,6 +41,7 @@ "ModelingV2Mode", "NoTarget", "modeling_v2_resolve", + "resolved_mode", "modeling_v2_resolve_on_all_ranks", "step_contract_enabled", ] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py index 3814683f7fd7..6eebdc4ebb75 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py @@ -122,6 +122,20 @@ class NoTarget(Exception): """ +# The mode each target class was last resolved under (the loader builds the class right after resolving it). +_RESOLVED_MODES: Dict[type, ModelingV2Mode] = {} + + +def resolved_mode(target: type) -> ModelingV2Mode: + """The mode ``modeling_v2_resolve`` returned ``target`` under; ``auto`` for a class built without it. + + A target reads it at construction: under ``require`` its in-target dispatch + raises where it would otherwise fall back to a built-in path inside the + shapes it claims. + """ + return _RESOLVED_MODES.get(target, ModelingV2Mode.AUTO) + + def tp_reaches_across_nodes(mapping) -> bool: """Whether an attention-TP allreduce over ``mapping`` can reach every rank through the workspaces the default targets build. @@ -197,7 +211,7 @@ def modeling_v2_resolve(config: "ModelConfig", llm_args: "TorchLlmArgs") -> Opti if route is None: known = ", ".join(sorted(registered_routes())) or "(none)" raise NoTarget(f"no route for this architecture; routed architectures: {known}") - return route(config, llm_args, _device_compute_capability()) + target = route(config, llm_args, _device_compute_capability()) except NoTarget as why: if mode is ModelingV2Mode.REQUIRE: raise ValueError(f"modeling_v2='require', but no target fits {arch!r}: {why}") from None @@ -205,6 +219,8 @@ def modeling_v2_resolve(config: "ModelConfig", llm_args: "TorchLlmArgs") -> Opti f"modeling_v2: no target fits {arch!r} ({why}); using the built-in implementation" ) return None + _RESOLVED_MODES[target] = mode + return target def _qualified(target: Optional[type]) -> Optional[str]: diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/README.md b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/README.md new file mode 100644 index 000000000000..120b5106727e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/README.md @@ -0,0 +1,14 @@ +# `glm_5_2_nvfp4__sm_103__tp8` + +GLM-5.2 NVFP4 on sm_103 for low-latency decode, with the attention heads and the experts +tensor-parallel over eight ranks and no attention data parallelism (tp8), and one-model +MTP-Eagle as a linear chain of five draft tokens. `../routing.py` routes only that layout here. + +The fast paths were designed and tuned for this layout, and the target runs every one that has an +engine hook (`TP8_PATHS` in `modeling.py`). Their code lives in this directory: the plan +(`dispatch.py`), the attention adoption (`attention.py`, `attention_forward.py`), the layer +composition (`layers.py`), the speculative worker (`spec_worker.py`), the step-loop hooks +(`runtime.py`), and one kernel package per fast path under `kernels/`. + +Other sm_103 layouts are served by other targets or by the built-in model: tep4 by +`../glm_5_2_nvfp4__sm_103__tep4`, and everything else by the built-in implementation. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/__init__.py new file mode 100644 index 000000000000..87156adf61ff --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GLM-5.2-NVFP4 / sm_103 / tp8.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/attention.py new file mode 100644 index 000000000000..a01b42b429e6 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/attention.py @@ -0,0 +1,569 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The attention modules of the glm_moe_dsa sm_103 tp8 target: the MLA, its DSA hooks, its DSA backend and +its indexer. + +``GlmMLA.adopt`` re-classes a layer's built-in attention modules in place -- the +``DeepseekV32Attention`` as ``GlmMLA``, its ``DSAHooks`` as ``GlmDsaHooks``, its +``DSATrtllmAttention`` as ``GlmDsaBackend`` and its ``Indexer`` as ``GlmIndexer`` -- and attaches +the attention fast paths the installed plan enables: the attention output tail and, with +``sparse_mla_decode``, that FMHA library ahead of the registry's on the layer's attention +backends. The backends build their FMHA libraries later, in ``create_weights``, so the library is +created when and where the built-in model creates its own. The classes add methods and +class-level defaults only, so the module tree and the weights stay the built-in model's. + +Each forward override runs the target's decode composition (``attention_forward``) when the +step is a decode step (``dispatch.Phase.DECODE``) and the built-in method otherwise. +""" + +from typing import Optional, Tuple + +import torch + +import tensorrt_llm +import tensorrt_llm.bindings +from tensorrt_llm._torch.attention.attention import _helix_cp_allgather_input +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + AttentionInputType, + AttentionMetadata, +) +from tensorrt_llm._torch.attention.backends.sparse.dsa.backend import DSATrtllmAttention +from tensorrt_llm._torch.attention.backends.sparse.dsa.indexer import Indexer +from tensorrt_llm._torch.attention.backends.sparse.dsa.metadata import DSAtrtllmAttentionMetadata +from tensorrt_llm._torch.attention.backends.sparse.dsa.module import DSAHooks +from tensorrt_llm._torch.distributed import AllReduceParams +from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV32Attention +from tensorrt_llm._torch.utils import Fp4QuantizedTensor + +from . import attention_forward, dispatch +from .kernels.indexer_topk import region + + +def reclass(module, builtin: type, cls: type) -> None: + """Re-class ``module``, which must be exactly the built-in class ``builtin``, as ``cls``. + + Raises: + TypeError: if ``module`` is not exactly ``builtin`` (a subclass of it is refused too). + """ + if type(module) is not builtin: + raise TypeError( + f"{cls.__name__} adopts the built-in {builtin.__name__}, got {type(module).__name__}" + ) + module.__class__ = cls + + +def _deferred_norm_view(hidden_states) -> Optional[torch.Tensor]: + """The tensor view of a deferred-norm packet, else None.""" + if dispatch.is_enabled("deferred_norm") and dispatch.get_impl("deferred_norm", "is_deferred")( + hidden_states + ): + return hidden_states.residual + return None + + +class GlmMLA(DeepseekV32Attention): + """A layer's MLA: the built-in attention, with the target's decode composition.""" + + # The attention output tail (``attention_output_tail``) and whether it also prefetches its + # weights into L2 (``attention_output_tail_prefetch``); None keeps the built-in output. + _attn_tail = None + _tail_prefetch = False + + @classmethod + def adopt(cls, mla: DeepseekV32Attention) -> "GlmMLA": + """Re-class a layer's built-in attention modules and attach its attention fast paths. + + Raises: + TypeError: if ``mla``, its DSA hooks, its DSA backend or its indexer is not exactly + the built-in class. + """ + reclass(mla, DeepseekV32Attention, cls) + mla._adopt_attention_tail() + GlmDsaHooks.adopt(mla.sparse_attn_hooks) + GlmDsaBackend.adopt(mla.mqa) + if mla.mqa.indexer is not None: + GlmIndexer.adopt(mla.mqa.indexer) + # The sparse-MLA decode library goes ahead of the registry's libraries on both attention + # backends, which build their FMHA libraries in ``create_weights``. + if dispatch.is_enabled("sparse_mla_decode"): + libraries = (dispatch.get_impl("sparse_mla_decode", "fmha_lib"),) + for backend in (mla.mha, mla.mqa): + if backend is not None: + backend.prepended_fmha_lib_classes = libraries + return mla + + def _adopt_attention_tail(self) -> None: + """The attention output tail: the fused generation W_UV and the local o_proj.""" + # attention_output_tail_prefetch uses the same tail object to stage the tail weights + # into L2 off the critical path while the attention runs. + self._tail_prefetch = dispatch.is_enabled("attention_output_tail_prefetch") + self._attn_tail = None + # attention_output_into_hop publishes the tail's partial into the post-attention hop's + # transport. + if ( + dispatch.is_enabled("attention_output_tail") + or self._tail_prefetch + or dispatch.is_enabled("attention_output_into_hop") + ): + self._attn_tail = dispatch.get_impl("attention_output_tail", "make_attn_tail")(self) + # oproj_weight_staging: the tail's o_proj stages the W_o rows it streams after its + # programmatic-dependent-launch wait before that wait. + if dispatch.is_enabled("oproj_weight_staging"): + self._attn_tail = dispatch.get_impl("oproj_weight_staging", "wrap_attn_tail")( + self, self._attn_tail + ) + + def forward_absorption_generation( + self, + q: torch.Tensor, + compressed_kv: torch.Tensor, + k_pe: torch.Tensor, + attn_metadata: AttentionMetadata, + output: torch.Tensor, + position_ids: Optional[torch.Tensor] = None, + latent_cache: Optional[torch.Tensor] = None, + absorbed_fused_q: Optional[torch.Tensor] = None, + deferred_indexer_k=None, + **kwargs, + ) -> torch.Tensor: + if dispatch.phase() is not dispatch.Phase.DECODE: + return super().forward_absorption_generation( + q=q, + compressed_kv=compressed_kv, + k_pe=k_pe, + attn_metadata=attn_metadata, + output=output, + position_ids=position_ids, + latent_cache=latent_cache, + **kwargs, + ) + return attention_forward.forward_absorption_generation( + self, + q, + compressed_kv, + k_pe, + attn_metadata, + output, + position_ids=position_ids, + latent_cache=latent_cache, + absorbed_fused_q=absorbed_fused_q, + deferred_indexer_k=deferred_indexer_k, + **kwargs, + ) + + def forward( + self, + position_ids: Optional[torch.Tensor], + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + all_reduce_params: Optional[AllReduceParams] = None, + latent_cache_gen: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if dispatch.phase() is not dispatch.Phase.DECODE: + return super().forward( + position_ids=position_ids, + hidden_states=hidden_states, + attn_metadata=attn_metadata, + all_reduce_params=all_reduce_params, + latent_cache_gen=latent_cache_gen, + ) + hidden_states = _helix_cp_allgather_input( + hidden_states, attn_metadata, self.mapping, self.layer_idx + ) + + # Unquantized view of the module input, used by create_mla_outputs and + # by _apply_output_gate. + output_hidden_states = hidden_states + deferred_view = _deferred_norm_view(hidden_states) + if deferred_view is not None: + # A deferred-norm packet: only the DSA query path consumes it; the outputs are + # allocated on its tensor view. + output_hidden_states = deferred_view + elif isinstance(hidden_states, Fp4QuantizedTensor): + assert hidden_states.unquantized_hidden_states is not None, ( + "MLA.forward received an Fp4QuantizedTensor without a " + "unquantized_hidden_states view" + ) + output_hidden_states = hidden_states.unquantized_hidden_states + + # The attention tail decides here, before the attention op, so the W_UV step inside the + # op and the projection below see one decision. arm() returns False under + # torch.compile and inside breakable CUDA graphs, where the op runs apart from the + # traced or captured projection. + attn_tail = getattr(self, "_attn_tail", None) + tail_armed = attn_tail is not None and attn_tail.arm(self, attn_metadata) + + if self.register_to_config: + attn_output = [ + torch.ops.trtllm.create_mla_outputs(output_hidden_states, self.layer_idx_str) + ] + self._forward_custom_op( + hidden_states, + position_ids, + attn_output, + latent_cache_gen, + ) + else: + attn_output = self._create_outputs( + hidden_states if deferred_view is None else deferred_view, attn_metadata + ) + self.forward_impl( + position_ids, + hidden_states, + attn_metadata, + attn_output=attn_output, + latent_cache_gen=latent_cache_gen, + ) + + attn_output[0] = self._apply_output_gate(output_hidden_states, attn_output[0]) + if tail_armed: + partial = attn_tail.project(self, attn_output[0]) + if partial is not None: + # The tail returns o_proj's rank-local GEMM output; reduce it as the + # ROW-parallel Linear.forward does for this o_proj (the tail takes only an + # o_proj without bias, LoRA or fused GEMM+all-reduce). + if self.o_proj.reduce_output: + partial = self.o_proj.all_reduce(partial, all_reduce_params=all_reduce_params) + return partial + return self._project_output(attn_output, position_ids, attn_metadata, all_reduce_params) + + +class GlmDsaHooks(DSAHooks): + """A layer's DSA hooks: the built-in sparse attention, with the target's decode forward.""" + + @classmethod + def adopt(cls, hooks: DSAHooks) -> "GlmDsaHooks": + """Re-class a layer's built-in DSA hooks. + + Raises: + TypeError: if ``hooks`` is not exactly the built-in ``DSAHooks``. + """ + reclass(hooks, DSAHooks, cls) + return hooks + + def forward( + self, + mla, + position_ids: Optional[torch.Tensor], + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + attn_output: list[torch.Tensor], + ) -> None: + if dispatch.phase() is not dispatch.Phase.DECODE: + return super().forward( + mla=mla, + position_ids=position_ids, + hidden_states=hidden_states, + attn_metadata=attn_metadata, + attn_output=attn_output, + ) + attention_forward.forward_sparse_attn( + mla, position_ids, hidden_states, attn_metadata, attn_output + ) + + def forward_custom_op( + self, + mla, + hidden_states: torch.Tensor, + position_ids: Optional[torch.Tensor], + attn_output: list[torch.Tensor], + latent_cache_gen: Optional[torch.Tensor], + ) -> bool: + if dispatch.phase() is not dispatch.Phase.DECODE: + return super().forward_custom_op( + mla=mla, + hidden_states=hidden_states, + position_ids=position_ids, + attn_output=attn_output, + latent_cache_gen=latent_cache_gen, + ) + attention_forward.forward_sparse_attn_custom_op( + mla, + hidden_states, + position_ids, + attn_output, + latent_cache_gen, + ) + return True + + +class GlmDsaBackend(DSATrtllmAttention): + """A layer's DSA attention backend, with the target's top-k bridge on decode steps.""" + + @classmethod + def adopt(cls, backend: DSATrtllmAttention) -> "GlmDsaBackend": + """Re-class a layer's built-in DSA attention backend. + + Raises: + TypeError: if ``backend`` is not exactly the built-in ``DSATrtllmAttention``. + """ + reclass(backend, DSATrtllmAttention, cls) + return backend + + def sparse_attn_predict( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + metadata: DSAtrtllmAttentionMetadata, + forward_args: AttentionForwardArgs, + ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: + """Run the DSA indexer and convert its local TopK to paged indices.""" + if dispatch.phase() is not dispatch.Phase.DECODE: + return super().sparse_attn_predict( + q=q, k=k, metadata=metadata, forward_args=forward_args + ) + is_generation = forward_args.attention_input_type == AttentionInputType.generation_only + sparse_backend_args = forward_args.sparse_backend_args + phase_start = metadata.num_ctx_tokens if is_generation else 0 + phase_end = metadata.num_tokens if is_generation else metadata.num_ctx_tokens + + shared_topk_indices = metadata.shared_topk_indices + # rope_kv_append: one fused bridge kernel performs the native top-k copy of this call + # and the local -> global remap. With cross_rank_topk_global_emit the cross-rank top-k + # writes them and records it on its packet. + bridge = None + bridge_path = ( + "cross_rank_topk_global_emit" + if dispatch.is_enabled("cross_rank_topk_global_emit") + else "rope_kv_append" + ) + if dispatch.is_enabled("rope_kv_append"): + bridge = dispatch.get_impl(bridge_path, "new_bridge_packet")(metadata, is_generation) + if self.indexer is None: + topk_indices = shared_topk_indices[phase_start:phase_end] + else: + indexer_intermediates = sparse_backend_args.indexer_intermediates + # The indexer's top-k region reads the bridge of this call only. + self.indexer._topk_bridge = bridge + try: + topk_indices = self.indexer.forward_from_projected( + metadata, + q, + indexer_intermediates, + is_generation=is_generation, + ) + finally: + self.indexer._topk_bridge = None + preserve_mtp_topk = metadata.in_mtp_draft_loop and self.indexer.mtp_index_share + if shared_topk_indices is not None and not preserve_mtp_topk: + if bridge is not None: + dispatch.get_impl("rope_kv_append", "defer_shared_copy")( + bridge, shared_topk_indices, phase_start, topk_indices + ) + else: + shared_topk_indices[ + phase_start : phase_start + topk_indices.shape[0], + : topk_indices.shape[1], + ].copy_(topk_indices) + + local_layer_idx = self.get_local_layer_idx(metadata) + kv_cache_dtype = getattr(metadata.kv_cache_manager, "dtype", None) + if kv_cache_dtype == tensorrt_llm.bindings.DataType.NVFP4 and not is_generation: + metadata._ensure_pool_view_cached() + page_index_scale, layer_offset = ( + metadata.kv_cache_manager.get_primary_pool_page_index_params(local_layer_idx) + ) + total_kv_tokens = metadata.num_ctx_mla_kv_tokens + if total_kv_tokens <= 0: + raise RuntimeError("NVFP4 DSA context has no active KV tokens") + head_dim = metadata.kv_cache_manager.head_dim + # Static sparse FMHA may prefetch up to TopK rows even when a + # short context has fewer valid KV tokens. The compaction itself + # can select at most min(total KV tokens, query/TopK pairs), so a + # late chunk with a long cached prefix does not need one fp8 row + # for every token in that prefix. + max_selected_tokens = min(total_kv_tokens, topk_indices.numel()) + scratch_capacity = max(max_selected_tokens, topk_indices.shape[1]) + # The previous layer's attention has already been enqueued on this + # stream. Drop its Python reference before allocating the next + # layer's scratch so the caching allocator can reuse that storage + # in stream order instead of transiently keeping two full buffers. + metadata.nvfp4_mla_context_fp8_scratch = None + scratch = torch.empty( + (scratch_capacity, 1, head_dim), + dtype=torch.float8_e4m3fn, + device=topk_indices.device, + ) + compact_indices = torch.empty_like(topk_indices) + torch.ops.trtllm.nvfp4_mla_context_kv_cache_gather( + metadata.host_kv_cache_pool_pointers, + metadata.host_kv_cache_pool_mapping, + topk_indices, + metadata._cached_req_idx_ctx, + metadata._cached_block_table_ctx, + metadata.ctx_kv_indptr[: metadata.num_contexts + 1], + scratch, + compact_indices, + self._nvfp4_mla_kv_scale_quant_orig, + local_layer_idx, + total_kv_tokens, + metadata._cached_tokens_per_block, + page_index_scale * metadata._cached_tokens_per_block, + layer_offset, + metadata.kv_cache_manager.mla_kv_cache_residual_dim, + metadata._cached_num_pool_tokens, + ) + metadata.nvfp4_mla_context_fp8_scratch = scratch + forward_args.sparse_runtime_params.aux_kv_cache_pool_ptr = scratch.data_ptr() + return compact_indices, None + + if bridge is not None: + topk_indices_global = dispatch.get_impl(bridge_path, "remap_topk_to_global")( + self, topk_indices, metadata, bridge + ) + else: + topk_indices_global = self._remap_topk_to_global(topk_indices, metadata, is_generation) + + if kv_cache_dtype == tensorrt_llm.bindings.DataType.NVFP4: + scratch = metadata.nvfp4_mla_fp8_scratch + compact_indices = metadata.nvfp4_mla_compact_indices + if scratch is None or compact_indices is None: + raise RuntimeError("NVFP4 MLA scratch buffers were not allocated") + num_rows = topk_indices_global.shape[0] + if num_rows > scratch.shape[0] or num_rows > compact_indices.shape[0]: + raise RuntimeError( + "NVFP4 MLA generation rows exceed the allocated MTP " + f"workspace: requested {num_rows}, capacity {scratch.shape[0]}" + ) + scratch = scratch[:num_rows] + compact_indices = compact_indices[:num_rows] + torch.ops.trtllm.nvfp4_mla_kv_cache_gather( + metadata.host_kv_cache_pool_pointers, + metadata.host_kv_cache_pool_mapping, + topk_indices_global, + scratch, + compact_indices, + self._nvfp4_mla_kv_scale_quant_orig, + local_layer_idx, + metadata.kv_cache_manager.mla_kv_cache_residual_dim, + metadata._cached_num_pool_tokens, + ) + # Static sparse MLA treats indices as offsets from kvPtr. Feed it + # compact offsets and the FP8 scratch base; TRTLLMGen is unchanged. + forward_args.sparse_runtime_params.aux_kv_cache_pool_ptr = scratch.data_ptr() + return compact_indices, None + + return topk_indices_global, None + + +class GlmIndexer(Indexer): + """A full-indexer layer's DSA indexer, with the target's decode top-k region.""" + + # The top-k bridge of the current ``GlmDsaBackend.sparse_attn_predict`` call (set only for + # that call, on a decode step); None keeps the built-in copies. + _topk_bridge = None + + @classmethod + def adopt(cls, indexer: Indexer) -> "GlmIndexer": + """Re-class a layer's built-in indexer. + + Raises: + TypeError: if ``indexer`` is not exactly the built-in ``Indexer``. + """ + reclass(indexer, Indexer, cls) + return indexer + + def _copy_reused_topk( + self, + metadata: DSAtrtllmAttentionMetadata, + topk_indices_buffer: torch.Tensor, + token_offset: int, + num_gen_tokens: int, + num_generations: int, + has_prefill: bool, + gvr_prior_indices: Optional[torch.Tensor], + ) -> None: + # rope_kv_append: the fused bridge reads the shared rows and writes this copy itself. + bridge = self._topk_bridge + if ( + bridge is not None + and not has_prefill + and num_gen_tokens == num_generations + and dispatch.get_impl("rope_kv_append", "defer_reuse_copy")( + bridge, + topk_indices_buffer, + token_offset, + metadata.shared_topk_indices[:num_generations, :], + gvr_prior_indices is not None, + ) + ): + return + super()._copy_reused_topk( + metadata, + topk_indices_buffer, + token_offset, + num_gen_tokens, + num_generations, + has_prefill, + gvr_prior_indices, + ) + + def _fill_decode_topk( + self, + metadata: DSAtrtllmAttentionMetadata, + q_fp8: torch.Tensor, + weights: torch.Tensor, + topk_indices_buffer: torch.Tensor, + token_offset: int, + num_gen_tokens: int, + ) -> bool: + if dispatch.phase() is not dispatch.Phase.DECODE: + return super()._fill_decode_topk( + metadata=metadata, + q_fp8=q_fp8, + weights=weights, + topk_indices_buffer=topk_indices_buffer, + token_offset=token_offset, + num_gen_tokens=num_gen_tokens, + ) + out = topk_indices_buffer[token_offset : token_offset + num_gen_tokens, :] + # The context-partitioned top-k reproduces the default selection only, so a + # deterministic top-k goes to the indexer_topk entry. + if not self.top_k.gvr_deterministic and region.cross_rank_topk_decode( + self, metadata, q_fp8, weights, token_offset, num_gen_tokens, out, self._topk_bridge + ): + return True + return region.indexer_topk_decode( + self, metadata, q_fp8, weights, token_offset, num_gen_tokens, out + ) + + def _store_mtp_shared_rows( + self, + topk_indices_buffer: torch.Tensor, + metadata: DSAtrtllmAttentionMetadata, + num_contexts: int, + num_generations: int, + num_gen_tokens: int, + token_offset: int, + has_prefill: bool, + has_decode: bool, + is_generation: Optional[bool], + ) -> None: + # rope_kv_append: the fused bridge stashes each request's last accepted row (the same + # selection as ``_mtp_last_accepted_rows``). + bridge = self._topk_bridge + if bridge is not None and has_decode and not has_prefill and is_generation: + num_accepted = metadata.mtp_num_accepted + if num_accepted is not None: + num_accepted = num_accepted[num_contexts : num_contexts + num_generations] + if dispatch.get_impl("rope_kv_append", "defer_mtp_row_stash")( + bridge, + metadata.shared_topk_indices, + num_contexts, + num_generations, + num_gen_tokens // num_generations, + num_accepted, + ): + return + super()._store_mtp_shared_rows( + topk_indices_buffer, + metadata, + num_contexts, + num_generations, + num_gen_tokens, + token_offset, + has_prefill, + has_decode, + is_generation, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/attention_forward.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/attention_forward.py new file mode 100644 index 000000000000..f18b936c2a42 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/attention_forward.py @@ -0,0 +1,744 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The MLA generation and DSA sparse-attention forward of the glm_moe_dsa sm_103 tp8 target on a decode step. + +These are the built-in functions of ``attention/mla.py`` and ``attention/backends/sparse/dsa`` +with the target's fast paths composed in: + +- the absorbed-query packet: one producer runs the DSA projections and the absorbed query + (``q_nope @ W_UK``) for a layer, so the generation skips its absorption BMM; +- the deferred-norm packet: the hidden state a post-MoE hop hands over before its RMSNorm, + normalized inside the next layer's projection (or materialized when no producer serves); +- the fused generation rope, which also performs the deferred indexer K-cache scatter; +- the attention output tail, which takes over the W_UV BMM when ``GlmMLA.forward`` armed it. + +``attention.GlmMLA`` and ``attention.GlmDsaHooks`` call them only on a decode step; every other +step runs the built-in functions. ``trtllm::glm_moe_dsa_mla_dsa_attn_inplace`` is the target's +copy of the built-in DSA op 2 (``trtllm::mla_dsa_attn_inplace``): the same op on this module's +``_forward_dsa_attn``. +""" + +from typing import List, Optional + +import torch + +from tensorrt_llm._torch.attention.attention import extract_extra_attrs +from tensorrt_llm._torch.attention.backends.interface import AttentionInputType, AttentionMetadata +from tensorrt_llm._torch.attention.backends.sparse.dsa.module import ( + forward_context_sparse_attn, + forward_dsa_proj, + forward_sparse_mla_kvcache_bf16, + should_use_short_mha, +) +from tensorrt_llm._torch.attention.backends.sparse.dsa.params import DSABackendForwardArgs +from tensorrt_llm._torch.attention.mla import fp8_block_scaling_bmm_out +from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel +from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import eager_on_graph +from tensorrt_llm._torch.utils import Fp4QuantizedTensor +from tensorrt_llm._utils import get_sm_version + +from . import dispatch + + +def _absorbed_query_packet(self, position_ids, hidden_states, attn_metadata): + """Return the absorbed-query packet of this layer's DSA projections, or None. + + The producers are tried in order; each decides its eligibility before any device work and + returns None outside its domain, in which case the next one, and finally the built-in + projection, runs: + + - a deferred-norm packet goes to the deferred-norm projection (the evict-hints variant when + selected), which normalizes it inside the projection; + - ``fused_kv_indexer_projection``: one projection launch for the fused [q_a | kv_a] and the + indexer's wk | weights_proj outputs; + - ``indexer_layer_front``: the query kernel of the full-indexer target layers and the draft + layer (the same packet bytes); + - ``absorbed_query`` (with ``shared_indexer_absorbed_query`` first on the shared-indexer + layers); + - ``draft_absorbed_query`` on the draft layer. + """ + if dispatch.is_enabled("deferred_norm") and dispatch.get_impl("deferred_norm", "is_deferred")( + hidden_states + ): + # The deferred-norm projection feeds the absorbed-query kernels; None makes the caller + # materialize the packet. + entry = ( + "deferred_norm_evict_hints" + if dispatch.is_enabled("deferred_norm_evict_hints") + else "deferred_norm" + ) + return dispatch.get_impl(entry, "dsa_proj")( + self, position_ids, hidden_states, attn_metadata + ) + if dispatch.is_enabled("fused_kv_indexer_projection"): + packet = dispatch.get_impl("fused_kv_indexer_projection", "dsa_proj")( + self, position_ids, hidden_states, attn_metadata + ) + if packet is not None: + return packet + if dispatch.is_enabled("indexer_layer_front"): + # The indexer-front query kernel on the full-indexer target layers and the draft layer + # (the same packet bytes); None runs the absorbed-query entries below. + packet = dispatch.get_impl("indexer_layer_front", "dsa_proj")( + self, position_ids, hidden_states, attn_metadata + ) + if packet is not None: + return packet + if dispatch.is_enabled("absorbed_query"): + if dispatch.is_enabled("shared_indexer_absorbed_query"): + packet = dispatch.get_impl("absorbed_query", "shared_dsa_proj")( + self, position_ids, hidden_states, attn_metadata + ) + if packet is not None: + return packet + packet = dispatch.get_impl("absorbed_query", "dsa_proj")( + self, position_ids, hidden_states, attn_metadata + ) + if packet is not None: + return packet + if dispatch.is_enabled("draft_absorbed_query"): + return dispatch.get_impl("absorbed_query", "mtp_dsa_proj")( + self, position_ids, hidden_states, attn_metadata + ) + return None + + +def _deferred_norm_materialize(hidden_states): + """The normalized hidden state of a deferred-norm packet (bit-identical to the post-MoE + hop's normalized output); any other input passes through.""" + if dispatch.is_enabled("deferred_norm") and dispatch.get_impl("deferred_norm", "is_deferred")( + hidden_states + ): + return hidden_states.materialize() + return hidden_states + + +def forward_absorption_generation( + self, + q: torch.Tensor, + compressed_kv: torch.Tensor, + k_pe: torch.Tensor, + attn_metadata: AttentionMetadata, + output: torch.Tensor, + position_ids: Optional[torch.Tensor] = None, + latent_cache: Optional[torch.Tensor] = None, + absorbed_fused_q: Optional[torch.Tensor] = None, + deferred_indexer_k=None, + **kwargs, +) -> torch.Tensor: + """Absorbed-MLA generation (``MLA.forward_absorption_generation``) on a decode step. + + ``absorbed_fused_q`` is the absorbed-query packet: a + ``[num_tokens, num_heads_tp, kv_lora_rank + qk_rope_head_dim]`` buffer whose + ``[..., :kv_lora_rank]`` already holds ``q_nope @ W_UK``. When given, only the absorption + BMM is skipped; RoPE, cache append, FP8 query quantization and scheduler buffers run as + usual. ``deferred_indexer_k`` is the indexer K-cache scatter ``_forward_dsa_attn`` deferred + into this call's rope. + """ + num_tokens = q.shape[0] + q_nope, q_pe = q.view([-1, self.num_heads_tp, self.qk_head_dim]).split( + [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 + ) + + # fused_q contains 1) the result of the following bmm with shape [num_tokens, num_heads, kv_lora_rank] + # 2) rope(q_pe) with shape [num_tokens, num_heads, qk_rope_head_dim]. rope is applied inside AttentionOp + num_seqs = attn_metadata.num_seqs + + # The rope kernel writes seqQOffset[batch_idx + 1] for every + # batch_idx < batch_size, and CUB-scans total_s_len / seq_len entries + # into seqKVOffsets (mlaKernels.cu:579, :650). `batch_size` there is the + # count of presented rows, which equals num_seqs only while every + # request contributes the same number of query tokens. Under the + # token-major presentation a row is a token, so size these by the token + # count. Uniform verification keeps the established request-major + # allocation rather than scaling these temporary buffers by K. + cu_seqlens_rows = ( + max(num_seqs, num_tokens) + 1 + if getattr(attn_metadata, "is_ragged_verify", False) + else num_seqs + 1 + ) + cu_q_seqlens = torch.empty(cu_seqlens_rows, dtype=torch.int32, device=q.device) + cu_kv_seqlens = torch.empty(cu_seqlens_rows, dtype=torch.int32, device=q.device) + fmha_scheduler_counter = torch.empty(1, dtype=torch.uint32, device=q.device) + use_fp8_mla = getattr(self.mqa, "has_fp8_kv_cache", False) or getattr( + self.mqa, "has_fp4_kv_cache", False + ) + + mla_bmm1_scale = None + mla_bmm2_scale = None + quant_q_buffer = None + if use_fp8_mla: + mla_bmm1_scale = torch.empty(2, dtype=torch.float32, device=q.device) + mla_bmm2_scale = torch.empty(1, dtype=torch.float32, device=q.device) + quant_q_buffer = torch.empty( + num_tokens, + self.num_heads_tp, + (self.kv_lora_rank + self.qk_rope_head_dim), + dtype=torch.uint8, + device=q.device, + ) + + if hasattr(self, "k_b_proj_trans"): + if absorbed_fused_q is not None: + fused_q = absorbed_fused_q + else: + fused_q = torch.empty( + [ + num_tokens, + self.num_heads_tp, + (self.kv_lora_rank + self.qk_rope_head_dim), + ], + dtype=q.dtype, + device=q.device, + ) + + fp4_mla = ( + self.sparse_params is None + and self.mqa.support_fp4_kv_cache() + and self.quant_config is not None + and self.quant_config.layer_quant_mode.has_fp4_kv_cache() + ) + if fp4_mla and latent_cache is None: + raise RuntimeError( + "FP4 MLA generation requires a latent cache for fused " + "Q quantization, RoPE, and cache update." + ) + fuse_fp4_q_quant = fp4_mla + fp4_rope_kwargs = {"fuse_fp4_q_quant": True} if fuse_fp4_q_quant else {} + + def _mla_gen_rope(): + # rope_kv_append: one fused launch writes the bytes of mla_rope_generation plus the + # deferred indexer K-cache scatter. Every check runs before any device work; + # otherwise the deferred scatter runs natively first and the native path follows. + if deferred_indexer_k is not None or dispatch.is_enabled("rope_kv_append"): + rope_fused = ( + not self.apply_rotary_emb + and self.kv_cache_dtype != "fp8_ds_mla" + and not fp4_rope_kwargs + and dispatch.get_impl("rope_kv_append", "rope_eligible")( + self.mqa, + attn_metadata, + fused_q, + q_pe, + latent_cache, + quant_q_buffer, + mla_bmm1_scale, + mla_bmm2_scale, + ) + ) + if rope_fused: + # indexer_layer_front: the same fused rope, forked to trigger its + # dependents after its last read (the same bytes). + rope_owner = "rope_kv_append" + if dispatch.is_enabled("indexer_layer_front"): + front_rope = dispatch.get_impl("indexer_layer_front", "rope_available")() + dispatch.log_capability( + "indexer_layer_front", + "rope", + front_rope, + "" if front_rope else "indexer_layer_front rope op not loaded", + ) + if front_rope: + rope_owner = "indexer_layer_front" + dispatch.get_impl(rope_owner, "mla_rope_generation")( + self.mqa, + fused_q, + q_pe, + latent_cache, + attn_metadata, + cu_q_seqlens, + cu_kv_seqlens, + fmha_scheduler_counter, + mla_bmm1_scale, + mla_bmm2_scale, + quant_q_buffer, + indexer_k=deferred_indexer_k, + ) + return + if deferred_indexer_k is not None: + dispatch.get_impl("rope_kv_append", "run_indexer_k_scatter_native")( + deferred_indexer_k + ) + if self.apply_rotary_emb: + # Non-fused backends (Vanilla / FlashInfer) do not fuse RoPE + # in the attention kernel. Reuse apply_rope, which rotates + # q's q_pe slice in place, then copy rotated k_pe into the + # latent cache that the backend appends. + assert position_ids is not None + assert latent_cache is not None + gen_position_ids = position_ids[..., attn_metadata.num_ctx_tokens :] + k_pe_rope = self.apply_rope(q, k_pe, gen_position_ids) + fused_q[..., self.kv_lora_rank :] = q_pe + latent_cache[..., self.kv_lora_rank :] = k_pe_rope + else: + # Fused backend (TRTLLM): RoPE, latent-cache append and the + # trtllm-gen scheduler buffers are all produced in-kernel. + if self.kv_cache_dtype == "fp8_ds_mla": + fused_q[..., self.kv_lora_rank :] = q_pe + else: + self.mqa.mla_rope_generation( + fused_q, + q_pe, + latent_cache, + attn_metadata, + cu_q_seqlens, + cu_kv_seqlens, + fmha_scheduler_counter, + mla_bmm1_scale, + mla_bmm2_scale, + quant_q_buffer, + **fp4_rope_kwargs, + ) + + if fuse_fp4_q_quant or use_fp8_mla: + rope_stream = None + else: + rope_stream = self.aux_stream + if absorbed_fused_q is not None: + _mla_gen_rope() + elif self.k_b_proj_trans.dtype == torch.bfloat16: + # [num_heads, num_tokens, self.qk_nope_head_dim] + q_nope_t = q_nope.transpose(0, 1) + # [num_heads, num_tokens, self.kv_lora_rank] + q_nope_out = fused_q[..., : self.kv_lora_rank].transpose(0, 1) + + # [num_heads, num_tokens, self.qk_nope_head_dim] x [num_heads, kv_lora_rank, qk_nope_head_dim] + # -> [num_heads, num_tokens, kv_lora_rank] -> [num_tokens, num_heads, kv_lora_rank] + # The output of bmm is written directly into fused_q + maybe_execute_in_parallel( + lambda: self._bmm_bf16_out( + q_nope_t, + self.k_b_proj_trans, + self.k_b_proj_trans.transpose(1, 2), + q_nope_out, + self._locality_domain_k_b_proj_trans_shards, + ), + _mla_gen_rope, + self.ln_events[0], + self.ln_events[1], + rope_stream, + ) + + elif self.k_b_proj_trans.dtype == torch.float8_e4m3fn: + # [num_heads, num_tokens, self.kv_lora_rank] + q_nope_out = fused_q[..., : self.kv_lora_rank].transpose(0, 1) + + maybe_execute_in_parallel( + lambda: fp8_block_scaling_bmm_out( + q_nope, + self.k_b_proj_trans, + self.k_b_proj_trans_scale, + q_nope_out, + self.k_b_proj_trans_dequant, + self.use_cute_dsl_blockscaling_bmm, + ), + _mla_gen_rope, + self.ln_events[0], + self.ln_events[1], + rope_stream, + ) + else: + raise NotImplementedError(f"Missing bmm impl for dtype: {self.k_b_proj_trans.dtype}.") + + fused_q = fused_q.view( + [ + num_tokens, + self.num_heads_tp * (self.kv_lora_rank + self.qk_rope_head_dim), + ] + ) + else: + raise RuntimeError("MLA absorption requires k_b_proj_trans") + + # attention_output_tail_prefetch: on a call that GlmMLA.forward armed, fork the L2 + # prefetch of the tail weights onto a private stream while the attention runs; the + # projection step joins it. No-op otherwise (decided before any device work). + attn_tail = getattr(self, "_attn_tail", None) + if getattr(self, "_tail_prefetch", False) and attn_tail is not None: + attn_tail.before_attention(self, point="attn") + + attn_out_latent = self._attn_forward_gen( + self.mqa, + fused_q, + None, + None, + position_ids, + attn_metadata, + attention_input_type=AttentionInputType.generation_only, + out_scale=self.out_scale, + latent_cache=latent_cache, # kvcache and k_pe + q_pe=q_pe, # used by `invokeMLARopeGeneration` + cu_q_seqlens=cu_q_seqlens, # used by `mlaGeneration` + cu_kv_seqlens=cu_kv_seqlens, # used by `mlaGeneration` + fmha_scheduler_counter=fmha_scheduler_counter, # used by `mlaGeneration` + mla_bmm1_scale=mla_bmm1_scale, # used by `mlaGeneration` + mla_bmm2_scale=mla_bmm2_scale, # used by `mlaGeneration` + quant_q_buffer=quant_q_buffer, # used by `mlaGeneration` + **kwargs, + ) + fused_q = None + + # note: if we do not have CP, then num_heads_tp_cp == num_heads_tp + assert ( + attn_out_latent.shape[0] == q.shape[0] + and attn_out_latent.shape[1] == self.num_heads_tp_cp * self.kv_lora_rank + ) + + # [seq, num_heads, kv_lora_rank] + attn_out_latent = attn_out_latent.view([-1, self.num_heads_tp_cp, self.kv_lora_rank]) + + # Acts only on a call that GlmMLA.forward armed; when the attention tail takes the call, + # its projection step consumes the latent and the native W_UV BMM is skipped. The latent + # shape is checked before any device work. + attn_tail = getattr(self, "_attn_tail", None) + if attn_tail is not None and attn_tail.defer_w_uv(self, attn_metadata, attn_out_latent, output): + return output + + attn_output = output.view([num_tokens, self.num_heads_tp_cp, self.v_head_dim]) + + if self.v_b_proj.dtype == torch.bfloat16: + # [num_heads, seq, kv_lora_rank] x [num_heads, kv_lora_rank, v_head_dim] + # -> [num_heads, seq, v_head_dim] + self._bmm_bf16_out( + attn_out_latent.transpose(0, 1), + self.v_b_proj, + self.v_b_proj.transpose(1, 2), + attn_output.transpose(0, 1), + self._locality_domain_v_b_proj_shards, + ) + elif self.v_b_proj.dtype == torch.float8_e4m3fn: + fp8_block_scaling_bmm_out( + attn_out_latent, + self.v_b_proj, + self.v_b_proj_scale, + attn_output.transpose(0, 1), + self.v_b_proj_dequant, + self.use_cute_dsl_blockscaling_bmm, + ) + else: + raise NotImplementedError(f"Missing bmm impl for dtype: {self.v_b_proj.dtype}.") + + return output + + +def forward_generation_sparse_attn( + self, + q: torch.Tensor, + compressed_kv: torch.Tensor, + k_pe: torch.Tensor, + attn_metadata: AttentionMetadata, + output: torch.Tensor, + position_ids: Optional[torch.Tensor] = None, + latent_cache: Optional[torch.Tensor] = None, + indexer_intermediates: Optional[List[torch.Tensor]] = None, + absorbed_fused_q: Optional[torch.Tensor] = None, + deferred_indexer_k=None, +) -> torch.Tensor: + if get_sm_version() >= 100: + return self.forward_absorption_generation( + q, + compressed_kv, + k_pe, + attn_metadata, + output, + position_ids=position_ids, + latent_cache=latent_cache, + absorbed_fused_q=absorbed_fused_q, + deferred_indexer_k=deferred_indexer_k, + sparse_backend_args=DSABackendForwardArgs(indexer_intermediates=indexer_intermediates), + ) + assert absorbed_fused_q is None, "absorbed_query preabsorbed queries require SM100+" + assert deferred_indexer_k is None + return forward_sparse_mla_kvcache_bf16( + self, + q, + latent_cache, + attn_metadata, + output, + indexer_intermediates, + is_generation=True, + ) + + +def forward_sparse_attn( + self, + position_ids: Optional[torch.Tensor], + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + attn_output: list[torch.Tensor], +) -> None: + """ + Forward pass for the MLA module with DSA (always in MQA mode). + Writes result into output tensor in-place. + + Delegates to forward_dsa_proj (token-wise projections) followed by + _forward_dsa_attn (batch-dependent attention dispatch). An absorbed-query packet replaces + forward_dsa_proj; a deferred-norm packet that no producer takes is materialized first. + + Args: + position_ids (Optional[torch.IntTensor]): The position IDs. + hidden_states (torch.Tensor): The hidden states. + attn_metadata (AttentionMetadata): The attention metadata. + attn_output: Algorithm-defined output tensors. DSA writes to the + standard output at index zero. + """ + output = attn_output[0] + query_packet = _absorbed_query_packet(self, position_ids, hidden_states, attn_metadata) + if query_packet is None: + hidden_states = _deferred_norm_materialize(hidden_states) + proj_outputs = forward_dsa_proj(self, position_ids, hidden_states, attn_metadata) + absorbed_fused_q = None + else: + proj_outputs, absorbed_fused_q = query_packet.proj_outputs, query_packet.fused_q + q, compressed_kv, k_pe, latent_cache = proj_outputs[:4] + indexer_intermediates = proj_outputs[4:] + _forward_dsa_attn( + self, + q, + compressed_kv, + k_pe, + latent_cache, + indexer_intermediates, + position_ids, + attn_metadata, + output, + absorbed_fused_q=absorbed_fused_q, + ) + + +def forward_sparse_attn_custom_op( + self, + hidden_states: torch.Tensor, + position_ids: Optional[torch.Tensor], + attn_output: list[torch.Tensor], + latent_cache_gen: Optional[torch.Tensor], +) -> None: + """Run DSA's two-stage custom ops for piecewise CUDA graph capture. + + The second op is the target's ``trtllm::glm_moe_dsa_mla_dsa_attn_inplace``. With an + absorbed-query packet the attention runs directly, without the ops. + """ + del latent_cache_gen + from tensorrt_llm._torch.attention.backends.sparse.dsa import custom_ops # noqa: F401 + + if ( + dispatch.is_enabled("absorbed_query") or dispatch.is_enabled("draft_absorbed_query") + ) and not isinstance(hidden_states, Fp4QuantizedTensor): + metadata, _ = extract_extra_attrs(self.layer_idx_str, "mla") + query_packet = _absorbed_query_packet(self, position_ids, hidden_states, metadata) + if query_packet is not None: + # The packet producers decline under torch.compile and breakable CUDA graphs, so + # the two-op split is not needed on this path. + proj_outputs = query_packet.proj_outputs + _forward_dsa_attn( + self, + *proj_outputs[:4], + proj_outputs[4:], + position_ids, + metadata, + attn_output[0], + absorbed_fused_q=query_packet.fused_q, + ) + return + hidden_states = _deferred_norm_materialize(hidden_states) + + if isinstance(hidden_states, Fp4QuantizedTensor): + proj_outputs = torch.ops.trtllm.mla_dsa_proj( + hidden_states.unquantized_hidden_states, + position_ids, + self.layer_idx_str, + hidden_states.fp4_tensor, + hidden_states.scaling_factor, + ) + else: + proj_outputs = torch.ops.trtllm.mla_dsa_proj( + hidden_states, position_ids, self.layer_idx_str + ) + q, compressed_kv, k_pe, latent_cache = proj_outputs[:4] + indexer_intermediates = proj_outputs[4:] + maybe_bcg_glm_moe_dsa_mla_dsa_attn_inplace( + q, + compressed_kv, + k_pe, + latent_cache, + indexer_intermediates, + position_ids, + self.layer_idx_str, + attn_output[0], + ) + + +def _forward_dsa_attn( + self, + q: torch.Tensor, + compressed_kv: torch.Tensor, + k_pe: torch.Tensor, + latent_cache: torch.Tensor, + indexer_intermediates: List[torch.Tensor], + position_ids: Optional[torch.Tensor], + attn_metadata: AttentionMetadata, + output: torch.Tensor, + absorbed_fused_q: Optional[torch.Tensor] = None, +) -> None: + """Batch-structure-dependent attention for DSA MLA (Op 2, not graph-captured). + + indexer_intermediates is [q_fp8, k_fp8, k_scale, weights, q_scale] + when the indexer ran in Op 1, or [] when short-MHA handled all tokens. + + All num_tokens slicing happens here (not in Op 1) because + num_tokens comes from batch-specific metadata and must not be + baked into CUDA graph capture. + """ + num_contexts = attn_metadata.num_contexts + num_generations = attn_metadata.num_generations + num_ctx_tokens = attn_metadata.num_ctx_tokens + num_tokens = attn_metadata.num_tokens + + # Slice Op 1 outputs to actual num_tokens (Op 1 operates on the + # full padded tensor for CUDA graph compatibility). + q = q[:num_tokens, ...] + compressed_kv = compressed_kv[:num_tokens, ...] + k_pe = k_pe[:num_tokens, ...] + latent_cache = latent_cache[:num_tokens, ...] + if position_ids is not None: + position_ids = position_ids[..., :num_tokens] + + use_short_mha_for_ctx = num_contexts > 0 and should_use_short_mha( + self, attn_metadata, position_ids + ) + + deferred_indexer_k = None + if not (use_short_mha_for_ctx and num_generations == 0) and self.mqa.indexer is not None: + q_fp8, k_fp8, k_scale, weights, q_scale = indexer_intermediates + # Slice indexer intermediates to actual num_tokens (they were + # computed on the full padded tensor in Op 1). + q_fp8 = q_fp8[:num_tokens, ...] + k_fp8 = k_fp8[:num_tokens, ...] + k_scale = k_scale[:num_tokens, ...] + weights = weights[:num_tokens, ...] + q_scale = q_scale[:num_tokens, ...] + indexer_intermediates = [q_fp8, k_fp8, k_scale, weights, q_scale] + # Update the full mixed batch once before context/generation prediction + # splits. Context chunk gathers must observe the newly appended keys. + # rope_kv_append: a pure-generation batch folds this scatter into the same layer's + # fused generation rope, which still runs before the indexer reads the K cache. The + # deferred scatter must reach the generation call below, which is checked right after + # that call. + if dispatch.is_enabled("rope_kv_append"): + deferred_indexer_k = dispatch.get_impl("rope_kv_append", "defer_indexer_k_scatter")( + self, k_fp8, k_scale, attn_metadata, forwards_packet=True + ) + if deferred_indexer_k is None: + self.mqa.indexer._update_k_cache(k_fp8, k_scale, attn_metadata) + else: + indexer_intermediates = None + + assert output is not None, "output must be provided" + + if num_contexts > 0: + q_ctx = q[:num_ctx_tokens, ...] + compressed_kv_ctx = compressed_kv[:num_ctx_tokens, ...] + k_pe_ctx = k_pe[:num_ctx_tokens, ...] + latent_cache_ctx = latent_cache[:num_ctx_tokens, ...] + ctx_position_ids = position_ids[..., :num_ctx_tokens] if position_ids is not None else None + if self.apply_rotary_emb: + assert ctx_position_ids is not None + k_pe_ctx = self.apply_rope(q_ctx, k_pe_ctx, ctx_position_ids) + + forward_context_sparse_attn( + self, + q_ctx, + compressed_kv_ctx, + k_pe_ctx, + attn_metadata, + output[:num_ctx_tokens, :], + latent_cache_ctx, + position_ids=ctx_position_ids, + indexer_intermediates=indexer_intermediates, + ) + + if num_generations > 0: + q_gen = q[num_ctx_tokens:, ...] + compressed_kv_gen = compressed_kv[num_ctx_tokens:, ...] + k_pe_gen = k_pe[num_ctx_tokens:, ...] + latent_cache_gen = latent_cache[num_ctx_tokens:, ...] + gen_position_ids = ( + position_ids[..., num_ctx_tokens:num_tokens] if position_ids is not None else None + ) + if self.apply_rotary_emb: + assert gen_position_ids is not None + k_pe_gen = self.apply_rope(q_gen, k_pe_gen, gen_position_ids) + + forward_generation_sparse_attn( + self, + q_gen, + compressed_kv_gen, + k_pe_gen, + attn_metadata, + output[num_ctx_tokens:num_tokens, :], + latent_cache=latent_cache_gen, + position_ids=gen_position_ids, + indexer_intermediates=indexer_intermediates, + absorbed_fused_q=( + absorbed_fused_q[num_ctx_tokens:num_tokens] + if absorbed_fused_q is not None + else None + ), + deferred_indexer_k=deferred_indexer_k, + ) + if deferred_indexer_k is not None: + # Raises unless the generation rope (fused or the native fallback) executed the + # deferred indexer K-cache scatter. + dispatch.get_impl("rope_kv_append", "check_indexer_k_scatter_consumed")(deferred_indexer_k) + + # Join the prev_topk copy forked in sparse_attn_indexer; CUDA graph + # capture requires the join within the same layer's forward. + if self.mqa.indexer is not None: + self.mqa.indexer.maybe_join_prev_topk_copy() + + +@torch.library.custom_op("trtllm::glm_moe_dsa_mla_dsa_attn_inplace", mutates_args=("output",)) +def glm_moe_dsa_mla_dsa_attn_inplace( + q: torch.Tensor, + compressed_kv: torch.Tensor, + k_pe: torch.Tensor, + latent_cache: torch.Tensor, + indexer_intermediates: List[torch.Tensor], + position_ids: Optional[torch.Tensor], + layer_idx: str, + output: torch.Tensor, +) -> None: + """Batch-structure-dependent attention dispatch for DSA MLA, on ``_forward_dsa_attn`` above. + + The built-in ``trtllm::mla_dsa_attn_inplace`` with the target's op-2 function: + indexer_intermediates is [q_fp8, k_fp8, k_scale, weights, q_scale] when the indexer ran in + Op 1, or [] when short-MHA handled all tokens. This op is excluded from CUDA graph capture. + """ + metadata, mla_layer = extract_extra_attrs(layer_idx, "mla") + _forward_dsa_attn( + mla_layer, + q, + compressed_kv, + k_pe, + latent_cache, + indexer_intermediates, + position_ids, + metadata, + output, + ) + + +@glm_moe_dsa_mla_dsa_attn_inplace.register_fake +def _glm_moe_dsa_mla_dsa_attn_inplace_fake( + q: torch.Tensor, + compressed_kv: torch.Tensor, + k_pe: torch.Tensor, + latent_cache: torch.Tensor, + indexer_intermediates: List[torch.Tensor], + position_ids: Optional[torch.Tensor], + layer_idx: str, + output: torch.Tensor, +) -> None: + """Model the in-place output mutation during fake-tensor propagation.""" + + +maybe_bcg_glm_moe_dsa_mla_dsa_attn_inplace = eager_on_graph(glm_moe_dsa_mla_dsa_attn_inplace) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/dispatch.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/dispatch.py new file mode 100644 index 000000000000..3aa079ae0962 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/dispatch.py @@ -0,0 +1,1209 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode fast paths of the glm_moe_dsa sm_103 tp8 target, which of them a process runs, and what happens +when one cannot. + +Composition. ``FAST_PATHS`` names every fast path of the family once, with the +paths it extends (``depends_on``): a path that extends another has no effect +without it. Each target serves some of them (``TargetPaths``), and every path +it serves is on by default. ``TRTLLM_STAIRCASE_GLM_LL_DISABLE`` takes a +comma-separated list of path names and marks those off in the plan together +with every path that depends on them, transitively; a name that is not a path +the target serves raises. ``resolve`` reads it once, at construction, into a +``Plan``. The fast paths read the installed plan through ``is_enabled``; +``DISABLE`` turns a path and its dependents off. + +Rank uniformity. Several fast paths are collectives across the model's ranks: a +rank that runs one while a peer runs the generic path waits forever. So every +rank must read the same knob values. ``check_rank_uniform`` allgathers them, +with the shared switches the kernels read and the modeling_v2 mode +(``rank_uniform_inputs``), once, at construction, before they are parsed and +before any fast path builds a collective workspace, and raises on every rank +with the per-rank table on a mismatch. The plan is a function of the knobs +alone, so equal knobs give every rank the same plan, and a value no rank can +parse fails on every rank alike. Once the model is built, +``check_rank_uniform_adopted`` compares the fast-path state every rank +attached to its modules the same way. + +Phases. ``step`` runs a model step under the ``Phase`` of the target that +takes it (``PREFILL`` or ``DECODE``); ``phase`` reads it, and is None outside +a step. + +Two kinds of not taking a fast path, kept apart: + +* **uncovered**: the step is outside the shapes the path takes (more rows than + its kernel holds, a KV length outside its tier, a prefill step). ``uncovered`` + logs the reason once per (path, reason); the caller runs the next tier or the + generic path in every mode. +* **missed**: the step is inside the shapes the path takes, but a precondition + is absent (a workspace, a compiled kernel, multi-node NVLink memory, a device + check). Under ``modeling_v2='require'`` ``missed`` raises + ``FastPathUnavailable`` naming the path and the precondition; otherwise it + logs once per (path, reason), counts the miss, and the caller falls back. A + miss decided per call never raises: it is counted and logged once. + +The kernels report their decisions through ``log_capability``, which also +classifies every decline (``classify_decline``) by where it is decided +(``DECLINE_KINDS``), the layers each path serves (``LAYER_CELLS``) and the +reason: a path turned off is neither kind, a decline decided once inside the +path's layers is a miss, and a decline decided per call is a miss only when +its reason names a missing resource. + +Selection. A kernel or a shared hook asks ``is_enabled(path)`` whether a fast +path runs, looks up its entry points with ``get_impl(path, name)`` (the kernel +package of the path registers them with ``register`` when it is imported) and +logs its per-rank decisions with ``log_capability``. ``is_enabled`` answers +from the installed plan (``install``); without one, which is the case of every +process that does not build a target of the family, no fast path runs. ``install`` logs a +record line (``feature_record``): every enabled path with its numerical variant +(``NUMERICAL_VARIANTS``), every set knob of the family and the shared +TensorRT-LLM switches the kernels read. + +Validation. ``TRTLLM_MODELING_V2_VALIDATE`` (read once, in ``resolve``) turns on +the diagnostics of the fast paths that have some (``validate_mode``): the +step-input stager's guard and the draft loop's in-pipeline native oracle. + +Knobs. The knobs of the family are ``KNOBS``: ``DISABLE``, the numerical variant +of ``routed_moe_into_hop`` and the cross-rank top-k wait timeout (read by the +C++ kernels). A target refuses a knob that only a path it does not serve reads +(``KNOB_PATHS``). The kernels run at fixed launch parameters; ``resolve`` refuses +every other ``TRTLLM_STAIRCASE_GLM_LL_*`` variable that is set. +""" + +from __future__ import annotations + +import enum +import functools +import hashlib +import importlib +import json +import os +import re +import threading +from dataclasses import dataclass, field +from typing import ( + Any, + Callable, + Dict, + FrozenSet, + Iterable, + List, + Mapping, + Optional, + Sequence, + Tuple, +) + +from tensorrt_llm._torch._experimental.modeling_v2._router_index import ( + MODELING_V2_VALIDATE_ENV, + ModelingV2Mode, + step_contract_enabled, +) +from tensorrt_llm.logger import logger + +#: Prefix of every knob of this target; every set variable with it is part of the rank-uniform +#: check. +KNOB_PREFIX = "TRTLLM_STAIRCASE_GLM_LL_" +#: Comma-separated fast-path names to turn off, with every path that depends on them. +DISABLE_ENV = KNOB_PREFIX + "DISABLE" +#: The numerical variant of ``routed_moe_into_hop`` (``NUMERICAL_VARIANTS``). +ACCUMULATION_ENV = KNOB_PREFIX + "ROUTED_MOE_INTO_HOP_ACCUMULATION" +#: How long the cross-rank top-k kernels wait for their peers before they trap, in ms (C++). +CROSS_RANK_TOPK_TIMEOUT_ENV = KNOB_PREFIX + "CROSS_RANK_TOPK_TIMEOUT_MS" +#: The knobs of this target; ``resolve`` refuses any other set ``KNOB_PREFIX`` variable. +KNOBS: FrozenSet[str] = frozenset({DISABLE_ENV, ACCUMULATION_ENV, CROSS_RANK_TOPK_TIMEOUT_ENV}) +#: A fast-path selector this target does not read; ``resolve`` refuses it. +FEATURE_ENV = KNOB_PREFIX + "FEATURES" +#: A narrowing of ``deferred_input_prep`` to some of its parts; ``resolve`` refuses it. +PARTS_ENV = KNOB_PREFIX + "DEFERRED_INPUT_PREP_PARTS" + +_VALIDATE_HINT = f"set {MODELING_V2_VALIDATE_ENV}=1 instead" +#: Knobs this target does not read, each with what to use instead; ``resolve`` refuses them. +REMOVED_KNOBS: Mapping[str, str] = { + FEATURE_ENV: f"every fast path of this target is on; {KNOB_PREFIX}DISABLE turns paths off", + PARTS_ENV: ( + "each part is a fast path of its own (prep = deferred_input_prep, store = " + f"sampler_store, stream = execution_stream_loop); list the parts to drop in " + f"{KNOB_PREFIX}DISABLE" + ), + KNOB_PREFIX + "DEFERRED_INPUT_PREP_DEBUG": _VALIDATE_HINT, + KNOB_PREFIX + "DRAFT_LOOP_VERIFY": _VALIDATE_HINT, + KNOB_PREFIX + "CROSS_LAYER_EXPERT_PREFETCH_STATS": "the hit counters are not reported", + KNOB_PREFIX + "CROSS_LAYER_EXPERT_PREFETCH_REPORT_EVERY": "the hit counters are not reported", + KNOB_PREFIX + "SHARED_INDEXER_ABSORBED_QUERY_FORCE_NATIVE": ( + f"list shared_indexer_absorbed_query in {KNOB_PREFIX}DISABLE to run absorbed_query on the " + "shared-indexer layers" + ), +} +#: What ``resolve`` says about any other set variable with the prefix that is not in ``KNOBS``. +_FIXED_HINT = "the fast paths of this target run at fixed launch parameters" + +_LOG_TAG = "[glm_moe_dsa]" + + +@dataclass(frozen=True) +class FastPath: + """One decode fast path: its name and the paths it extends.""" + + name: str + depends_on: Tuple[str, ...] = () + + +# The attention query front, the indexer and top-k selection, sparse attention, +# the attention output and its cross-rank hop, the MoE region and its hop, the +# dense layers, the draft (MTP) layer and loop, and the engine-side step work. +FAST_PATHS: Tuple[FastPath, ...] = ( + # The target layers' query front: the query projections and the absorbed query. + FastPath("absorbed_query"), + # The same query front on the draft layer. + FastPath("draft_absorbed_query"), + # The query front of layers that reuse an earlier layer's indexer selection. + FastPath("shared_indexer_absorbed_query", depends_on=("absorbed_query",)), + # One projection launch for the latent KV and the indexer key and weights. + FastPath("fused_kv_indexer_projection", depends_on=("absorbed_query",)), + # The query front of layers that run the full indexer, with its rope and + # top-k exchange started early. + FastPath( + "indexer_layer_front", + depends_on=( + "absorbed_query", + "draft_absorbed_query", + "indexer_topk", + "rope_kv_append", + "cross_rank_topk", + ), + ), + # The next layer's RMSNorm deferred across the post-MoE hop into its projection. + FastPath("deferred_norm", depends_on=("absorbed_query", "post_moe_hop")), + # The deferred-norm hop and projection chain with cache-eviction hints. + FastPath("deferred_norm_evict_hints", depends_on=("deferred_norm",)), + # Rope, FP8 query, latent-KV append, indexer-K scatter and the top-k to + # KV-index bridge in one launch per layer. + FastPath("rope_kv_append"), + # Indexer logits and top-k selection on decode steps. + FastPath("indexer_topk"), + # Top-k over the KV pages each rank owns, merged across the ranks. + FastPath("cross_rank_topk"), + # The cross-rank top-k with a cluster merge of the per-rank candidates. + FastPath("cross_rank_topk_cluster_merge", depends_on=("cross_rank_topk",)), + # The cluster merge also emits the top-k to KV-index bridge. + FastPath( + "cross_rank_topk_global_emit", + depends_on=("rope_kv_append", "cross_rank_topk_cluster_merge"), + ), + # Sparse MLA attention for the decode rows. + FastPath("sparse_mla_decode"), + # The attention output projection and its local tail. + FastPath("attention_output_tail"), + # The attention tail's prefetch schedule: the next weights are fetched while the current + # attention runs. + FastPath("attention_output_tail_prefetch", depends_on=("attention_output_tail",)), + # The attention output published straight into the post-attention hop. + FastPath( + "attention_output_into_hop", depends_on=("attention_output_tail", "post_attention_hop") + ), + # Output-projection weight staging inside an already-waiting kernel of the attention window. + FastPath("oproj_weight_staging", depends_on=("attention_output_into_hop", "sparse_mla_decode")), + # Cross-rank sum, residual add, RMSNorm and router after attention, in one hop. + FastPath("post_attention_hop"), + # The MoE routing (top-8) computed inside the post-attention hop. + FastPath("hop_routing", depends_on=("attention_output_into_hop", "routed_moe")), + # The router GEMM at the pre-MoE boundary. + FastPath("router_gemm"), + # The routed NVFP4 experts in one launch. + FastPath("routed_moe"), + # The routed-expert result written straight into the post-MoE hop. + FastPath("routed_moe_into_hop", depends_on=("post_moe_hop", "routed_moe")), + # The shared expert's weights prefetched into L2 by the routed-expert launch. + FastPath("shared_expert_prefetch", depends_on=("routed_moe_into_hop", "hop_routing")), + # Cross-rank finalize, reduction and RMSNorm after the MoE, in one hop. + FastPath("post_moe_hop"), + # The next MoE layer's predicted experts prefetched into L2. + FastPath("cross_layer_expert_prefetch", depends_on=("routed_moe",)), + # The dense layers' small-row MLP. + FastPath("dense_mlp"), + # The draft layer's input normalization and packing. + FastPath("draft_front"), + # The draft layer's bf16 routed experts. + FastPath("draft_routed_moe"), + # Only the accepted rows of the first draft step through the post-attention hop. + FastPath("draft_live_rows", depends_on=("post_attention_hop", "draft_loop")), + # Draft-token selection at the end of each draft step. + FastPath("draft_tail"), + # The draft loop's embedding, bookkeeping and sparse-attention metadata. + FastPath("draft_loop"), + # The decode step's metadata prologue and the target embedding in one launch. + FastPath("decode_prologue", depends_on=("draft_loop",)), + # The decode step's input preparation recorded and issued as one launch. + FastPath("deferred_input_prep"), + # The speculative sampler's store update and host copy in one launch. + FastPath("sampler_store"), + # The overlap executor loop on the execution stream. + FastPath("execution_stream_loop"), + # Freed host pages returned to the operating system lazily on decode ranks. + FastPath("host_page_retention", depends_on=("post_attention_hop",)), +) + +FAST_PATH_NAMES: Tuple[str, ...] = tuple(p.name for p in FAST_PATHS) +_BY_NAME: Mapping[str, FastPath] = {p.name: p for p in FAST_PATHS} + + +def _dependents() -> Dict[str, FrozenSet[str]]: + """Every path that depends on each path, directly or transitively.""" + direct: Dict[str, List[str]] = {name: [] for name in FAST_PATH_NAMES} + for path in FAST_PATHS: + for base in path.depends_on: + direct[base].append(path.name) + closure = {} + for name in FAST_PATH_NAMES: + seen, todo = set(), list(direct[name]) + while todo: + dependent = todo.pop() + if dependent not in seen: + seen.add(dependent) + todo.extend(direct[dependent]) + closure[name] = frozenset(seen) + return closure + + +_DEPENDENTS = _dependents() + + +def parse_disable(raw: Optional[str], served: Iterable[str] = FAST_PATH_NAMES) -> FrozenSet[str]: + """The fast-path names a ``DISABLE`` value lists (case and blanks ignored). + + ``served`` is the fast paths of the target that reads the value (every fast path by default). + + Raises: + ValueError: naming every entry that is not one of ``served``. + """ + if not raw: + return frozenset() + names = {tok.strip().lower() for tok in raw.split(",") if tok.strip()} + served = frozenset(served) + unknown = sorted(names - served) + if unknown: + raise ValueError( + f"{DISABLE_ENV} names {unknown}, which are not fast paths of this target; " + f"the fast paths are {[name for name in FAST_PATH_NAMES if name in served]}" + ) + return frozenset(names) + + +def disable_closure(names: FrozenSet[str]) -> FrozenSet[str]: + """``names`` and every path that depends on one of them.""" + closed = set(names) + for name in names: + closed |= _DEPENDENTS[name] + return frozenset(closed) + + +#: Every knob of this target set in an environment, as (name, value) sorted by name. +Knobs = Tuple[Tuple[str, str], ...] + + +def read_knobs(environ: Optional[Mapping[str, str]] = None) -> Knobs: + """Every knob of the family set in ``environ`` (the process environment by default).""" + env = os.environ if environ is None else environ + return tuple(sorted((k, v) for k, v in env.items() if k.startswith(KNOB_PREFIX))) + + +def knob_digest(knobs: Knobs) -> str: + """A short digest of ``knobs`` for the log line.""" + return hashlib.sha256(repr(knobs).encode()).hexdigest()[:16] + + +@dataclass(frozen=True) +class Plan: + """The composition this process resolved: which fast paths run, and under which mode.""" + + mode: ModelingV2Mode + # Paths the DISABLE knob names. + disabled: FrozenSet[str] + # Paths turned off only because a path they depend on is disabled. + disabled_dependents: FrozenSet[str] + # Every knob of this target set in the environment, as (name, value). + knobs: Knobs + # Whether the fast paths run their diagnostics (TRTLLM_MODELING_V2_VALIDATE). + validate: bool = False + # What the target that resolved the plan serves (``TargetPaths``); None serves every path. + target: Optional["TargetPaths"] = None + + @property + def require(self) -> bool: + return self.mode is ModelingV2Mode.REQUIRE + + @property + def served(self) -> FrozenSet[str]: + """The fast paths the target serves; the others are off in every plan of the target.""" + return frozenset(FAST_PATH_NAMES) if self.target is None else self.target.served + + def enabled(self, name: str) -> bool: + """Whether fast path ``name`` runs in this process. + + Raises: + KeyError: if ``name`` is not a fast path, so a misspelled call site fails. + """ + if name not in _BY_NAME: + raise KeyError(f"{name!r} is not a fast path of this target") + return ( + name in self.served + and name not in self.disabled + and name not in self.disabled_dependents + ) + + def describe(self) -> str: + """One line naming the mode, the number of fast paths that run and what is off.""" + on = sum(self.enabled(name) for name in FAST_PATH_NAMES) + words = f"modeling_v2={self.mode.value}, {on}/{len(self.served)} fast paths" + if self.disabled: + words += f", disabled {sorted(self.disabled)}" + if self.disabled_dependents: + words += f" and their dependents {sorted(self.disabled_dependents)}" + return words + f", knobs {knob_digest(self.knobs)}" + + +#: The knobs that only one fast path reads; a target that does not serve the path refuses them. +KNOB_PATHS: Mapping[str, str] = {ACCUMULATION_ENV: "routed_moe_into_hop"} + + +def _refuse_removed_knobs(knobs: Knobs, served: Iterable[str] = FAST_PATH_NAMES) -> None: + """Raise if a variable with the knob prefix that is not in ``KNOBS`` is set, or a knob of a + fast path outside ``served``, naming each with what to use instead.""" + served = frozenset(served) + refused = [ + f"{name} ({REMOVED_KNOBS.get(name, _FIXED_HINT)})" for name, _ in knobs if name not in KNOBS + ] + refused += [ + f"{name} (this target does not serve {KNOB_PATHS[name]})" + for name, _ in knobs + if name in KNOB_PATHS and KNOB_PATHS[name] not in served + ] + if refused: + raise ValueError(f"this target does not read {'; '.join(refused)}; unset them") + + +def resolve( + mode: ModelingV2Mode, + environ: Optional[Mapping[str, str]] = None, + target: Optional["TargetPaths"] = None, +) -> Plan: + """Read the knobs once and resolve the composition of ``target``; arms ``missed`` for ``mode``. + + ``target`` is what the target serves (every fast path, with the default tables, when None). + + Raises: + ValueError: if ``DISABLE`` names something that is not a fast path the target serves, or + if a variable with the knob prefix that is not in ``KNOBS`` (or a knob of a path the + target does not serve) is set. + """ + knobs = read_knobs(environ) + served = FAST_PATH_NAMES if target is None else target.served + _refuse_removed_knobs(knobs, served) + disabled = parse_disable(dict(knobs).get(DISABLE_ENV), served) + plan = Plan( + mode=mode, + disabled=disabled, + disabled_dependents=disable_closure(disabled) - disabled, + knobs=knobs, + validate=step_contract_enabled(), + target=target, + ) + set_require(plan.require) + logger.info(f"{_LOG_TAG} dispatch: {plan.describe()}") + return plan + + +# -------------------------------------------------------------------------------------------------- +# Rank uniformity +# -------------------------------------------------------------------------------------------------- + +#: Gathers one rank's value from every rank of a ``world_size``-rank group, or returns None +#: (having logged why) when no such group is available. +Allgather = Callable[[Any, int], Optional[List[Any]]] + +#: The key under which ``rank_uniform_inputs`` records the modeling_v2 mode. +MODE_KEY = "modeling_v2" + + +def rank_uniform_inputs( + knobs: Knobs, mode: ModelingV2Mode, environ: Optional[Mapping[str, str]] = None +) -> Knobs: + """What every rank must agree on before the model is built, as (name, value) sorted by name. + + ``knobs``, the shared switches the kernels read to choose a collective or a protocol + (``SHARED_ENV``) and the modeling_v2 validation switch, where set in ``environ`` (the + process environment by default), and ``mode``. + """ + env = os.environ if environ is None else environ + shared = [(name, env[name]) for name in SHARED_ENV + (MODELING_V2_VALIDATE_ENV,) if name in env] + return tuple(sorted(list(knobs) + shared + [(MODE_KEY, mode.value)])) + + +def _world_allgather(value: Any, world_size: int) -> Optional[List[Any]]: + """``value`` from every rank of the session's MPI communicator. + + Returns None, with a warning, without MPI or when the communicator is not + the ``world_size`` ranks of this model: gathering over any other group + would either prove nothing or wait on ranks that never build this model. + """ + from tensorrt_llm._utils import mpi_comm, mpi_disabled + + if mpi_disabled(): + logger.warning( + f"{_LOG_TAG} dispatch: no MPI communicator, so the ranks are not checked for agreement" + ) + return None + comm = mpi_comm() + if comm.Get_size() != world_size: + logger.warning( + f"{_LOG_TAG} dispatch: the MPI communicator has {comm.Get_size()} ranks, not the " + f"{world_size} of this model, so the ranks are not checked for agreement" + ) + return None + return comm.allgather(value) + + +def _rank_table(gathered: Sequence[Knobs]) -> str: + knob_names = sorted({name for knobs in gathered for name, _ in knobs}) + lines = [] + for rank, knobs in enumerate(gathered): + values = dict(knobs) + knob_words = ", ".join(f"{name}={values.get(name, '')!r}" for name in knob_names) + lines.append(f" rank {rank}: knobs {knob_digest(knobs)} ({knob_words or 'none set'})") + return "\n".join(lines) + + +def check_rank_uniform( + knobs: Knobs, world_size: int, allgather: Optional[Allgather] = None +) -> None: + """Raise on every rank unless every rank read the same knob values. + + Called once, at construction, on every rank, before the knobs are parsed: + the allgather is itself a collective, so a rank must not raise before it. + ``knobs`` is ``rank_uniform_inputs`` of the knobs. ``allgather`` defaults to + the session's MPI communicator. + + Raises: + RuntimeError: with one line per rank, when the ranks disagree. + """ + gathered = (allgather or _world_allgather)(knobs, world_size) + if gathered is None or all(peer == knobs for peer in gathered): + return + raise RuntimeError( + f"{_LOG_TAG} the ranks read different knob values, so they would resolve different " + "fast paths; a collective fast path would wait forever on the ranks that skip it. " + f"Set the {KNOB_PREFIX}* variables, {', '.join(SHARED_ENV)}, " + f"{MODELING_V2_VALIDATE_ENV} and {MODE_KEY} identically on every rank (before the " + f"ranks start):\n{_rank_table(gathered)}" + ) + + +#: The fast-path state of a built model, module by module in construction order: per module, +#: its path and, per fast-path object attached to it, the attribute, the object's class and its +#: rank-independent flags. +AdoptedState = Tuple[Tuple[str, Tuple[Tuple[str, ...], ...]], ...] + + +def check_rank_uniform_adopted( + state: AdoptedState, world_size: int, allgather: Optional[Allgather] = None +) -> None: + """Raise on every rank unless every rank attached the same fast-path state to its modules. + + Called once, on every rank, after the fast paths are attached; a fast path that engaged on + some ranks only would run its collective on those ranks alone. + + Raises: + RuntimeError: naming the first module whose state differs, with one line per rank. + """ + gathered = (allgather or _world_allgather)(state, world_size) + if gathered is None or all(peer == state for peer in gathered): + return + depth = max(len(peer) for peer in gathered) + first = next( + i + for i in range(depth) + if len({peer[i] if i < len(peer) else None for peer in gathered}) > 1 + ) + lines = [] + for rank, peer in enumerate(gathered): + lines.append(f" rank {rank}: {peer[first] if first < len(peer) else ''}") + raise RuntimeError( + f"{_LOG_TAG} the ranks attached different fast-path state to their modules; a " + "collective fast path would wait forever on the ranks that skip it. First module that " + "differs:\n" + "\n".join(lines) + ) + + +# -------------------------------------------------------------------------------------------------- +# Fallback policy +# -------------------------------------------------------------------------------------------------- + + +class FastPathUnavailable(RuntimeError): + """A step inside a fast path's shapes could not take it under ``modeling_v2='require'``.""" + + +_lock = threading.Lock() +_require = [False] +_logged: set = set() +_missed: Dict[str, int] = {} + + +def set_require(require: bool) -> None: + """Whether ``missed`` raises (``modeling_v2='require'``); set by ``resolve``.""" + _require[0] = bool(require) + + +def require_mode() -> bool: + return _require[0] + + +def _log_once(level: Callable[[str], None], path: str, reason: str, message: str) -> None: + key = (path, reason) + with _lock: + if key in _logged: + return + _logged.add(key) + level(message) + + +def missed(path: str, rows: Optional[int], reason: str, *, per_call: bool = False) -> None: + """``path`` takes this step's shape but cannot run: raise under require, else log once and count. + + A miss decided per call (``per_call``) never raises: a step that is already being served + falls back, and the miss is counted and logged once. The caller falls back after it returns. + """ + at = "" if rows is None else f" at {rows} rows" + if _require[0] and not per_call: + raise FastPathUnavailable( + f"modeling_v2='require': the GLM-5.2 decode fast path {path!r} takes this step{at} " + f"but cannot run: {reason}" + ) + with _lock: + _missed[path] = _missed.get(path, 0) + 1 + _log_once( + logger.warning, + path, + reason, + f"{_LOG_TAG} {path}/missed{at}: {reason}; falling back (logged once per reason)", + ) + + +def uncovered(path: str, reason: str) -> None: + """The step is outside ``path``'s shapes: log once per (path, reason); the next tier runs.""" + _log_once(logger.info, path, reason, f"{_LOG_TAG} {path}/uncovered: {reason} (logged once)") + + +def missed_counts() -> Dict[str, int]: + """Misses per fast path in this process: per-call misses in every mode, the others under + ``auto`` (``require`` raises on them instead).""" + with _lock: + return dict(_missed) + + +def reset_fallback_state() -> None: + """Forget the logged reasons, the miss counts and the installed plan.""" + with _lock: + _logged.clear() + _missed.clear() + _INSTALLED[0] = None + _VALIDATE[0] = False + _TARGET[0] = DEFAULT_TARGET + classify_decline.cache_clear() + + +# -------------------------------------------------------------------------------------------------- +# Phases +# -------------------------------------------------------------------------------------------------- + + +class Phase(enum.Enum): + """Which target takes a model step.""" + + PREFILL = "prefill" + DECODE = "decode" + + +_PHASE: List[Optional[Phase]] = [None] + + +class _Step: + """The scope of one model step under a phase; restores the enclosing phase on exit.""" + + __slots__ = ("_phase", "_previous") + + def __init__(self, phase: Phase) -> None: + self._phase = phase + self._previous: Optional[Phase] = None + + def __enter__(self) -> None: + self._previous = _PHASE[0] + _PHASE[0] = self._phase + + def __exit__(self, *exc_info: Any) -> None: + _PHASE[0] = self._previous + + +def step(phase: Phase) -> _Step: + """A context manager that runs the enclosed model step under ``phase``; scopes nest.""" + return _Step(phase) + + +def phase() -> Optional[Phase]: + """The phase of the model step being run, or None outside a step.""" + return _PHASE[0] + + +# -------------------------------------------------------------------------------------------------- +# Decline classes +# -------------------------------------------------------------------------------------------------- + + +class Decline(enum.Enum): + """What a fast path's decline means.""" + + #: The path is turned off: neither a miss nor outside its shapes. + DISABLED = "disabled" + #: Outside the path's shapes (a layer it does not serve, a row count, a step kind). + UNCOVERED = "uncovered" + #: Decided once, at construction or in the warm-up, inside the path's layers, with a + #: precondition absent: ``missed``, which raises under ``require``. + MISSED = "missed" + #: Decided per call, with a resource absent: counted and logged once, never raised. + MISSED_PER_CALL = "missed per call" + + +_TARGET_LAYERS = frozenset(range(78)) +_DRAFT_LAYERS = frozenset({78}) +# The decoder layers that run the full indexer: the three dense layers and every fourth layer +# from layer 6. The other target layers reuse an earlier layer's top-k selection. +_FULL_INDEXER_LAYERS = frozenset((0, 1, 2) + tuple(range(6, 78, 4))) +_SHARED_INDEXER_LAYERS = _TARGET_LAYERS - _FULL_INDEXER_LAYERS + +#: Per fast path and per layer-keyed capability key family (the key without its layer index), +#: the layers on which the path engages on every rank when this target serves its checkpoint. +#: A decline on a layer outside its target's cells is outside the path's shapes. +LAYER_CELLS: Mapping[str, Mapping[str, FrozenSet[int]]] = { + "absorbed_query": {"layer": _TARGET_LAYERS}, + "draft_absorbed_query": {"layer": _DRAFT_LAYERS}, + "shared_indexer_absorbed_query": {"layer": _SHARED_INDEXER_LAYERS}, + # Only the shared-indexer layer that the last dense layer feeds. + "fused_kv_indexer_projection": {"layer": frozenset({3})}, + "indexer_layer_front": {"target_layer": _FULL_INDEXER_LAYERS, "draft_layer": _DRAFT_LAYERS}, + # Boundaries into the shared-indexer layers that a MoE layer feeds. + "deferred_norm": {"scope[shared]_into_layer": _SHARED_INDEXER_LAYERS - {3}}, + "deferred_norm_evict_hints": {"projection_layer": _SHARED_INDEXER_LAYERS - {3}}, + "sparse_mla_decode": { + "layer": _TARGET_LAYERS | _DRAFT_LAYERS, + "call:layer": _TARGET_LAYERS | _DRAFT_LAYERS, + }, + "attention_output_tail": {"layer": _TARGET_LAYERS}, + "oproj_weight_staging": {"layer": _TARGET_LAYERS}, +} + +#: Reason fragments that turn a path off rather than decline it. +DISABLED_REASONS: Tuple[str, ...] = ("DISABLE",) +#: Reason fragments that name a layer kind or a composition the path does not serve. +LAYER_KIND_REASONS: Tuple[str, ...] = ( + "target layers only", + "draft (MTP) layers only", + "full indexer layer", + "shared-indexer layer", + "layers not selected (", + "draft or unindexed layer", + "not a target layer", + "is a full-indexer layer", + "not a DSA layer", + "not MLA", + "needs the post_attention_hop hop", + "no post_attention_hop hop", + "post_attention_hop without sparse_mla_decode", +) +#: Reason fragments (any case) of a per-call decline that name a missing resource. +RESOURCE_REASONS: Tuple[str, ...] = ( + "workspace", + "compile failed", + "not compiled", + "not loaded", + "not built", + "failed", + "unavailable", + "not warmed up", + "before CUDA graph capture", + "packet does not match", + "does not forward", + "another rank", + "MNNVL", + "fabric", +) + + +def _any_of(fragments: Iterable[str], flags: int = 0) -> "re.Pattern[str]": + return re.compile("|".join(map(re.escape, fragments)), flags) + + +_DISABLED_RE = _any_of(DISABLED_REASONS) +_LAYER_KIND_RE = _any_of(LAYER_KIND_REASONS) +_RESOURCE_RE = _any_of(RESOURCE_REASONS, re.IGNORECASE) + + +@dataclass(frozen=True) +class DeclineSite: + """Capability keys of a fast path that are decided alike. + + ``key`` matches a whole key; a key that names a layer captures its index as ``layer``, and + ``cells`` names the ``LAYER_CELLS`` family of those layers. A ``per_call`` site decides on + every call; the others decide once, at construction or in the warm-up. ``uncovered`` + matches reasons that are outside the path's shapes at this site. + """ + + path: str + key: "re.Pattern[str]" + per_call: bool = False + cells: Optional[str] = None + uncovered: Optional["re.Pattern[str]"] = None + + +def _site( + path: str, + key: str, + per_call: bool = False, + cells: Optional[str] = None, + uncovered: Optional[str] = None, +) -> DeclineSite: + return DeclineSite( + path, + re.compile(key), + per_call, + cells, + None if uncovered is None else re.compile(re.escape(uncovered)), + ) + + +_LAYER = r"(?P\d+)" +_EXPERT_PREFETCH = r"expert_prefetch\[.*\]" + +#: Every capability key the fast paths log, by where and when it is decided. +DECLINE_KINDS: Tuple[DeclineSite, ...] = ( + _site("absorbed_query", "form"), + _site("absorbed_query", f"layer{_LAYER}", cells="layer"), + _site("draft_absorbed_query", f"layer{_LAYER}", cells="layer"), + _site("shared_indexer_absorbed_query", f"layer{_LAYER}", cells="layer"), + _site("fused_kv_indexer_projection", f"layer{_LAYER}", cells="layer"), + _site("indexer_layer_front", f"target_layer{_LAYER}", cells="target_layer"), + _site("indexer_layer_front", f"draft_layer{_LAYER}", cells="draft_layer"), + _site("indexer_layer_front", "rope"), + _site("indexer_layer_front", r"cp_indexer_decode_.+", per_call=True), + _site("deferred_norm", "deferred_hop"), + _site( + "deferred_norm", rf"scope\[shared\]_into_layer{_LAYER}", cells="scope[shared]_into_layer" + ), + _site("deferred_norm", r"deferred_hop_routed\[.*\]"), + _site("deferred_norm", f"projection_layer{_LAYER}", cells="projection_layer"), + _site("deferred_norm_evict_hints", "deferred_hop"), + _site("deferred_norm_evict_hints", r"deferred_hop_routed\[.*\]"), + _site("deferred_norm_evict_hints", f"projection_layer{_LAYER}", cells="projection_layer"), + _site("rope_kv_append", "rope|scatter_defer|bridge", per_call=True), + _site("indexer_topk", "indexer_decode", per_call=True), + _site("cross_rank_topk", "cp_indexer_decode", per_call=True), + _site("cross_rank_topk_cluster_merge", "cp_indexer_decode", per_call=True), + _site("cross_rank_topk_global_emit", "cp_indexer_decode", per_call=True), + _site("sparse_mla_decode", f"layer{_LAYER}", cells="layer"), + _site("sparse_mla_decode", r"request:.*", per_call=True), + _site("sparse_mla_decode", f"call:layer{_LAYER}", per_call=True, cells="call:layer"), + _site("attention_output_tail", f"layer{_LAYER}", cells="layer"), + _site("attention_output_tail_prefetch", f"layer{_LAYER}", cells="layer"), + _site("attention_output_tail_prefetch", "schedule"), + _site("attention_output_into_hop", f"layer{_LAYER}", cells="layer"), + _site("attention_output_into_hop", "into_hop_transport"), + _site("oproj_weight_staging", f"layer{_LAYER}", cells="layer"), + _site("post_attention_hop", r"post_attn_hop(:.*)?"), + _site("hop_routing", "routing_hop"), + _site("hop_routing", r"routed_moe_fork rows=.*", per_call=True), + _site("hop_routing", _EXPERT_PREFETCH), + _site("router_gemm", r"router:.*"), + _site("routed_moe", r"rows=.*", per_call=True), + _site("routed_moe", _EXPERT_PREFETCH), + _site("routed_moe_into_hop", r"rows=.*", per_call=True), + _site("routed_moe_into_hop", r"post_moe_routed_input.*"), + _site("shared_expert_prefetch", r"routed rows=.*", per_call=True), + _site("shared_expert_prefetch", _EXPERT_PREFETCH), + _site("post_moe_hop", r"cluster_kernel(_compile)?|post_moe_hop\[.*\]"), + _site("cross_layer_expert_prefetch", r"expert_prefetch|fused_predict|xc_prefetch|link\[.*\]"), + _site("cross_layer_expert_prefetch", r"spec_experts(_declined)?", per_call=True), + _site("dense_mlp", r"dense_mlp:.*", per_call=True), + _site("draft_front", "draft_front"), + _site("draft_front", r"draft_front_call(_fallback)?", per_call=True), + _site("draft_routed_moe", r"rows=.*", per_call=True), + _site("draft_live_rows", r"live_row_hop(:.*)?"), + _site("draft_live_rows", r"live_row_hop_call(:.*)?|step0:.*", per_call=True), + _site("draft_tail", r"draft_tail\[tp\d+\]"), + _site("draft_tail", r"draft_tail\[(step: .*|graph|eager)\]", per_call=True), + _site("draft_loop", r"(rows|draft_step|finalize|draft_loop):.*", per_call=True), + _site("draft_loop", "draft_embedding"), + _site("decode_prologue", r"target_(prologue|embedding)(\[.*\])?|take_embedding", per_call=True), + _site("deferred_input_prep", r"prep(\[.*\])?", per_call=True), + # Before the overlap executor loop declares its domain every call declines; the decision + # is final only once the loop has declared it. + _site( + "deferred_input_prep", + "loop_domain", + uncovered="no overlap executor loop declared its domain", + ), + _site("deferred_input_prep", "engaged"), + _site("sampler_store", "engaged"), + _site("execution_stream_loop", "engaged"), + _site("host_page_retention", "host_policy"), +) + +_SITES_BY_PATH: Mapping[str, Tuple[DeclineSite, ...]] = { + name: tuple(s for s in DECLINE_KINDS if s.path == name) for name in FAST_PATH_NAMES +} + + +@dataclass(frozen=True) +class TargetPaths: + """What one target of the family serves of the fast paths, and where their code lives. + + ``served`` is the fast paths the target runs: a path outside it is off in every plan of the + target, ``DISABLE`` may not name it and a knob only it reads is refused. ``layer_cells`` is + ``LAYER_CELLS`` for the target's layers and shapes; ``uncovered`` lists, per path, reason + fragments of declines that are outside the target's shapes at every site of the path. + ``kernel_modules`` names the module of each path whose kernel package is the target's own; + every other path's package is the target's (``KERNELS_PACKAGE``). + + Raises: + ValueError: if a name is not a fast path, or a served path extends one the target does not + serve. + """ + + served: FrozenSet[str] = frozenset(FAST_PATH_NAMES) + layer_cells: Mapping[str, Mapping[str, FrozenSet[int]]] = field( + default_factory=lambda: LAYER_CELLS + ) + uncovered: Mapping[str, Tuple[str, ...]] = field(default_factory=dict) + kernel_modules: Mapping[str, str] = field(default_factory=dict) + + def __post_init__(self) -> None: + named = set(self.served) | set(self.layer_cells) | set(self.uncovered) + unknown = sorted((named | set(self.kernel_modules)) - set(FAST_PATH_NAMES)) + if unknown: + raise ValueError(f"{unknown} are not fast paths of the family") + lacking = sorted( + f"{name} extends {base}" + for name in self.served + for base in _BY_NAME[name].depends_on + if base not in self.served + ) + if lacking: + raise ValueError(f"a served fast path extends one the target does not serve: {lacking}") + + +#: Every fast path, with the default tables. +DEFAULT_TARGET = TargetPaths() +# The tables of the installed plan's target. +_TARGET: List[TargetPaths] = [DEFAULT_TARGET] + + +def _decline_site(path: str, key: str) -> Optional[Tuple[DeclineSite, "re.Match[str]"]]: + """The site of ``path`` whose key pattern matches all of ``key``, with the match.""" + for site in _SITES_BY_PATH.get(path, ()): + match = site.key.fullmatch(key) + if match is not None: + return site, match + return None + + +@functools.lru_cache(maxsize=4096) +def classify_decline(path: str, key: str, reason: str) -> Decline: + """What a decline of fast path ``path`` at capability ``key`` for ``reason`` means. + + In order: a reason that turns the path off is ``DISABLED``; a key no site of the path + matches is ``UNCOVERED``; a per-call decline is ``MISSED_PER_CALL`` when its reason names a + missing resource and ``UNCOVERED`` otherwise; a decline decided once is ``UNCOVERED`` when + its reason names a layer kind or a composition the path does not serve, or its layer is + outside the path's cells, and ``MISSED`` otherwise. The cells and the target's uncovered + reasons are the installed plan's target's (``TargetPaths``). Memoized: a repeated decline + costs one lookup. + """ + if _DISABLED_RE.search(reason): + return Decline.DISABLED + found = _decline_site(path, key) + if found is None: + return Decline.UNCOVERED + site, match = found + if site.uncovered is not None and site.uncovered.search(reason): + return Decline.UNCOVERED + if any(fragment in reason for fragment in _TARGET[0].uncovered.get(path, ())): + return Decline.UNCOVERED + if site.per_call: + return Decline.MISSED_PER_CALL if _RESOURCE_RE.search(reason) else Decline.UNCOVERED + if _LAYER_KIND_RE.search(reason): + return Decline.UNCOVERED + if site.cells is not None: + cells = _TARGET[0].layer_cells.get(path, {}).get(site.cells, frozenset()) + if int(match["layer"]) not in cells: + return Decline.UNCOVERED + return Decline.MISSED + + +def _record_decline(path: str, key: str, reason: str) -> None: + kind = classify_decline(path, key, reason) + if kind is Decline.MISSED: + missed(path, None, f"{key}: {reason}" if reason else key) + elif kind is Decline.MISSED_PER_CALL: + missed(path, None, reason or key, per_call=True) + + +# -------------------------------------------------------------------------------------------------- +# Selection +# -------------------------------------------------------------------------------------------------- + +#: The parts of the decode step's input preparation, each a fast path of its own: the step-input +#: stager, the sampler's store and the overlap loop on the execution stream. +PART_PATHS: Mapping[str, str] = { + "prep": "deferred_input_prep", + "store": "sampler_store", + "stream": "execution_stream_loop", +} +_PART_ONLY_PATHS = frozenset(PART_PATHS.values()) - {PART_PATHS["prep"]} + +#: Every fast path except the parts that only ``deferred_input_prep`` brings along. +SELECTOR_PATHS: Tuple[str, ...] = tuple(n for n in FAST_PATH_NAMES if n not in _PART_ONLY_PATHS) + +#: The family's kernel packages, one directory per fast path under ``kernels/``. +KERNELS_PACKAGE = f"{__package__}.kernels" +#: The fast paths whose code lives in the kernel directory of the path they extend. +KERNEL_DIRECTORY: Mapping[str, str] = { + "draft_absorbed_query": "absorbed_query", + "shared_indexer_absorbed_query": "absorbed_query", + "cross_rank_topk": "indexer_topk", + "attention_output_tail_prefetch": "attention_output_tail", + "attention_output_into_hop": "attention_output_tail", + "routed_moe_into_hop": "routed_moe", + "sampler_store": "deferred_input_prep", + "execution_stream_loop": "deferred_input_prep", +} + + +def _bases() -> Dict[str, FrozenSet[str]]: + """Every path each path extends, directly or transitively.""" + closure: Dict[str, FrozenSet[str]] = {} + + def visit(name: str) -> FrozenSet[str]: + if name not in closure: + bases = set() + for base in _BY_NAME[name].depends_on: + bases |= {base} | visit(base) + closure[name] = frozenset(bases) + return closure[name] + + for name in FAST_PATH_NAMES: + visit(name) + return closure + + +#: The paths each path has no effect without (the transitive closure of ``depends_on``). +REQUIRED_FEATURES: Mapping[str, FrozenSet[str]] = { + name: bases for name, bases in _bases().items() if bases +} + +# Numerical variants (a variant has its own name). Per fast path with variants: the knob that +# selects one (None: a single variant) and the variant names, default first. A path without an +# entry has one numerical form, equal to the generic path's or held to a frozen error budget. +NUMERICAL_VARIANTS: Dict[str, Tuple[Optional[str], Tuple[str, ...]]] = { + # Exact top-k over each page owner's indexer K, identical on every rank. + "cross_rank_topk": (None, ("rank-consistent",)), + # The cluster merge's rows are byte-identical to the cross-rank top-k's (same set and order). + "cross_rank_topk_cluster_merge": (None, ("rank-consistent",)), + # Its id rows are the cluster merge's; its pool and copy rows are the values of the + # rope/KV-append top-k bridge. + "cross_rank_topk_global_emit": (None, ("rank-consistent",)), + # What the routed experts write into the post-MoE hop: fp32 sums the top-k contributions in + # another order (nondeterministic from run to run); identity is bit-exact. + "routed_moe_into_hop": (ACCUMULATION_ENV, ("fp32", "identity")), + # The next layer's RMSNorm deferred into its projection (changes rounding), only on boundaries + # into shared-indexer layers: full-indexer layers keep the post-MoE hop's normalization. + "deferred_norm": (None, ("shared",)), + # One projection launch for the fused-A GEMM and the indexer key/weights GEMM: kva keeps the + # fused-A order (bit-identical); kw (full-indexer and draft layers) takes the same mma.sync + # order instead of the cuBLAS split-K GEMM (rounds differently). This target runs it on the + # dense-fed shared-indexer layer, which computes kva only. + "fused_kv_indexer_projection": (None, ("kw-fused-a-order",)), +} + +#: Prefixes of the set variables the record line lists (``side_selectors``). +SIDE_SELECTOR_PREFIXES = (KNOB_PREFIX,) +#: Shared TensorRT-LLM switches that the kernels also read. +SHARED_ENV = ("TLLM_TEST_MNNVL", "TRTLLM_ENABLE_PDL", "TRTLLM_FORCE_MNNVL_AR") + + +def feature_variants( + paths: Iterable[str], environ: Optional[Mapping[str, str]] = None +) -> Dict[str, str]: + """The numerical variant of every path in ``paths`` that has some. + + Raises: + ValueError: if a knob names an unknown variant of a path in ``paths``. + """ + env = os.environ if environ is None else environ + variants = {} + for path in sorted(paths): + if path not in NUMERICAL_VARIANTS: + continue + name, choices = NUMERICAL_VARIANTS[path] + raw = env.get(name) if name else None + value = (raw or "").strip().lower() or choices[0] + if value not in choices: + raise ValueError( + f"{name}={raw!r} is not a {path} variant; known variants are {list(choices)} " + f"(unset selects {choices[0]!r})" + ) + variants[path] = value + return variants + + +def side_selectors(environ: Optional[Mapping[str, str]] = None) -> Dict[str, str]: + """Every set knob of the family.""" + env = os.environ if environ is None else environ + return {name: env[name] for name in sorted(env) if name.startswith(SIDE_SELECTOR_PREFIXES)} + + +def shared_env(environ: Optional[Mapping[str, str]] = None) -> Dict[str, str]: + """The set shared TensorRT-LLM switches that the kernels read.""" + env = os.environ if environ is None else environ + return {name: env[name] for name in SHARED_ENV if name in env} + + +def feature_record(paths: Iterable[str], environ: Optional[Mapping[str, str]] = None) -> str: + """The one-line record of a selection. + + ``paths=`` lists every path as ``path`` or ``path[variant]``; ``side_selectors=`` and + ``shared_env=`` are JSON objects of the set variables. + + Raises: + ValueError: as ``feature_variants``. + """ + variants = feature_variants(paths, environ) + listed = ",".join(f"{p}[{variants[p]}]" if p in variants else p for p in sorted(paths)) + + def _json(values: Dict[str, str]) -> str: + return json.dumps(values, sort_keys=True, separators=(",", ":")) + + return ( + f"paths={listed or ''} " + f"side_selectors={_json(side_selectors(environ))} " + f"shared_env={_json(shared_env(environ))}" + ) + + +def _log_info(message: str) -> None: + logger.info(message) + + +# The fast paths of the installed plan; None while no plan is installed. +_INSTALLED: List[Optional[FrozenSet[str]]] = [None] +# Whether the installed plan runs the fast paths' diagnostics. +_VALIDATE: List[bool] = [False] + + +def install(plan: Plan, environ: Optional[Mapping[str, str]] = None) -> None: + """Answer ``is_enabled`` from ``plan`` from now on, and log its record line once. + + Raises: + ValueError: as ``feature_variants``, for a knob that names an unknown variant of an + enabled path. + """ + paths = frozenset(n for n in FAST_PATH_NAMES if plan.enabled(n)) + record = feature_record(paths, environ) + _INSTALLED[0] = paths + _VALIDATE[0] = plan.validate + _TARGET[0] = DEFAULT_TARGET if plan.target is None else plan.target + classify_decline.cache_clear() + _log_info(f"{_LOG_TAG} dispatch: record {record}") + + +def installed() -> Optional[FrozenSet[str]]: + """The fast paths of the installed plan, or None.""" + return _INSTALLED[0] + + +def enabled_features() -> FrozenSet[str]: + """The fast paths that run in this process: the installed plan's, else none.""" + plan_paths = _INSTALLED[0] + return plan_paths if plan_paths is not None else frozenset() + + +def validate_mode() -> bool: + """Whether the installed plan runs the fast paths' diagnostics (False without a plan).""" + return _VALIDATE[0] + + +def is_enabled(path: str) -> bool: + """Whether fast path ``path`` runs in this process. + + Raises: + ValueError: if ``path`` is not a fast path of the family. + """ + if path not in _BY_NAME: + raise ValueError(f"{path!r} is not a fast path of this target") + return path in enabled_features() + + +# Entry points per fast path. The kernel package of a path registers them when it is imported; +# a call site looks one up with ``get_impl`` only after ``is_enabled`` returned True. +_REGISTRY: Dict[str, Dict[str, Callable]] = {} + +# Per-rank capability decisions, logged once per (path, key). +_CAPABILITY_LOGGED: Dict[Tuple[str, str], bool] = {} + + +def register(path: str, name: str, fn: Callable) -> Callable: + """Register ``fn`` as entry point ``name`` of fast path ``path``. + + Raises: + ValueError: if ``path`` is not a fast path of the family. + """ + if path not in _BY_NAME: + raise ValueError(f"{path!r} is not a fast path of this target") + _REGISTRY.setdefault(path, {})[name] = fn + return fn + + +def get_impl(path: str, name: str) -> Callable: + """Import the kernel package of ``path`` on first use and return its entry ``name``. + + The package is the installed plan's target's own (``TargetPaths.kernel_modules``) or the + target's. + + Raises: + KeyError: if the package does not register ``name`` for ``path``. + """ + if path not in _REGISTRY: + module = _TARGET[0].kernel_modules.get(path) + importlib.import_module(module or f"{KERNELS_PACKAGE}.{KERNEL_DIRECTORY.get(path, path)}") + entries = _REGISTRY.get(path, {}) + if name not in entries: + raise KeyError( + f"fast path {path!r} does not register {name!r}; registered: {sorted(entries)}" + ) + return entries[name] + + +def log_capability(path: str, key: str, eligible: bool, reason: str = "") -> bool: + """Log a per-rank capability decision once per (path, key) and return ``eligible``. + + Every decline, logged or not, is also classified (``classify_decline``): a miss decided + once goes to ``missed``, which raises under ``require``; a miss decided per call is counted + and logged once; a decline outside the path's shapes or of a path turned off adds nothing + to the capability line. + """ + if (path, key) not in _CAPABILITY_LOGGED: + _CAPABILITY_LOGGED[(path, key)] = eligible + state = "engaged" if eligible else "native fallback" + suffix = f" ({reason})" if reason else "" + _log_info(f"{_LOG_TAG} {path}/{key}: {state}{suffix}") + if not eligible: + _record_decline(path, key, reason) + return eligible diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/__init__.py new file mode 100644 index 000000000000..45b6ee48d0ee --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kernels of the target's fast paths, one package per fast path.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/absorbed_query/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/absorbed_query/__init__.py new file mode 100644 index 000000000000..11a8ee3db664 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/absorbed_query/__init__.py @@ -0,0 +1,48 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""absorbed_query: target query/projection region (the query_front contract, token-wise part of the +index_selection contract). + +Registered entry points (``dispatch.get_impl("absorbed_query", name)``): + +* ``dsa_proj(mla, position_ids, hidden_states, attn_metadata)`` -> + ``Optional[AbsorbedQueryPacket]`` for target layers (``absorbed_query``); ``None`` + means "run the native path". +* ``mtp_dsa_proj(mla, position_ids, hidden_states, attn_metadata)`` -> the + same packet for the one-model MTP draft layer at every draft step + (package ``draft_absorbed_query``, which extends absorbed_query to the draft layer); ``None`` for + target layers and outside the support domain. +* ``shared_dsa_proj(mla, position_ids, hidden_states, attn_metadata)`` -> + the preabsorbed packet of a shared-indexer target layer from one kernel + (package ``shared_indexer_absorbed_query``, which extends absorbed_query's shared-layer path); + ``None`` for full-indexer and draft layers and outside the support domain. +""" + +from ...dispatch import register +from .query_proj import AbsorbedQueryPacket, AbsorbedQueryProjection, dsa_proj, mtp_dsa_proj +from .shared_indexer import SharedIndexerProjection, shared_dsa_proj + +register("absorbed_query", "dsa_proj", dsa_proj) +register("absorbed_query", "mtp_dsa_proj", mtp_dsa_proj) +register("absorbed_query", "shared_dsa_proj", shared_dsa_proj) + +__all__ = [ + "AbsorbedQueryPacket", + "AbsorbedQueryProjection", + "SharedIndexerProjection", + "dsa_proj", + "mtp_dsa_proj", + "shared_dsa_proj", +] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/absorbed_query/kernels.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/absorbed_query/kernels.py new file mode 100644 index 000000000000..be6c61752cd6 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/absorbed_query/kernels.py @@ -0,0 +1,1764 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Triton kernels of the absorbed_query/projection package. + +``qproj_kernel`` consumes the native fused ``kv_a_proj_with_mqa`` output +``kva = [q_a | compressed_kv | k_pe]`` of ``R <= BLOCK_M`` token rows (and, +with ``PRENORM``, the native ``q_a_layernorm`` output) and produces, in one +launch: + +* ``q = q_b_proj(rmsnorm(q_a))`` ``[R, H * 256]`` (BF16, the native ``q``), +* optionally ``qr = rmsnorm(q_a)`` ``[R, 2048]`` for the native indexer, +* ``latent = [rmsnorm(compressed_kv) | k_pe]`` ``[R, 576]``, +* optionally the absorbed query ``fused_q[:, h, :512] = q_nope[:, h] @ W_UK[h]`` + written into the caller's ``[R, H, 576]`` buffer. + +``qproj_indexer_kernel`` additionally produces the token-wise indexer +outputs of full indexer layers. + +Numerics follow the native chain bit for bit on the default path +(``PRENORM``, the native kv_a GEMM, ``q_a_layernorm``, fused ``wk`` / +``weights_proj`` GEMM and ``k_norm`` as bridges): ``q`` and the indexer +``wq_b`` output are FP32-accumulated MMA chains in ascending K order (equal to +the native GEMMs), the ``compressed_kv`` RMSNorm reproduces the native +FlashInfer kernel's reduction order (``_fi_rmsnorm_rows_exact``), the +indexer RoPE its FMA pattern (``_rope_pairs``) and the FP8 tail the native +``fusedCatFp8Kernel``; ``q`` is rounded to BF16 before the absorption (native +intermediate). The experimental paths (own hidden GEMV, prefetching kernel's +in-kernel q_a norm, ``PRENORM=False``) are not bit-exact. + +Program roles, in program-id order (each kernel documents its own list): +``q_b`` producer tiles ``[0, N_QB)`` (each ``BN`` output columns over the full +``K``; ``TPH = HEAD_DIM // BN`` tiles per head, the first ``NOPE // BN`` of a +head are its no-RoPE tiles), then indexer ``wq_b`` producer tiles, then the +programs that only consume. + +Inter-program synchronization contract +-------------------------------------- +* A producer stores its tile, executes a CTA barrier and then one thread + releases the tile with ``atom.add.release.gpu`` on a per-head counter. A + consumer spins with ``atom.add.acquire.gpu`` (adding 0) until the counter + reaches the number of producer tiles it needs, executes a CTA barrier and + only then reads the tiles, with ``ld.global.cg`` (L2, never a stale L1 line). +* Every counter is per module (per layer) and is zero between launches. The + last consumer of a counter resets it (relaxed exchange) after every + consumer of that counter has read its tiles: for the absorption counters a + second ``done`` counter (``acq_rel``) counts the consumers of a head; an + indexer-head counter has a single consumer. Counters are touched only after + ``griddepcontrol.wait``, so under PDL a launch never observes the counters + of an earlier launch that has not completed. +* Forward progress: every spin-wait waits only on counters released by + programs with a *lower* program id. Absorption job ``j`` (head + ``j // TPH``, ``BC`` latent columns) is run, with ``ABSORB_IN_PRODUCER`` + and ``PRODUCER_JOB_SHIFT == TPH``, by producer ``j + TPH`` (a producer of + the next head) after it has published its own tile, and the last head's + jobs by ``TPH`` extra consumer programs placed after all producers. So the + lowest-numbered unfinished program never waits, and a launch completes + with a single resident program slot when programs are dispatched in id + order, and on any dispatch order when the whole grid is resident at once + (the default pipelined grids, <= 234 programs of <= 96 KB shared memory, + fit on an idle sm_103 GPU; the experimental prefetching kernel's do not). + ``PRODUCER_JOB_SHIFT == 0`` is the legacy layout (producer ``p`` runs job + ``p`` of its own head group), which needs ``NOPE // BN`` co-resident + programs of one head group; it is kept only as a stress-test reference. +* ``DELAY_NS > 0`` (tests only) delays a deterministic subset of producers + before they publish and of consumers before they finish, by up to + ``DELAY_NS`` nanoseconds of ``%globaltimer``, to stress the protocol. +""" + +import triton +import triton.language as tl + + +@triton.jit +def _rms_rcp( + src_ptr, stride_src, rows, rmask, N: tl.constexpr, BLOCK_M: tl.constexpr, BK: tl.constexpr, eps +): + cols = tl.arange(0, BK) + ss = tl.zeros([BLOCK_M], dtype=tl.float32) + for k0 in tl.range(0, N, BK, num_stages=N // BK): + x = tl.load( + src_ptr + rows[:, None] * stride_src + (k0 + cols)[None, :], + mask=rmask[:, None], + other=0.0, + ).to(tl.float32) + ss += tl.sum(x * x, axis=1) + return tl.rsqrt(ss / N + eps) + + +@triton.jit +def _rms_rcp_partials( + ss_ptr, + group: tl.constexpr, + n_parts, + N: tl.constexpr, + BLOCK_M: tl.constexpr, + MAX_PARTS: tl.constexpr, + eps, +): + parts = tl.arange(0, MAX_PARTS) + rows = tl.arange(0, BLOCK_M) + ss = tl.load( + ss_ptr + parts[:, None] * 2 * BLOCK_M + group * BLOCK_M + rows[None, :], + mask=(parts < n_parts)[:, None], + other=0.0, + ) + return tl.rsqrt(tl.sum(ss, axis=0) / N + eps) + + +# -------------------------------------------------------------------------- +# Bit-exact reproduction of the native FlashInfer CuTe-DSL RMSNorm (the +# ``kv_a_layernorm`` / ``q_a_layernorm`` kernels, H in {512, 2048}): one warp +# per row, lane t owns columns 256 * b + 8 * t + v (b < H / 256, v < 8), a +# sequential per-lane fold of the squares in column order (blocks outer, +# columns inner), a butterfly over the 32 lanes (xor offsets 1, 2, 4, 8, 16), +# rstd = rsqrt(sum / H + eps), y = bf16((x * rstd) * (w + 0)). Every add is +# explicit (reshape/split, sums over axes of size 2), so the result does not +# depend on the compiled layout or on the rows per program. The same +# arithmetic as ``shared_indexer._fi_rmsnorm_rows_vec`` (shared_indexer_absorbed_query), checked bit +# for bit against the native kernel on real and wide-range rows. +# -------------------------------------------------------------------------- + + +@triton.jit +def _fi_butterfly32(s, ROWS: tl.constexpr): + """Sum of ``s [ROWS, 32]`` over the lanes in shuffle-xor butterfly order + (offsets 1, 2, 4, 8, 16): a balanced pairwise tree in lane order.""" + s = tl.sum(tl.reshape(s, [ROWS, 16, 2]), axis=2) + s = tl.sum(tl.reshape(s, [ROWS, 8, 2]), axis=2) + s = tl.sum(tl.reshape(s, [ROWS, 4, 2]), axis=2) + s = tl.sum(tl.reshape(s, [ROWS, 2, 2]), axis=2) + return tl.sum(s, axis=1) + + +@triton.jit +def _v8_leaves(x, ROWS: tl.constexpr): + """The 8 columns ``v`` of ``x [ROWS, 32, 8]`` as ``[ROWS, 32]`` tensors, + in ``v`` order (register-level reshape/split, no data movement).""" + e, o = tl.split(tl.reshape(x, [ROWS, 32, 4, 2])) + ee, eo = tl.split(tl.reshape(e, [ROWS, 32, 2, 2])) + oe, oo = tl.split(tl.reshape(o, [ROWS, 32, 2, 2])) + v0, v4 = tl.split(ee) + v2, v6 = tl.split(eo) + v1, v5 = tl.split(oe) + v3, v7 = tl.split(oo) + return v0, v1, v2, v3, v4, v5, v6, v7 + + +@triton.jit +def _fi_rmsnorm_rows_exact( + x_ptr, stride_x, w_ptr, y_ptr, stride_y, rows, rmask, eps, H: tl.constexpr, ROWS: tl.constexpr +): + """``y[rows, :H] = bf16(rmsnorm(x[rows, :H]) * w)`` bit for bit as the + native FlashInfer RMSNorm (see above).""" + lanes = tl.arange(0, 32) + vv = tl.arange(0, 8) + offs = rows[:, None, None] * stride_x + 8 * lanes[None, :, None] + vv[None, None, :] + s = tl.zeros([ROWS, 32], dtype=tl.float32) + for b in tl.static_range(H // 256): + x = tl.load(x_ptr + offs + 256 * b, mask=rmask[:, None, None], other=0.0).to(tl.float32) + v0, v1, v2, v3, v4, v5, v6, v7 = _v8_leaves(x, ROWS) + s = s + v0 * v0 + s = s + v1 * v1 + s = s + v2 * v2 + s = s + v3 * v3 + s = s + v4 * v4 + s = s + v5 * v5 + s = s + v6 * v6 + s = s + v7 * v7 + rstd = tl.rsqrt(_fi_butterfly32(s, ROWS) / H + eps) + cols = tl.arange(0, H) + x = tl.load( + x_ptr + rows[:, None] * stride_x + cols[None, :], mask=rmask[:, None], other=0.0 + ).to(tl.float32) + w = tl.load(w_ptr + cols).to(tl.float32) + 0.0 + y = (x * rstd[:, None]) * w[None, :] + tl.store( + y_ptr + rows[:, None] * stride_y + cols[None, :], y.to(tl.bfloat16), mask=rmask[:, None] + ) + + +@triton.jit +def _stress_delay(pid, salt: tl.constexpr, DELAY_NS: tl.constexpr): + """Test-only skew: spin ``(h / 10) * DELAY_NS`` ns, ``h = (37 pid + salt) + mod 11`` (so some programs do not wait and others wait up to DELAY_NS).""" + if DELAY_NS > 0: + wait_ns = ((pid * 37 + salt) % 11) * (DELAY_NS // 10) + t0 = tl.extra.cuda.globaltimer() + while tl.extra.cuda.globaltimer() - t0 < wait_ns: + pass + + +@triton.jit +def _ld_wuk(a, valid, wuk_ptr, KV_LORA: tl.constexpr, NOPE: tl.constexpr, BC: tl.constexpr): + """``W_UK[h, c0:c0+BC, :NOPE]`` of absorption job ``a`` as [NOPE, BC] + (split 128 + NOPE-128 along NOPE); nothing is read unless ``valid``.""" + chunks_per_head: tl.constexpr = KV_LORA // BC + h = a // chunks_per_head + c0 = (a % chunks_per_head) * BC + cc = tl.arange(0, BC) + w_base = wuk_ptr + h * (KV_LORA * NOPE) + (c0 + cc)[None, :] * NOPE + w_lo = tl.load(w_base + tl.arange(0, 128)[:, None], mask=valid, other=0.0) + w_hi = tl.load(w_base + 128 + tl.arange(0, NOPE - 128)[:, None], mask=valid, other=0.0) + return w_lo, w_hi + + +@triton.jit +def _qb_program( + pid, + kva_ptr, + stride_kva, + gq_ptr, + wqb_ptr, + q_ptr, + qr_ptr, + sync_ptr, + ss_ptr, + n_parts, + rows, + rmask, + eps, + Q_LORA: tl.constexpr, + HEAD_DIM: tl.constexpr, + NOPE: tl.constexpr, + H: tl.constexpr, + BLOCK_M: tl.constexpr, + BN: tl.constexpr, + BK: tl.constexpr, + WRITE_QR: tl.constexpr, + DO_ABSORB: tl.constexpr, + SS_PARTIALS: tl.constexpr, + MAX_PARTS: tl.constexpr, + PRENORM: tl.constexpr, + DELAY_NS: tl.constexpr = 0, +): + if PRENORM: + # kva_ptr/stride_kva address the already normalized BF16 q_a rows. + _prenorm_gemm( + pid, + kva_ptr, + stride_kva, + wqb_ptr, + q_ptr, + H * HEAD_DIM, + rows, + rmask, + Q_LORA, + BLOCK_M, + BN, + BK, + ) + else: + _norm_gemm( + pid, + kva_ptr, + stride_kva, + gq_ptr, + wqb_ptr, + q_ptr, + qr_ptr, + ss_ptr, + n_parts, + rows, + rmask, + eps, + Q_LORA, + HEAD_DIM, + H, + BLOCK_M, + BN, + BK, + WRITE_QR, + SS_PARTIALS, + MAX_PARTS, + ) + if DO_ABSORB: + if (pid % (HEAD_DIM // BN)) < (NOPE // BN): + _stress_delay(pid, 1, DELAY_NS) + # Publish this no-RoPE tile of head pid // (HEAD_DIM // BN): every + # thread's store is ordered before the release by the barrier. + tl.debug_barrier() + tl.atomic_add(sync_ptr + pid // (HEAD_DIM // BN), 1, sem="release", scope="gpu") + + +@triton.jit +def _prenorm_gemm( + t, + x_ptr, + stride_x, + w_ptr, + out_ptr, + OUT_COLS: tl.constexpr, + rows, + rmask, + K: tl.constexpr, + BLOCK_M: tl.constexpr, + BN: tl.constexpr, + BK: tl.constexpr, +): + """``out[:, t*BN:(t+1)*BN] = x @ W[t*BN:(t+1)*BN]^T`` for BF16 ``x``.""" + cols = tl.arange(0, BK) + n_offs = t * BN + tl.arange(0, BN) + acc = tl.zeros([BLOCK_M, BN], dtype=tl.float32) + for k0 in tl.range(0, K, BK): + x = tl.load( + x_ptr + rows[:, None] * stride_x + (k0 + cols)[None, :], mask=rmask[:, None], other=0.0 + ) + w = tl.load(w_ptr + n_offs[None, :] * K + (k0 + cols)[:, None]) + acc = tl.dot(x, w, acc) + tl.store( + out_ptr + rows[:, None] * OUT_COLS + n_offs[None, :], + acc.to(tl.bfloat16), + mask=rmask[:, None], + ) + + +@triton.jit +def _norm_gemm( + pid, + kva_ptr, + stride_kva, + gq_ptr, + wqb_ptr, + q_ptr, + qr_ptr, + ss_ptr, + n_parts, + rows, + rmask, + eps, + Q_LORA: tl.constexpr, + HEAD_DIM: tl.constexpr, + H: tl.constexpr, + BLOCK_M: tl.constexpr, + BN: tl.constexpr, + BK: tl.constexpr, + WRITE_QR: tl.constexpr, + SS_PARTIALS: tl.constexpr, + MAX_PARTS: tl.constexpr, +): + if SS_PARTIALS: + rcp = _rms_rcp_partials(ss_ptr, 0, n_parts, Q_LORA, BLOCK_M, MAX_PARTS, eps) + else: + rcp = _rms_rcp(kva_ptr, stride_kva, rows, rmask, Q_LORA, BLOCK_M, BK, eps) + cols = tl.arange(0, BK) + n_offs = pid * BN + tl.arange(0, BN) + acc = tl.zeros([BLOCK_M, BN], dtype=tl.float32) + for k0 in tl.range(0, Q_LORA, BK): + x = tl.load( + kva_ptr + rows[:, None] * stride_kva + (k0 + cols)[None, :], + mask=rmask[:, None], + other=0.0, + ).to(tl.float32) + g = tl.load(gq_ptr + k0 + cols).to(tl.float32) + xn = (x * rcp[:, None] * g[None, :]).to(tl.bfloat16) + if WRITE_QR: + if pid == 0: + tl.store( + qr_ptr + rows[:, None] * Q_LORA + (k0 + cols)[None, :], xn, mask=rmask[:, None] + ) + w = tl.load(wqb_ptr + n_offs[None, :] * Q_LORA + (k0 + cols)[:, None]) + acc = tl.dot(xn, w, acc) + tl.store( + q_ptr + rows[:, None] * (H * HEAD_DIM) + n_offs[None, :], + acc.to(tl.bfloat16), + mask=rmask[:, None], + ) + + +@triton.jit +def _absorb_program( + a, + wuk_ptr, + q_ptr, + fq_ptr, + stride_fq_row, + sync_ptr, + rows, + rmask, + KV_LORA: tl.constexpr, + ROPE: tl.constexpr, + HEAD_DIM: tl.constexpr, + NOPE: tl.constexpr, + H: tl.constexpr, + BN: tl.constexpr, + BC: tl.constexpr, + USE_PDL: tl.constexpr, + DELAY_NS: tl.constexpr = 0, +): + chunks_per_head: tl.constexpr = KV_LORA // BC + h = a // chunks_per_head + c0 = (a % chunks_per_head) * BC + cc = tl.arange(0, BC) + d_lo = tl.arange(0, 128) + d_hi = tl.arange(0, NOPE - 128) + # W_UK[h] is k_b_proj_trans[h] = [KV_LORA, NOPE]; load the [NOPE, BC] + # slice for this program before waiting so it overlaps the q_b GEMM. + w_base = wuk_ptr + h * (KV_LORA * NOPE) + (c0 + cc)[None, :] * NOPE + w_lo = tl.load(w_base + d_lo[:, None]) + w_hi = tl.load(w_base + 128 + d_hi[:, None]) + # Weights are never written by an earlier grid, so only the counter and + # buffer accesses below have to wait for the programmatic dependency. + if USE_PDL: + tl.extra.cuda.gdc_wait() + + ready = tl.atomic_add(sync_ptr + h, 0, sem="acquire", scope="gpu") + while ready < NOPE // BN: + ready = tl.atomic_add(sync_ptr + h, 0, sem="acquire", scope="gpu") + tl.debug_barrier() + + q_base = q_ptr + rows[:, None] * (H * HEAD_DIM) + h * HEAD_DIM + q_lo = tl.load(q_base + d_lo[None, :], mask=rmask[:, None], other=0.0, cache_modifier=".cg") + q_hi = tl.load( + q_base + 128 + d_hi[None, :], mask=rmask[:, None], other=0.0, cache_modifier=".cg" + ) + acc = tl.dot(q_lo, w_lo) + acc = tl.dot(q_hi, w_hi, acc) + tl.store( + fq_ptr + rows[:, None] * stride_fq_row + h * (KV_LORA + ROPE) + (c0 + cc)[None, :], + acc.to(tl.bfloat16), + mask=rmask[:, None], + ) + + # Every program of head h has read the q tiles once this count reaches + # chunks_per_head; the last one resets both counters for the next launch. + _stress_delay(a, 3, DELAY_NS) + tl.debug_barrier() + done = tl.atomic_add(sync_ptr + H + h, 1, sem="acq_rel", scope="gpu") + if done == chunks_per_head - 1: + tl.atomic_xchg(sync_ptr + h, 0, sem="relaxed", scope="gpu") + tl.atomic_xchg(sync_ptr + H + h, 0, sem="relaxed", scope="gpu") + + +@triton.jit +def _kv_program( + kva_ptr, + stride_kva, + gkv_ptr, + latent_ptr, + ss_ptr, + n_parts, + rows, + rmask, + eps, + Q_LORA: tl.constexpr, + KV_LORA: tl.constexpr, + ROPE: tl.constexpr, + BLOCK_M: tl.constexpr, + SS_PARTIALS: tl.constexpr, + MAX_PARTS: tl.constexpr, +): + if SS_PARTIALS: + # Experimental own hidden GEMV only (not bit-exact: its partial sums + # follow the GEMV's column split). + kc = tl.arange(0, KV_LORA) + x = tl.load( + kva_ptr + rows[:, None] * stride_kva + Q_LORA + kc[None, :], + mask=rmask[:, None], + other=0.0, + ).to(tl.float32) + rcp = _rms_rcp_partials(ss_ptr, 1, n_parts, KV_LORA, BLOCK_M, MAX_PARTS, eps) + g = tl.load(gkv_ptr + kc).to(tl.float32) + y = (x * rcp[:, None] * g[None, :]).to(tl.bfloat16) + tl.store( + latent_ptr + rows[:, None] * (KV_LORA + ROPE) + kc[None, :], y, mask=rmask[:, None] + ) + else: + # The native kv_a_layernorm, bit for bit. + _fi_rmsnorm_rows_exact( + kva_ptr + Q_LORA, + stride_kva, + gkv_ptr, + latent_ptr, + KV_LORA + ROPE, + rows, + rmask, + eps, + KV_LORA, + BLOCK_M, + ) + rc = tl.arange(0, ROPE) + k_pe = tl.load( + kva_ptr + rows[:, None] * stride_kva + Q_LORA + KV_LORA + rc[None, :], + mask=rmask[:, None], + other=0.0, + ) + tl.store( + latent_ptr + rows[:, None] * (KV_LORA + ROPE) + KV_LORA + rc[None, :], + k_pe, + mask=rmask[:, None], + ) + + +@triton.jit +def qproj_kernel( + kva_ptr, + stride_kva, + gq_ptr, + gkv_ptr, + wqb_ptr, + wuk_ptr, + q_ptr, + qr_ptr, + latent_ptr, + fq_ptr, + stride_fq_row, + sync_ptr, + ss_ptr, + n_parts, + R, + eps_q, + eps_kv, + Q_LORA: tl.constexpr, + KV_LORA: tl.constexpr, + ROPE: tl.constexpr, + HEAD_DIM: tl.constexpr, + NOPE: tl.constexpr, + H: tl.constexpr, + BLOCK_M: tl.constexpr, + BN: tl.constexpr, + BK: tl.constexpr, + BC: tl.constexpr, + N_QB: tl.constexpr, + N_ABS: tl.constexpr, + WRITE_QR: tl.constexpr, + DO_ABSORB: tl.constexpr, + SS_PARTIALS: tl.constexpr, + MAX_PARTS: tl.constexpr, + USE_PDL: tl.constexpr, + PRENORM: tl.constexpr = False, + qn_ptr=None, + stride_qn=0, + ABSORB_IN_PRODUCER: tl.constexpr = False, + PRODUCER_JOB_SHIFT: tl.constexpr = 0, + ABS_JOB0: tl.constexpr = 0, + DELAY_NS: tl.constexpr = 0, +): + """Programs: ``[0, N_QB)`` q_b producers (with ``ABSORB_IN_PRODUCER``, + producer ``p >= PRODUCER_JOB_SHIFT`` then runs absorption job + ``p - PRODUCER_JOB_SHIFT``); ``[N_QB, N_QB + N_ABS)`` absorption jobs + ``ABS_JOB0 + i``; the last program normalizes ``compressed_kv``.""" + pid = tl.program_id(0) + rows = tl.arange(0, BLOCK_M) + rmask = rows < R + if pid < N_QB: + if ABSORB_IN_PRODUCER: + job = tl.maximum(pid - PRODUCER_JOB_SHIFT, 0) + w_lo, w_hi = _ld_wuk(job, pid >= PRODUCER_JOB_SHIFT, wuk_ptr, KV_LORA, NOPE, BC) + if USE_PDL: + tl.extra.cuda.gdc_wait() + if PRENORM: + _qb_program( + pid, + qn_ptr, + stride_qn, + gq_ptr, + wqb_ptr, + q_ptr, + qr_ptr, + sync_ptr, + ss_ptr, + n_parts, + rows, + rmask, + eps_q, + Q_LORA, + HEAD_DIM, + NOPE, + H, + BLOCK_M, + BN, + BK, + False, + DO_ABSORB, + False, + MAX_PARTS, + True, + DELAY_NS, + ) + else: + _qb_program( + pid, + kva_ptr, + stride_kva, + gq_ptr, + wqb_ptr, + q_ptr, + qr_ptr, + sync_ptr, + ss_ptr, + n_parts, + rows, + rmask, + eps_q, + Q_LORA, + HEAD_DIM, + NOPE, + H, + BLOCK_M, + BN, + BK, + WRITE_QR, + DO_ABSORB, + SS_PARTIALS, + MAX_PARTS, + False, + DELAY_NS, + ) + if ABSORB_IN_PRODUCER: + if pid >= PRODUCER_JOB_SHIFT: + _absorb_job( + job, + w_lo, + w_hi, + q_ptr, + fq_ptr, + stride_fq_row, + sync_ptr, + rows, + rmask, + KV_LORA, + ROPE, + HEAD_DIM, + NOPE, + H, + BN, + BC, + DELAY_NS, + ) + elif pid < N_QB + N_ABS: + _absorb_program( + ABS_JOB0 + pid - N_QB, + wuk_ptr, + q_ptr, + fq_ptr, + stride_fq_row, + sync_ptr, + rows, + rmask, + KV_LORA, + ROPE, + HEAD_DIM, + NOPE, + H, + BN, + BC, + USE_PDL, + DELAY_NS, + ) + else: + if USE_PDL: + tl.extra.cuda.gdc_wait() + _kv_program( + kva_ptr, + stride_kva, + gkv_ptr, + latent_ptr, + ss_ptr, + n_parts, + rows, + rmask, + eps_kv, + Q_LORA, + KV_LORA, + ROPE, + BLOCK_M, + SS_PARTIALS, + MAX_PARTS, + ) + + +@triton.jit +def _gemv_cols( + x_ptr, + stride_x, + w_ptr, + out_ptr, + stride_o, + ss_ptr, + R, + col0, + n_cols, + SS_SPLIT: tl.constexpr, + SS_END: tl.constexpr, + WRITE_SS: tl.constexpr, + CPP: tl.constexpr, + K: tl.constexpr, + BLOCK_M: tl.constexpr, + BN: tl.constexpr, + BK: tl.constexpr, +): + rows = tl.arange(0, BLOCK_M) + rmask = rows < R + lane = tl.arange(0, BN) + cols = col0 + lane + cmask = (lane < CPP) & (cols < n_cols) + kk = tl.arange(0, BK) + acc = tl.zeros([BLOCK_M, BN], dtype=tl.float32) + for k0 in tl.range(0, K, BK): + x = tl.load( + x_ptr + rows[:, None] * stride_x + (k0 + kk)[None, :], mask=rmask[:, None], other=0.0 + ) + w = tl.load(w_ptr + cols[None, :] * K + (k0 + kk)[:, None], mask=cmask[None, :], other=0.0) + acc = tl.dot(x, w, acc) + out = acc.to(tl.bfloat16) + tl.store( + out_ptr + rows[:, None] * stride_o + cols[None, :], + out, + mask=rmask[:, None] & cmask[None, :], + ) + if WRITE_SS: + # Per-program partial sums of squares of the BF16 outputs over the + # RMSNorm groups [0, SS_SPLIT) and [SS_SPLIT, SS_END); the query + # kernel reduces them instead of re-reading the rows. + sq = out.to(tl.float32) + sq = sq * sq + in_a = cmask & (cols < SS_SPLIT) + in_b = cmask & (cols >= SS_SPLIT) & (cols < SS_END) + ss_a = tl.sum(tl.where(in_a[None, :], sq, 0.0), axis=1) + ss_b = tl.sum(tl.where(in_b[None, :], sq, 0.0), axis=1) + pid = tl.program_id(0) + tl.store(ss_ptr + pid * 2 * BLOCK_M + rows, ss_a) + tl.store(ss_ptr + pid * 2 * BLOCK_M + BLOCK_M + rows, ss_b) + + +@triton.jit +def hidden_gemv_kernel( + x_ptr, + stride_x, + w0_ptr, + w1_ptr, + out0_ptr, + stride_o0, + out1_ptr, + stride_o1, + ss_ptr, + R, + N0, + N1, + P0: tl.constexpr, + CPP0: tl.constexpr, + CPP1: tl.constexpr, + SS_SPLIT: tl.constexpr, + SS_END: tl.constexpr, + WRITE_SS: tl.constexpr, + K: tl.constexpr, + BLOCK_M: tl.constexpr, + BN: tl.constexpr, + BK: tl.constexpr, + USE_PDL: tl.constexpr, +): + """``out0 = x @ W0^T`` and ``out1 = x @ W1^T`` for ``R <= BLOCK_M`` rows. + + ``W0 [N0, K]`` (the native fused ``kv_a_proj_with_mqa`` weight) and the + optional ``W1 [N1, K]`` (the indexer's fused ``wk``/``weights_proj`` + weight) share the input. Programs ``[0, P0)`` own ``CPP0`` consecutive + columns of ``W0``, the rest ``CPP1`` columns of ``W1`` (``BN`` is the + next power of two), each over the full ``K``, so the work is split + evenly over about one wave. FP32 accumulation, BF16 outputs. With + ``WRITE_SS``, ``W0`` programs also emit RMSNorm partial sums (see + ``_gemv_cols``) into ``ss_ptr [P0, 2, BLOCK_M]``. + """ + pid = tl.program_id(0) + if USE_PDL: + tl.extra.cuda.gdc_wait() + if pid < P0: + _gemv_cols( + x_ptr, + stride_x, + w0_ptr, + out0_ptr, + stride_o0, + ss_ptr, + R, + pid * CPP0, + N0, + SS_SPLIT, + SS_END, + WRITE_SS, + CPP0, + K, + BLOCK_M, + BN, + BK, + ) + else: + _gemv_cols( + x_ptr, + stride_x, + w1_ptr, + out1_ptr, + stride_o1, + ss_ptr, + R, + (pid - P0) * CPP1, + N1, + 0, + 0, + False, + CPP1, + K, + BLOCK_M, + BN, + BK, + ) + + +# -------------------------------------------------------------------------- +# Full-indexer layers: the same query kernel plus the token-wise indexer +# projection (the index_selection contract, token-wise part), reproducing the native +# ``Indexer.pre_indexer_proj`` chain for BF16 projections and FP8 keys bit +# for bit: +# q = wq_b(qr); rope(q[:, :, :64]); fused_cat_fp8(q) (per head) +# k = k_norm(wk(hidden)); rope(k[:, :64]); fused_cat_fp8(k) +# weights = weights_proj(hidden) * q_scale * softmax_scale * n_heads^-0.5 +# The fused wk/weights_proj GEMM and ``k_norm`` (LayerNorm under +# torch.compile) run natively as bridges; the kernel consumes their outputs. +# -------------------------------------------------------------------------- + + +@triton.jit +def _iq_program( + t, + kva_ptr, + stride_kva, + gq_ptr, + wiq_ptr, + iq_ptr, + isync_ptr, + ss_ptr, + n_parts, + rows, + rmask, + eps, + Q_LORA: tl.constexpr, + IHD: tl.constexpr, + IH: tl.constexpr, + BLOCK_M: tl.constexpr, + BN: tl.constexpr, + BK: tl.constexpr, + SS_PARTIALS: tl.constexpr, + MAX_PARTS: tl.constexpr, + PRENORM: tl.constexpr, + DELAY_NS: tl.constexpr = 0, +): + if PRENORM: + _prenorm_gemm( + t, kva_ptr, stride_kva, wiq_ptr, iq_ptr, IH * IHD, rows, rmask, Q_LORA, BLOCK_M, BN, BK + ) + else: + _iq_norm_gemm( + t, + kva_ptr, + stride_kva, + gq_ptr, + wiq_ptr, + iq_ptr, + ss_ptr, + n_parts, + rows, + rmask, + eps, + Q_LORA, + IHD, + IH, + BLOCK_M, + BN, + BK, + SS_PARTIALS, + MAX_PARTS, + ) + _stress_delay(t, 5, DELAY_NS) + tl.debug_barrier() + tl.atomic_add(isync_ptr + t // (IHD // BN), 1, sem="release", scope="gpu") + + +@triton.jit +def _iq_norm_gemm( + t, + kva_ptr, + stride_kva, + gq_ptr, + wiq_ptr, + iq_ptr, + ss_ptr, + n_parts, + rows, + rmask, + eps, + Q_LORA: tl.constexpr, + IHD: tl.constexpr, + IH: tl.constexpr, + BLOCK_M: tl.constexpr, + BN: tl.constexpr, + BK: tl.constexpr, + SS_PARTIALS: tl.constexpr, + MAX_PARTS: tl.constexpr, +): + if SS_PARTIALS: + rcp = _rms_rcp_partials(ss_ptr, 0, n_parts, Q_LORA, BLOCK_M, MAX_PARTS, eps) + else: + rcp = _rms_rcp(kva_ptr, stride_kva, rows, rmask, Q_LORA, BLOCK_M, BK, eps) + cols = tl.arange(0, BK) + n_offs = t * BN + tl.arange(0, BN) + acc = tl.zeros([BLOCK_M, BN], dtype=tl.float32) + for k0 in tl.range(0, Q_LORA, BK): + x = tl.load( + kva_ptr + rows[:, None] * stride_kva + (k0 + cols)[None, :], + mask=rmask[:, None], + other=0.0, + ).to(tl.float32) + g = tl.load(gq_ptr + k0 + cols).to(tl.float32) + xn = (x * rcp[:, None] * g[None, :]).to(tl.bfloat16) + w = tl.load(wiq_ptr + n_offs[None, :] * Q_LORA + (k0 + cols)[:, None]) + acc = tl.dot(xn, w, acc) + tl.store( + iq_ptr + rows[:, None] * (IH * IHD) + n_offs[None, :], + acc.to(tl.bfloat16), + mask=rmask[:, None], + ) + + +@triton.jit +def _fp8_ue8m0_scale(amax): + """``2^ceil(log2(max(amax, 1e-12) / 448))`` and its exact inverse, as in + the native ``fusedCatFp8Kernel``.""" + ratio = tl.maximum(amax, 1.0e-12) * (1.0 / 448.0) + bits = ratio.to(tl.uint32, bitcast=True) + exp_bits = bits & 0x7F800000 + exp_bits = tl.where((bits & 0x007FFFFF) != 0, exp_bits + 0x00800000, exp_bits) + scale = exp_bits.to(tl.float32, bitcast=True) + inv_scale = ((254 << 23) - exp_bits).to(tl.float32, bitcast=True) + return scale, inv_scale + + +@triton.jit +def _mul_rn(a, b): + """FP32 multiply rounded once (never contracted into an FMA).""" + return tl.inline_asm_elementwise( + "mul.rn.f32 $0, $1, $2;", "=r,r,r", [a, b], dtype=tl.float32, is_pure=True, pack=1 + ) + + +@triton.jit +def _fma_rn(a, b, c): + """``a * b + c`` with a single rounding.""" + return tl.inline_asm_elementwise( + "fma.rn.f32 $0, $1, $2, $3;", "=r,r,r,r", [a, b, c], dtype=tl.float32, is_pure=True, pack=1 + ) + + +@triton.jit +def _rope_pairs(x_even, x_odd, cos, sin): + """Interleaved (GPT-J) RoPE on BF16 pairs, FP32 math, BF16 result, with + the exact arithmetic of the native indexer RoPE (FlashInfer + ``vec_apply_llama_rope_cos_sin_interleave_reuse_half``: ``x[i] * cos + + (+-x[i ^ 1]) * sin`` compiled as ``fma(x[i], cos, rn((+-x[i ^ 1]) * + sin))``), pinned with explicit ``fma.rn`` / ``mul.rn``. Checked bit for + bit against the native kernel on real and synthetic rows.""" + o_even = _fma_rn(x_even, cos, -_mul_rn(x_odd, sin)) + o_odd = _fma_rn(x_odd, cos, _mul_rn(x_even, sin)) + return (o_even.to(tl.bfloat16).to(tl.float32), o_odd.to(tl.bfloat16).to(tl.float32)) + + +@triton.jit +def _qpost_program( + i, + iq_ptr, + isync_ptr, + kw_ptr, + stride_kw, + pos_ptr, + cos_sin_ptr, + qfp8_ptr, + qscale_ptr, + wout_ptr, + w_scale, + rows, + rmask, + IHD: tl.constexpr, + IH: tl.constexpr, + ROPE_I: tl.constexpr, + BN: tl.constexpr, + USE_PDL: tl.constexpr, + DELAY_NS: tl.constexpr = 0, +): + half: tl.constexpr = ROPE_I // 2 + nope: tl.constexpr = IHD - ROPE_I + j = tl.arange(0, half) + c = tl.arange(0, nope) + pos = tl.load(pos_ptr + rows, mask=rmask, other=0).to(tl.int64) + cs_row = cos_sin_ptr + pos[:, None] * ROPE_I + cos = tl.load(cs_row + j[None, :], mask=rmask[:, None], other=0.0) + sin = tl.load(cs_row + half + j[None, :], mask=rmask[:, None], other=0.0) + if USE_PDL: + tl.extra.cuda.gdc_wait() + w = tl.load(kw_ptr + rows * stride_kw + IHD + i, mask=rmask, other=0.0).to(tl.float32) + + ready = tl.atomic_add(isync_ptr + i, 0, sem="acquire", scope="gpu") + while ready < IHD // BN: + ready = tl.atomic_add(isync_ptr + i, 0, sem="acquire", scope="gpu") + tl.debug_barrier() + + base = iq_ptr + rows[:, None] * (IH * IHD) + i * IHD + xe = tl.load(base + 2 * j[None, :], mask=rmask[:, None], other=0.0, cache_modifier=".cg").to( + tl.float32 + ) + xo = tl.load( + base + 2 * j[None, :] + 1, mask=rmask[:, None], other=0.0, cache_modifier=".cg" + ).to(tl.float32) + xn = tl.load( + base + ROPE_I + c[None, :], mask=rmask[:, None], other=0.0, cache_modifier=".cg" + ).to(tl.float32) + oe, oo = _rope_pairs(xe, xo, cos, sin) + amax = tl.maximum( + tl.maximum(tl.max(tl.abs(oe), axis=1), tl.max(tl.abs(oo), axis=1)), + tl.max(tl.abs(xn), axis=1), + ) + scale, inv = _fp8_ue8m0_scale(amax) + out = qfp8_ptr + rows[:, None] * (IH * IHD) + i * IHD + tl.store(out + 2 * j[None, :], (oe * inv[:, None]).to(tl.float8e4nv), mask=rmask[:, None]) + tl.store(out + 2 * j[None, :] + 1, (oo * inv[:, None]).to(tl.float8e4nv), mask=rmask[:, None]) + tl.store(out + ROPE_I + c[None, :], (xn * inv[:, None]).to(tl.float8e4nv), mask=rmask[:, None]) + tl.store(qscale_ptr + rows * IH + i, scale, mask=rmask) + tl.store(wout_ptr + rows * IH + i, w * scale * w_scale, mask=rmask) + # Single consumer per head: reset for the next launch. + _stress_delay(i, 7, DELAY_NS) + tl.atomic_xchg(isync_ptr + i, 0, sem="relaxed", scope="gpu") + + +@triton.jit +def _kpost_program( + kln_ptr, + stride_kln, + pos_ptr, + cos_sin_ptr, + kfp8_ptr, + kscale_ptr, + rows, + rmask, + IHD: tl.constexpr, + ROPE_I: tl.constexpr, + USE_PDL: tl.constexpr, +): + """Indexer key tail on the native ``k_norm`` output ``kln [R, IHD]`` + (the LayerNorm itself stays native: its inductor reduction order + depends on the compiled configuration): native RoPE arithmetic + (``_rope_pairs``) and the native ``fusedCatFp8Kernel`` + quantization.""" + half: tl.constexpr = ROPE_I // 2 + nope: tl.constexpr = IHD - ROPE_I + j = tl.arange(0, half) + c = tl.arange(0, nope) + pos = tl.load(pos_ptr + rows, mask=rmask, other=0).to(tl.int64) + cs_row = cos_sin_ptr + pos[:, None] * ROPE_I + cos = tl.load(cs_row + j[None, :], mask=rmask[:, None], other=0.0) + sin = tl.load(cs_row + half + j[None, :], mask=rmask[:, None], other=0.0) + if USE_PDL: + tl.extra.cuda.gdc_wait() + base = kln_ptr + rows[:, None] * stride_kln + ye = tl.load(base + 2 * j[None, :], mask=rmask[:, None], other=0.0).to(tl.float32) + yo = tl.load(base + 2 * j[None, :] + 1, mask=rmask[:, None], other=0.0).to(tl.float32) + yn = tl.load(base + ROPE_I + c[None, :], mask=rmask[:, None], other=0.0).to(tl.float32) + oe, oo = _rope_pairs(ye, yo, cos, sin) + amax = tl.maximum( + tl.maximum(tl.max(tl.abs(oe), axis=1), tl.max(tl.abs(oo), axis=1)), + tl.max(tl.abs(yn), axis=1), + ) + scale, inv = _fp8_ue8m0_scale(amax) + out = kfp8_ptr + rows[:, None] * IHD + tl.store(out + 2 * j[None, :], (oe * inv[:, None]).to(tl.float8e4nv), mask=rmask[:, None]) + tl.store(out + 2 * j[None, :] + 1, (oo * inv[:, None]).to(tl.float8e4nv), mask=rmask[:, None]) + tl.store(out + ROPE_I + c[None, :], (yn * inv[:, None]).to(tl.float8e4nv), mask=rmask[:, None]) + tl.store(kscale_ptr + rows, scale, mask=rmask) + + +@triton.jit +def qproj_indexer_kernel( + kva_ptr, + stride_kva, + gq_ptr, + gkv_ptr, + wqb_ptr, + wuk_ptr, + q_ptr, + latent_ptr, + fq_ptr, + stride_fq_row, + sync_ptr, + ss_ptr, + n_parts, + R, + eps_q, + eps_kv, + wiq_ptr, + iq_ptr, + kw_ptr, + stride_kw, + kln_ptr, + stride_kln, + pos_ptr, + cos_sin_ptr, + qfp8_ptr, + qscale_ptr, + kfp8_ptr, + kscale_ptr, + wout_ptr, + w_scale, + Q_LORA: tl.constexpr, + KV_LORA: tl.constexpr, + ROPE: tl.constexpr, + HEAD_DIM: tl.constexpr, + NOPE: tl.constexpr, + H: tl.constexpr, + IHD: tl.constexpr, + IH: tl.constexpr, + ROPE_I: tl.constexpr, + BLOCK_M: tl.constexpr, + BN: tl.constexpr, + BK: tl.constexpr, + BC: tl.constexpr, + N_QB: tl.constexpr, + N_IQ: tl.constexpr, + N_ABS: tl.constexpr, + DO_ABSORB: tl.constexpr, + SS_PARTIALS: tl.constexpr, + MAX_PARTS: tl.constexpr, + USE_PDL: tl.constexpr, + PRENORM: tl.constexpr = False, + qn_ptr=None, + stride_qn=0, + ABSORB_IN_PRODUCER: tl.constexpr = False, + PRODUCER_JOB_SHIFT: tl.constexpr = 0, + ABS_JOB0: tl.constexpr = 0, + DELAY_NS: tl.constexpr = 0, +): + """``qproj_kernel`` plus the indexer's token-wise projection. + + Program roles, in program-id order (every consumer after the producers + it waits on): ``q_b`` tiles (with ``ABSORB_IN_PRODUCER``, producer + ``p >= PRODUCER_JOB_SHIFT`` then runs absorption job + ``p - PRODUCER_JOB_SHIFT``), indexer ``wq_b`` tiles (publish per indexer + head), absorption jobs ``ABS_JOB0 + i``, indexer query post-processing + (RoPE, FP8, weights; one program per indexer head), ``compressed_kv`` + norm, indexer key post-processing. ``sync_ptr`` holds ``2 * H`` + absorption counters followed by ``IH`` indexer-head counters. + """ + pid = tl.program_id(0) + rows = tl.arange(0, BLOCK_M) + rmask = rows < R + if pid < N_QB: + if ABSORB_IN_PRODUCER: + job = tl.maximum(pid - PRODUCER_JOB_SHIFT, 0) + w_lo, w_hi = _ld_wuk(job, pid >= PRODUCER_JOB_SHIFT, wuk_ptr, KV_LORA, NOPE, BC) + if USE_PDL: + tl.extra.cuda.gdc_wait() + if PRENORM: + _qb_program( + pid, + qn_ptr, + stride_qn, + gq_ptr, + wqb_ptr, + q_ptr, + q_ptr, + sync_ptr, + ss_ptr, + n_parts, + rows, + rmask, + eps_q, + Q_LORA, + HEAD_DIM, + NOPE, + H, + BLOCK_M, + BN, + BK, + False, + DO_ABSORB, + False, + MAX_PARTS, + True, + DELAY_NS, + ) + else: + _qb_program( + pid, + kva_ptr, + stride_kva, + gq_ptr, + wqb_ptr, + q_ptr, + q_ptr, + sync_ptr, + ss_ptr, + n_parts, + rows, + rmask, + eps_q, + Q_LORA, + HEAD_DIM, + NOPE, + H, + BLOCK_M, + BN, + BK, + False, + DO_ABSORB, + SS_PARTIALS, + MAX_PARTS, + False, + DELAY_NS, + ) + if ABSORB_IN_PRODUCER: + if pid >= PRODUCER_JOB_SHIFT: + _absorb_job( + job, + w_lo, + w_hi, + q_ptr, + fq_ptr, + stride_fq_row, + sync_ptr, + rows, + rmask, + KV_LORA, + ROPE, + HEAD_DIM, + NOPE, + H, + BN, + BC, + DELAY_NS, + ) + elif pid < N_QB + N_IQ: + if USE_PDL: + tl.extra.cuda.gdc_wait() + if PRENORM: + _iq_program( + pid - N_QB, + qn_ptr, + stride_qn, + gq_ptr, + wiq_ptr, + iq_ptr, + sync_ptr + 2 * H, + ss_ptr, + n_parts, + rows, + rmask, + eps_q, + Q_LORA, + IHD, + IH, + BLOCK_M, + BN, + BK, + False, + MAX_PARTS, + True, + DELAY_NS, + ) + else: + _iq_program( + pid - N_QB, + kva_ptr, + stride_kva, + gq_ptr, + wiq_ptr, + iq_ptr, + sync_ptr + 2 * H, + ss_ptr, + n_parts, + rows, + rmask, + eps_q, + Q_LORA, + IHD, + IH, + BLOCK_M, + BN, + BK, + SS_PARTIALS, + MAX_PARTS, + False, + DELAY_NS, + ) + elif pid < N_QB + N_IQ + N_ABS: + _absorb_program( + ABS_JOB0 + pid - N_QB - N_IQ, + wuk_ptr, + q_ptr, + fq_ptr, + stride_fq_row, + sync_ptr, + rows, + rmask, + KV_LORA, + ROPE, + HEAD_DIM, + NOPE, + H, + BN, + BC, + USE_PDL, + DELAY_NS, + ) + elif pid < N_QB + N_IQ + N_ABS + IH: + _qpost_program( + pid - N_QB - N_IQ - N_ABS, + iq_ptr, + sync_ptr + 2 * H, + kw_ptr, + stride_kw, + pos_ptr, + cos_sin_ptr, + qfp8_ptr, + qscale_ptr, + wout_ptr, + w_scale, + rows, + rmask, + IHD, + IH, + ROPE_I, + BN, + USE_PDL, + DELAY_NS, + ) + elif pid == N_QB + N_IQ + N_ABS + IH: + if USE_PDL: + tl.extra.cuda.gdc_wait() + _kv_program( + kva_ptr, + stride_kva, + gkv_ptr, + latent_ptr, + ss_ptr, + n_parts, + rows, + rmask, + eps_kv, + Q_LORA, + KV_LORA, + ROPE, + BLOCK_M, + SS_PARTIALS, + MAX_PARTS, + ) + else: + _kpost_program( + kln_ptr, + stride_kln, + pos_ptr, + cos_sin_ptr, + kfp8_ptr, + kscale_ptr, + rows, + rmask, + IHD, + ROPE_I, + USE_PDL, + ) + + +# -------------------------------------------------------------------------- +# Prefetching variant (shared and full indexer layers). +# +# The cp.async-pipelined producers above are latency bound: each CTA keeps +# only (num_stages - 1) K-tiles in flight and starts its weight stream only +# after the RMSNorm pre-pass. Here every producer issues its whole [K, BN] +# weight tile (and its absorption W_UK slice) as independent loads at +# program start -- before ``griddepcontrol.wait``, so under programmatic +# dependent launch the weight stream overlaps the preceding native kv_a +# GEMM -- then waits, loads the eight q_a row tiles at once, normalizes and +# multiplies. +# Each producer program of head h + 1 also performs one absorption job of +# head h after publishing its own tile (produce-then-consume), and TPH extra +# programs after all producers absorb the last head, so every wait targets +# lower program ids (see the module docstring). +# -------------------------------------------------------------------------- + + +@triton.jit +def _ld_w(w_ptr, n_offs, j: tl.constexpr, K: tl.constexpr, BK: tl.constexpr): + k = j * BK + tl.arange(0, BK) + return tl.load(w_ptr + n_offs[None, :] * K + k[:, None]) + + +@triton.jit +def _ld_x(kva_ptr, stride_kva, rows, rmask, j: tl.constexpr, BK: tl.constexpr): + k = j * BK + tl.arange(0, BK) + return tl.load( + kva_ptr + rows[:, None] * stride_kva + k[None, :], mask=rmask[:, None], other=0.0 + ) + + +@triton.jit +def _sq(x): + xf = x.to(tl.float32) + return tl.sum(xf * xf, axis=1) + + +@triton.jit +def _nx(x, rcp, gq_ptr, j: tl.constexpr, BK: tl.constexpr): + g = tl.load(gq_ptr + j * BK + tl.arange(0, BK)).to(tl.float32) + return (x.to(tl.float32) * rcp[:, None] * g[None, :]).to(tl.bfloat16) + + +@triton.jit +def _prefetch_tile_gemm( + t, + w_ptr, + out_ptr, + OUT_COLS: tl.constexpr, + kva_ptr, + stride_kva, + gq_ptr, + qr_ptr, + rows, + rmask, + eps, + Q_LORA: tl.constexpr, + BLOCK_M: tl.constexpr, + BN: tl.constexpr, + WRITE_QR: tl.constexpr, + USE_PDL: tl.constexpr, +): + """One [R, BN] tile of ``rmsnorm(q_a) @ W[t*BN:(t+1)*BN]^T`` with the + whole weight tile requested before the grid-dependency wait.""" + BK: tl.constexpr = Q_LORA // 8 + n_offs = t * BN + tl.arange(0, BN) + w0 = _ld_w(w_ptr, n_offs, 0, Q_LORA, BK) + w1 = _ld_w(w_ptr, n_offs, 1, Q_LORA, BK) + w2 = _ld_w(w_ptr, n_offs, 2, Q_LORA, BK) + w3 = _ld_w(w_ptr, n_offs, 3, Q_LORA, BK) + w4 = _ld_w(w_ptr, n_offs, 4, Q_LORA, BK) + w5 = _ld_w(w_ptr, n_offs, 5, Q_LORA, BK) + w6 = _ld_w(w_ptr, n_offs, 6, Q_LORA, BK) + w7 = _ld_w(w_ptr, n_offs, 7, Q_LORA, BK) + if USE_PDL: + tl.extra.cuda.gdc_wait() + x0 = _ld_x(kva_ptr, stride_kva, rows, rmask, 0, BK) + x1 = _ld_x(kva_ptr, stride_kva, rows, rmask, 1, BK) + x2 = _ld_x(kva_ptr, stride_kva, rows, rmask, 2, BK) + x3 = _ld_x(kva_ptr, stride_kva, rows, rmask, 3, BK) + x4 = _ld_x(kva_ptr, stride_kva, rows, rmask, 4, BK) + x5 = _ld_x(kva_ptr, stride_kva, rows, rmask, 5, BK) + x6 = _ld_x(kva_ptr, stride_kva, rows, rmask, 6, BK) + x7 = _ld_x(kva_ptr, stride_kva, rows, rmask, 7, BK) + ss = _sq(x0) + _sq(x1) + _sq(x2) + _sq(x3) + _sq(x4) + _sq(x5) + _sq(x6) + _sq(x7) + rcp = tl.rsqrt(ss / Q_LORA + eps) + xn0 = _nx(x0, rcp, gq_ptr, 0, BK) + xn1 = _nx(x1, rcp, gq_ptr, 1, BK) + xn2 = _nx(x2, rcp, gq_ptr, 2, BK) + xn3 = _nx(x3, rcp, gq_ptr, 3, BK) + xn4 = _nx(x4, rcp, gq_ptr, 4, BK) + xn5 = _nx(x5, rcp, gq_ptr, 5, BK) + xn6 = _nx(x6, rcp, gq_ptr, 6, BK) + xn7 = _nx(x7, rcp, gq_ptr, 7, BK) + if WRITE_QR: + if t == 0: + kk = tl.arange(0, BK) + qr_row = qr_ptr + rows[:, None] * Q_LORA + kk[None, :] + tl.store(qr_row + 0 * BK, xn0, mask=rmask[:, None]) + tl.store(qr_row + 1 * BK, xn1, mask=rmask[:, None]) + tl.store(qr_row + 2 * BK, xn2, mask=rmask[:, None]) + tl.store(qr_row + 3 * BK, xn3, mask=rmask[:, None]) + tl.store(qr_row + 4 * BK, xn4, mask=rmask[:, None]) + tl.store(qr_row + 5 * BK, xn5, mask=rmask[:, None]) + tl.store(qr_row + 6 * BK, xn6, mask=rmask[:, None]) + tl.store(qr_row + 7 * BK, xn7, mask=rmask[:, None]) + acc = tl.dot(xn0, w0) + acc = tl.dot(xn1, w1, acc) + acc = tl.dot(xn2, w2, acc) + acc = tl.dot(xn3, w3, acc) + acc = tl.dot(xn4, w4, acc) + acc = tl.dot(xn5, w5, acc) + acc = tl.dot(xn6, w6, acc) + acc = tl.dot(xn7, w7, acc) + tl.store( + out_ptr + rows[:, None] * OUT_COLS + n_offs[None, :], + acc.to(tl.bfloat16), + mask=rmask[:, None], + ) + + +@triton.jit +def _absorb_job( + a, + w_lo, + w_hi, + q_ptr, + fq_ptr, + stride_fq_row, + sync_ptr, + rows, + rmask, + KV_LORA: tl.constexpr, + ROPE: tl.constexpr, + HEAD_DIM: tl.constexpr, + NOPE: tl.constexpr, + H: tl.constexpr, + BN: tl.constexpr, + BC: tl.constexpr, + DELAY_NS: tl.constexpr = 0, +): + chunks_per_head: tl.constexpr = KV_LORA // BC + h = a // chunks_per_head + c0 = (a % chunks_per_head) * BC + cc = tl.arange(0, BC) + d_lo = tl.arange(0, 128) + d_hi = tl.arange(0, NOPE - 128) + ready = tl.atomic_add(sync_ptr + h, 0, sem="acquire", scope="gpu") + while ready < NOPE // BN: + ready = tl.atomic_add(sync_ptr + h, 0, sem="acquire", scope="gpu") + tl.debug_barrier() + q_base = q_ptr + rows[:, None] * (H * HEAD_DIM) + h * HEAD_DIM + q_lo = tl.load(q_base + d_lo[None, :], mask=rmask[:, None], other=0.0, cache_modifier=".cg") + q_hi = tl.load( + q_base + 128 + d_hi[None, :], mask=rmask[:, None], other=0.0, cache_modifier=".cg" + ) + acc = tl.dot(q_lo, w_lo) + acc = tl.dot(q_hi, w_hi, acc) + tl.store( + fq_ptr + rows[:, None] * stride_fq_row + h * (KV_LORA + ROPE) + (c0 + cc)[None, :], + acc.to(tl.bfloat16), + mask=rmask[:, None], + ) + _stress_delay(a, 3, DELAY_NS) + tl.debug_barrier() + done = tl.atomic_add(sync_ptr + H + h, 1, sem="acq_rel", scope="gpu") + if done == chunks_per_head - 1: + tl.atomic_xchg(sync_ptr + h, 0, sem="relaxed", scope="gpu") + tl.atomic_xchg(sync_ptr + H + h, 0, sem="relaxed", scope="gpu") + + +@triton.jit +def qproj_prefetch_kernel( + kva_ptr, + stride_kva, + gq_ptr, + gkv_ptr, + wqb_ptr, + wuk_ptr, + q_ptr, + qr_ptr, + latent_ptr, + fq_ptr, + stride_fq_row, + sync_ptr, + R, + eps_q, + eps_kv, + wiq_ptr, + iq_ptr, + kw_ptr, + stride_kw, + kln_ptr, + stride_kln, + pos_ptr, + cos_sin_ptr, + qfp8_ptr, + qscale_ptr, + kfp8_ptr, + kscale_ptr, + wout_ptr, + w_scale, + Q_LORA: tl.constexpr, + KV_LORA: tl.constexpr, + ROPE: tl.constexpr, + HEAD_DIM: tl.constexpr, + NOPE: tl.constexpr, + H: tl.constexpr, + IHD: tl.constexpr, + IH: tl.constexpr, + ROPE_I: tl.constexpr, + BLOCK_M: tl.constexpr, + BN: tl.constexpr, + BC: tl.constexpr, + N_QB: tl.constexpr, + N_IQ: tl.constexpr, + WRITE_QR: tl.constexpr, + DO_ABSORB: tl.constexpr, + FUSE_INDEXER: tl.constexpr, + USE_PDL: tl.constexpr, + N_ABS: tl.constexpr = 0, + PRODUCER_JOB_SHIFT: tl.constexpr = 0, + ABS_JOB0: tl.constexpr = 0, + DELAY_NS: tl.constexpr = 0, +): + """Prefetching query kernel. Program roles in program-id order: + + * ``[0, N_QB)``: ``q_b`` tile ``pid`` (prefetched), then (``DO_ABSORB``, + ``pid >= PRODUCER_JOB_SHIFT``) absorption job ``pid - + PRODUCER_JOB_SHIFT`` (``N_QB == H * KV_LORA // BC``; its ``W_UK`` slice + is also prefetched before the wait); + * ``[N_QB, N_QB + N_IQ)`` (``FUSE_INDEXER``): indexer ``wq_b`` tiles; + * ``N_ABS`` absorption jobs ``ABS_JOB0 + i``; + * ``IH`` indexer-query post-processing programs (``FUSE_INDEXER``); + * the ``compressed_kv`` norm / latent program; + * the indexer-key program (``FUSE_INDEXER``). + """ + pid = tl.program_id(0) + rows = tl.arange(0, BLOCK_M) + rmask = rows < R + if pid < N_QB: + if DO_ABSORB: + job = tl.maximum(pid - PRODUCER_JOB_SHIFT, 0) + w_lo, w_hi = _ld_wuk(job, pid >= PRODUCER_JOB_SHIFT, wuk_ptr, KV_LORA, NOPE, BC) + _prefetch_tile_gemm( + pid, + wqb_ptr, + q_ptr, + H * HEAD_DIM, + kva_ptr, + stride_kva, + gq_ptr, + qr_ptr, + rows, + rmask, + eps_q, + Q_LORA, + BLOCK_M, + BN, + WRITE_QR, + USE_PDL, + ) + if DO_ABSORB: + if (pid % (HEAD_DIM // BN)) < (NOPE // BN): + _stress_delay(pid, 1, DELAY_NS) + tl.debug_barrier() + tl.atomic_add(sync_ptr + pid // (HEAD_DIM // BN), 1, sem="release", scope="gpu") + if pid >= PRODUCER_JOB_SHIFT: + _absorb_job( + job, + w_lo, + w_hi, + q_ptr, + fq_ptr, + stride_fq_row, + sync_ptr, + rows, + rmask, + KV_LORA, + ROPE, + HEAD_DIM, + NOPE, + H, + BN, + BC, + DELAY_NS, + ) + elif pid < N_QB + N_IQ: + t = pid - N_QB + _prefetch_tile_gemm( + t, + wiq_ptr, + iq_ptr, + IH * IHD, + kva_ptr, + stride_kva, + gq_ptr, + qr_ptr, + rows, + rmask, + eps_q, + Q_LORA, + BLOCK_M, + BN, + False, + USE_PDL, + ) + _stress_delay(t, 5, DELAY_NS) + tl.debug_barrier() + tl.atomic_add(sync_ptr + 2 * H + t // (IHD // BN), 1, sem="release", scope="gpu") + elif pid < N_QB + N_IQ + N_ABS: + _absorb_program( + ABS_JOB0 + pid - N_QB - N_IQ, + wuk_ptr, + q_ptr, + fq_ptr, + stride_fq_row, + sync_ptr, + rows, + rmask, + KV_LORA, + ROPE, + HEAD_DIM, + NOPE, + H, + BN, + BC, + USE_PDL, + DELAY_NS, + ) + elif pid < N_QB + N_IQ + N_ABS + IH: + _qpost_program( + pid - N_QB - N_IQ - N_ABS, + iq_ptr, + sync_ptr + 2 * H, + kw_ptr, + stride_kw, + pos_ptr, + cos_sin_ptr, + qfp8_ptr, + qscale_ptr, + wout_ptr, + w_scale, + rows, + rmask, + IHD, + IH, + ROPE_I, + BN, + USE_PDL, + DELAY_NS, + ) + elif pid == N_QB + N_IQ + N_ABS + IH: + if USE_PDL: + tl.extra.cuda.gdc_wait() + _kv_program( + kva_ptr, + stride_kva, + gkv_ptr, + latent_ptr, + sync_ptr, + 0, + rows, + rmask, + eps_kv, + Q_LORA, + KV_LORA, + ROPE, + BLOCK_M, + False, + 16, + ) + else: + _kpost_program( + kln_ptr, + stride_kln, + pos_ptr, + cos_sin_ptr, + kfp8_ptr, + kscale_ptr, + rows, + rmask, + IHD, + ROPE_I, + USE_PDL, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/absorbed_query/query_proj.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/absorbed_query/query_proj.py new file mode 100644 index 000000000000..84ac9c10b892 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/kernels/absorbed_query/query_proj.py @@ -0,0 +1,886 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""absorbed_query: target query/projection region of GLM-5.x DSA MLA (the query_front contract). + +``dsa_proj`` is the package entry point. It replaces the token-wise work of +the native ``forward_dsa_proj`` for one generation-only target layer and +returns a :class:`AbsorbedQueryPacket`. ``mtp_dsa_proj`` (``draft_absorbed_query``) does +the same for the one-model MTP draft layer at every draft step. The packet +holds: + +* ``proj_outputs``: the native 4/9-tensor list ``[q, compressed_kv, k_pe, + latent_cache(, q_fp8, k_fp8, k_scale, weights, q_scale)]`` with native + shapes, dtypes and values (``compressed_kv``/``k_pe`` are views of + ``latent_cache``); +* ``fused_q``: in the preabsorbed form, the ``[R, H, kv_lora + rope]`` buffer + whose ``[..., :kv_lora]`` already holds ``q_nope @ W_UK`` (BF16). The + caller hands it to the native ``forward_absorption_generation`` in place of + its own allocation and skips only the absorption BMM; native RoPE, cache + append, FP8 query quantization and scheduler buffers still write the rest. + ``None`` in the native-q form. + +The kv_a GEMM and, on full indexer layers, ``Indexer.pre_indexer_proj`` run +natively around the absorbed_query kernel (adapter bridges). ``dsa_proj`` returns ``None`` +before any GPU work when the call is outside the support domain, so the +caller runs the unchanged native path. +""" + +from dataclasses import dataclass +from typing import List, Optional, Tuple + +import torch + +from ...dispatch import log_capability +from . import kernels + +#: How the form's capability line names it (``